KAFKA-20803: Demultiplex batched persister responses once (WriteStateHandler) by Shekharrajak · Pull Request #22835 · apache/kafka

Shekharrajak · 2026-07-15T04:25:25Z

fix: https://issues.apache.org/jira/browse/KAFKA-20803

PersisterStateManager coalesces many partitions into a single
WRITE_SHARE_GROUP_STATE RPC. On completion, every per-partition handler
re-scanned the full combined response to find its own slice — N handlers

O(N) scan = O(N ^ 2 ), all on the single send thread, stalling the
send loop.

This PR demultiplexes the combined response once into a topicId ->
partition -> result index (PartitionResultIndex, built in the send loop
and shared across the handlers of one coalesced request) and converts
WriteStateHandler to an O(1) lookup.

Scope is deliberately limited to WRITE: it is the hot path, since every
acknowledgement/commit drives a write. readState runs once per
SharePartition lifetime, initializeState once per group-topic-partition,
and delete/read-summary are admin-invoked — for those, building an index
would cost more than the lookup it saves, so they keep the existing
scan.

Reviewers: Sushant Mahajan smahajan@confluent.io, Andrew Schofield
aschofield@confluent.io

github-actions · 2026-07-23T04:07:15Z

A label of 'needs-attention' was automatically added to this PR in order to raise the
attention of the committers. Once this issue has been triaged, the triage label
should be removed to prevent this automation from happening again.

Shekharrajak · 2026-07-28T16:24:37Z

CI check failure looking related to #22777

…iteStateHandler

Shekharrajak · 2026-07-28T16:48:04Z

+        ToIntFunction<P> partition
+    ) {
+        Map<Uuid, Map<Integer, P>> index = new HashMap<>();
+        for (T result : results) {


Flattens results -> partitions into a nested HashMap for O(1) per-partition lookup. Generic so any RPC handler can reuse it.

Shekharrajak · 2026-07-28T16:49:00Z

+        }
+
+        @SuppressWarnings("unchecked")
+        protected final <R> R lookupPartitionResult(ClientResponse response) {


the shared O(1) path: use the shared index if present, else build lazily

Shekharrajak · 2026-07-28T16:51:05Z

-                                        super.resetCoordinatorNode();
-                                        timer.add(new PersisterTimerTask(writeStateBackoff.backOff(), this));
-                                        return;
+                    WriteShareGroupStateResponseData.PartitionResult partitionResult = lookupPartitionResult(response);


the combined response is demultiplexed once into a topicId -> partition -> result map, and each
handler does an O(1) lookup -> O(N).

About heap : lookupPartitionResult drops each handler's reference once it is consumed. so it will be alive for one callback only. No memory leak.

Copilot

Pull request overview

This PR addresses the send-thread O(N²) cost when batched share-state RPC responses are completed by introducing shared, one-time demultiplexing of combined responses and converting WriteStateHandler to use O(1) per-partition lookup.

Changes:

Added a generic indexByTopicPartition helper and per-handler plumbing (buildResultIndex / lookupPartitionResult) to support indexing combined responses once.
Updated WriteStateHandler to use the shared index instead of re-scanning the combined response.
Updated the send-loop completion path to build the index once per combined response and share it across handlers.

Suppressed comments (1)

server-common/src/main/java/org/apache/kafka/server/share/persister/PersisterStateManager.java:1600

The batched-request completion callback now calls buildResultIndex(response) before invoking handler.onComplete(response). If the client ever completes a request with response == null (which PersisterStateManagerHandler.onComplete explicitly handles), this will throw a NullPointerException and prevent all per-partition handlers from completing/failing their futures.

                                        // Demux the combined response once and share it across handlers (KAFKA-20803).
                                        Object sharedResultIndex = handlersPerGroup.isEmpty()
                                            ? null
                                            : handlersPerGroup.get(0).buildResultIndex(response);
                                        handlersPerGroup.forEach(handler1 -> {

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

smjn · 2026-07-31T06:19:50Z

+        protected final <R> R lookupPartitionResult(ClientResponse response) {
+            Object index = sharedResultIndex;
+            sharedResultIndex = null;
+            Map<Uuid, Map<Integer, R>> resultIndex =
+                (Map<Uuid, Map<Integer, R>>) (index != null ? index : buildResultIndex(response));
+            if (resultIndex == null) {
+                return null;
+            }
+            Map<Integer, R> byPartition = resultIndex.get(partitionKey().topicId());
+            return byPartition == null ? null : byPartition.get(partitionKey().partition());


looks legit

smjn

@Shekharrajak
Overall looks good.

Though I would not be too keen on modifying other RPCs to use the lookup as their frequency is extremely small(for instance a SharePartition calls readState once in its lifetime). Delete/ReadSummary are admin invoked and Initialize is once per group-topic-partition.

smjn

Slightly tangential but related - could you also add trace log in coalesceRequests to the tune of:

log.trace("Received coalesce request for {} handlers of RPC type {}", handlers.size(), rpcType);

It will help in spot checking.

smjn

Since most of the changes are within static context, do you think the cache can be refactored into its self contained inner class?

AndrewJSchofield

Thanks for the PR.

AndrewJSchofield · 2026-07-31T10:25:46Z

        private Consumer<ClientResponse> onCompleteCallback;
        protected final SharePartitionKey partitionKey;
+        // Combined-response index shared across handlers; set before onComplete, read-and-cleared in lookupPartitionResult (KAFKA-20803).
+        protected Object sharedResultIndex;


Hmm. Object here is not ideal.

Updated, thanks!

…handling

Shekharrajak · 2026-08-01T10:24:59Z

Thanks @smjn @AndrewJSchofield few updates

Trace log - added in coalesceRequests: LOG.trace("Received coalesce request for {} handlers of RPC type {}", handlers.size(), rpcType).
the index is now PartitionResultIndex
with build/get
Keeping this PR to WRITE only and updated the description

Shekharrajak · 2026-08-01T10:37:57Z

+                writeStateResult(topicId, writePartitionResult(0, Errors.NONE))))));
+
+        assertWriteStateResult(present, topicId, 0, Errors.NONE);
+        assertWriteStateResult(missing, topicId, 1, Errors.UNKNOWN_SERVER_ERROR);


the missing partition fails

Shekharrajak · 2026-08-01T10:39:40Z

+        assertEquals(1, requests.size());
+
+        // a null response must still fail every handler in the batch rather than throwing
+        assertDoesNotThrow(() -> requests.iterator().next().handler.onComplete(null));


null response

Shekharrajak · 2026-08-01T10:42:05Z

+                    writePartitionResult(0, Errors.FENCED_STATE_EPOCH),
+                    writePartitionResult(1, Errors.NONE))))));
+
+        assertWriteStateResult(topic1Partition0, topicId1, 0, Errors.NONE);


topic1 p0 and topic2 p0 are both in the same batch. Here the index key on topicId and partition helping.

Shekharrajak · 2026-08-01T10:43:54Z

+            for (T result : results) {
+                Map<Integer, P> byPartition = index.computeIfAbsent(topicId.apply(result), t -> new HashMap<>());
+                for (P p : partitions.apply(result)) {
+                    byPartition.put(partition.applyAsInt(p), p);


map inserts once and O(1) lookup

Shekharrajak · 2026-08-01T11:11:03Z

+        // that index rather than rescanning the response
+        verify(handlers.get(0), times(1)).buildResultIndex(any());
+        for (int i = 1; i < handlers.size(); i++) {
+            verify(handlers.get(i), never()).buildResultIndex(any());


If sharedResultIndex was not handed to it, it is forced into that fallback and builds its own index. Here https://github.com/apache/kafka/pull/22835/changes#diff-14030e8d09d19bdb8e5b37a07657bef62db5cb9d76dce9bb0ebcf3058f6fd002R319

So one call on handler 0, zero across the other handlers

github-actions Bot added triage PRs from the community core Kafka Broker labels Jul 15, 2026

AndrewJSchofield requested a review from smjn July 16, 2026 18:52

AndrewJSchofield added KIP-932 Queues for Kafka ci-approved labels Jul 16, 2026

github-actions Bot added the needs-attention label Jul 23, 2026

KAFKA-20803: Demultiplex batched persister responses once; convert Wr…

6d188ec

…iteStateHandler

Shekharrajak force-pushed the KAFKA-20803-write-state-handler branch from 70955e0 to 6d188ec Compare July 28, 2026 16:39

Shekharrajak commented Jul 28, 2026

View reviewed changes

github-actions Bot removed needs-attention triage PRs from the community labels Jul 29, 2026

smjn requested a review from Copilot July 31, 2026 05:26

Copilot started reviewing on behalf of smjn July 31, 2026 05:27 View session

Copilot AI reviewed Jul 31, 2026

View reviewed changes

This comment was marked as outdated.

Sign in to view

smjn reviewed Jul 31, 2026

View reviewed changes

smjn self-requested a review July 31, 2026 05:52

smjn suggested changes Jul 31, 2026

View reviewed changes

smjn reviewed Jul 31, 2026

View reviewed changes

AndrewJSchofield self-requested a review July 31, 2026 10:22

AndrewJSchofield requested changes Jul 31, 2026

View reviewed changes

Shekharrajak added 2 commits August 1, 2026 11:53

KAFKA-20803: Type the shared response index and harden null response …

2c912a0

…handling

KAFKA-20803: Test combined write-state response demultiplexing

bffc4d2

Shekharrajak commented Aug 1, 2026

View reviewed changes

KAFKA-20803: Assert the combined response is indexed once per batch

d2da6a1

Shekharrajak commented Aug 1, 2026

View reviewed changes

Uh oh!

Conversation

Shekharrajak commented Jul 15, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions Bot commented Jul 23, 2026

Uh oh!

Shekharrajak commented Jul 28, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Choose a reason for hiding this comment

Uh oh!

This comment was marked as outdated.

Uh oh!

smjn left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

smjn left a comment

Choose a reason for hiding this comment

Uh oh!

smjn left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

AndrewJSchofield left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Shekharrajak commented Aug 1, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Shekharrajak commented Jul 15, 2026 •

edited by github-actions Bot

Loading

smjn left a comment •

edited

Loading

smjn left a comment •

edited

Loading