KAFKA-18452: Implemented batch size in acquired records #18459

apoorvmittal10 · 2025-01-09T11:23:57Z

The PR implements a way to divide acquired batches into batch size as desirable by client.

The BatchSize is the soft limit and should allign the batches in response and cached state in broker at the log batch boundaries i.e.

say produce log batch is from 0-50, 51-100

for batch size from client - 60: then response will be 0-100
for batch size from client - 40: then response will be 0-50, 51-100

This means that broker will not break the batches rather align them in response as per desirable batch size from client so complete batches are acknowledged which shall avoid offset based tracking on broker.

Committer Checklist (excluded from commit message)

Verify design and implementation
Verify test coverage and CI build status
Verify documentation (including upgrade notes)

adixitconfluent

Took a pass of non testing files. A couple of minor comments.

adixitconfluent · 2025-01-09T19:39:02Z

core/src/main/java/kafka/server/share/SharePartition.java

@@ -589,6 +589,7 @@ public long nextFetchOffset() {
    @SuppressWarnings("cyclomaticcomplexity") // Consider refactoring to avoid suppression
    public ShareAcquiredRecords acquire(
        String memberId,
+        int batchSize,


description of param batchSize is missing in the javadoc for this function

The javadoc for this new param is still missing.

Missed it, done.

adixitconfluent · 2025-01-10T12:14:44Z

core/src/main/java/kafka/server/share/SharePartition.java

+                    .setDeliveryCount((short) 1));
+            } else {
+                // The batch is split into multiple batches considering batch size.
+                long currentFirstOffset = firstAcquiredOffset;


shouldn't the variable name for currentFirstOffset be previousBatchFirstOffset, sounds more intuitive to me. Wdyt?

I don't think that's an improvement, I'm afraid. I find the code easier to follow with the current name.

adixitconfluent

Took a pass of test files as well. Left minor comments.

adixitconfluent · 2025-01-10T12:25:38Z

core/src/test/java/kafka/server/share/SharePartitionManagerTest.java

@@ -146,6 +146,7 @@ public class SharePartitionManagerTest {
    private static final short MAX_FETCH_RECORDS = 500;
    private static final int DELAYED_SHARE_FETCH_MAX_WAIT_MS = 2000;
    private static final int DELAYED_SHARE_FETCH_TIMEOUT_MS = 3000;
+    private static final int BATCH_SIZE = 500;


We have defined the same BATCH_SIZE = 500 in SharePartitionTest, DelayedShareFetchTest, ShareFetchUtilsTest and ShareFetchTest, can we reuse it?

I don't think it's necessary to do that in tests, it could be for individual classes itself.

adixitconfluent · 2025-01-10T12:33:48Z

core/src/test/java/kafka/server/share/SharePartitionTest.java

+    }
+
+    @Test
+    public void testAcquireWithBatchSizeWithBatchSize() {


the test name seems incorrect

My bad, corrected it.

junrao

@apoorvmittal10 : Thanks for the PR. Left a few comments.

junrao · 2025-01-13T18:59:25Z

core/src/main/java/kafka/server/share/SharePartition.java

@@ -589,6 +589,7 @@ public long nextFetchOffset() {
    @SuppressWarnings("cyclomaticcomplexity") // Consider refactoring to avoid suppression
    public ShareAcquiredRecords acquire(
        String memberId,
+        int batchSize,


The javadoc for this new param is still missing.

junrao · 2025-01-13T22:46:46Z

core/src/main/java/kafka/server/share/SharePartition.java

+                        result.add(new AcquiredRecords()
+                            .setFirstOffset(currentFirstOffset)
+                            .setLastOffset(lastAcquiredOffset)
+                            .setDeliveryCount((short) 1));


Could we just break here since the code outside the loop does the same thing?

Yes. Agreed.

Yeah, agree. I have simplified the if/else statement as well to make common insertion in array.

junrao · 2025-01-13T22:47:48Z

core/src/main/java/kafka/server/share/SharePartition.java

        String memberId,
        Iterable<? extends RecordBatch> batches,
        long firstOffset,
        long lastOffset,
+        int batchSize,


Currently, maxFetchRecords is the maxPollRecords. What benefits do we get by setting batchSize smaller than maxFetchRecords? To allow acks to be sent quicker?

Yes. That's it.

And not splitting batches on broker, as partial acknowledgement for a large batch results in offset based tracking for InflightBatch, so alligning them as per max.poll.records.

junrao · 2025-01-13T22:53:02Z

core/src/main/java/kafka/server/share/SharePartition.java

-                AcquiredRecords acquiredRecords = acquireNewBatchRecords(memberId, fetchPartitionData.records.batches(),
-                    firstBatch.baseOffset(), lastBatch.lastOffset(), maxFetchRecords);
-                return ShareAcquiredRecords.fromAcquiredRecords(acquiredRecords);
+                return acquireNewBatchRecords(memberId, fetchPartitionData.records.batches(),


This is an existing issue. In shareFetchResponse, we always set records to all fetch batches. However, some of the batches could be excluded during acquisition. So, it's wasteful to return excluded batches in shareFetchResponse.

Created: https://issues.apache.org/jira/browse/KAFKA-18522. I ll adjust the fetch records batches in response accordingly.

junrao · 2025-01-13T23:24:25Z

core/src/main/java/kafka/server/share/SharePartitionManager.java

@@ -244,13 +244,14 @@ public CompletableFuture<Map<TopicIdPartition, PartitionData>> fetchMessages(
        String groupId,
        String memberId,
        FetchParams fetchParams,
+        int batchSize,


missing javadoc for the new param

My bad, done.

junrao · 2025-01-14T00:05:05Z

core/src/test/java/kafka/server/share/SharePartitionTest.java

+                Optional.empty(), OptionalLong.empty(), Optional.empty(), OptionalInt.empty(), false)),
+            25 /* Gap of 3 records will also be added to first batch */);
+
+        // Fetch expected records from 3 batches, but change the first expected record to include gap offsets.


There are 4 batches, instead of 3.

Thanks for pointing out, done.

junrao · 2025-01-14T00:17:51Z

core/src/main/java/kafka/server/share/SharePartition.java

+        long firstAcquiredOffset,
+        long lastAcquiredOffset,
+        String memberId,
+        int batchSize


This is an existing issue. Will it be better for BatchSize to be a group level config instead of in each share fetch request?

We've discussed this. I would like to assess the efficacy of the approach of having the batch size information known on client and broker, and then we can decide on the best way to control it. In practice, almost all of the broker code would be the same whether it's a config or part of the request, so not too difficult to switch if we want later.

AndrewJSchofield

Thanks for the PR. Overall, it makes sense to me. Just a few minor comments to address.

AndrewJSchofield · 2025-01-14T11:39:03Z

core/src/main/java/kafka/server/share/SharePartition.java

+                    }
+
+                    // Track last batch for filling the final batch, if required.
+                    lastBatch = batch;


nit: I'd add a blank line here for legibility.

Refactored this method and lastBatch is not required now.

AndrewJSchofield · 2025-01-14T11:46:12Z

core/src/main/java/kafka/server/share/SharePartition.java

+                }
+            }
+
+            result.forEach(record -> {


The variable record is really an instance of AcquiredRecords. I suggest renaming it to acquiredRecords because record.firstOffset() and record.lastOffset() doesn't make sense for something is currently named as if it was a single record.

AndrewJSchofield · 2025-01-14T12:06:18Z

core/src/main/java/kafka/server/share/SharePartition.java

+                        result.add(new AcquiredRecords()
+                            .setFirstOffset(currentFirstOffset)
+                            .setLastOffset(lastAcquiredOffset)
+                            .setDeliveryCount((short) 1));


Yes. Agreed.

AndrewJSchofield · 2025-01-14T12:51:36Z

core/src/main/java/kafka/server/share/SharePartition.java

+        Iterable<? extends RecordBatch> batches,
+        long firstAcquiredOffset,
+        long lastAcquiredOffset,
+        String memberId,


nit: Almost always, memberId is the first argument. I suggest making it the first argument here too.

AndrewJSchofield · 2025-01-14T13:09:35Z

core/src/main/java/kafka/server/share/SharePartition.java

+                    .setDeliveryCount((short) 1));
+            } else {
+                // The batch is split into multiple batches considering batch size.
+                long currentFirstOffset = firstAcquiredOffset;


I don't think that's an improvement, I'm afraid. I find the code easier to follow with the current name.

junrao

@apoorvmittal10 : Thanks for the updated PR. A few more comments.

junrao · 2025-01-14T17:27:58Z

core/src/main/java/kafka/server/share/SharePartition.java

@@ -580,6 +580,7 @@ public long nextFetchOffset() {
     * fetched from the leader.
     *
     * @param memberId           The member id of the client that is fetching the record.
+     * @param batchSize          The batch size of the individual acquired records batch.


Could we add that the unit of batch is record?

Done, also in SharePartitionManager.

junrao · 2025-01-14T17:38:13Z

core/src/main/java/kafka/server/share/SharePartition.java

+                // as lastOffset call of RecordBatch is expensive (loads headers).
+                for (RecordBatch batch : batches) {
+                    long batchBaseOffset = batch.baseOffset();
+                    if (batchBaseOffset < firstAcquiredOffset) {


It seems that we could eliminate this statement.

That is so right, I have added multiple test to validate skipping in different scenarios.

junrao · 2025-01-14T17:54:06Z

core/src/test/java/kafka/server/share/SharePartitionTest.java

+    @Test
+    public void testAcquireWithBatchSizeAndMultipleBatches() {
+        SharePartition sharePartition = SharePartitionBuilder.builder().withState(SharePartitionState.ACTIVE).build();
+        // Create 3 batches of records.


Created 4 batches instead of 3.

Sorry, fixed.

junrao · 2025-01-14T18:00:05Z

core/src/test/java/kafka/server/share/SharePartitionTest.java

+    }
+
+    @Test
+    public void testAcquireWithBatchSizeAndMaxFetchRecords() {


We don't have a test that combines multiple record batches into a single acquired inFlightBatch. Should we add one?

I think that's being validated in testAcquireWithMultipleBatchesAndMaxFetchRecords where we provide 3 batches and 1 single response by combining first 2 batches arrive in response. Am I missing something here?

Thanks. This sounds good then.

junrao

@apoorvmittal10 : Thanks for the updated PR. Just a minor comment.

junrao · 2025-01-14T19:34:16Z

core/src/test/java/kafka/server/share/SharePartitionTest.java

+    }
+
+    @Test
+    public void testAcquireWithBatchSizeAndMaxFetchRecords() {


Thanks. This sounds good then.

junrao · 2025-01-14T19:36:05Z

core/src/test/java/kafka/server/share/SharePartitionTest.java

+    }
+
+    @Test
+    public void testAcquireSingleBatchWithBatchSizeAndEndOffsetAheadBatchFirstOffset() {


The endOffset is actually after the firstOffset of the first batch. Ditto in the two other tests below.

Yeah, that's what I wanted to test. When endOffset is ahed of first batch offset then do we acquire the batches/offsets correctly.

Do I need to have better method name so it's not confusing?

OK. Perhaps use larger since ahead can be ambiguous?

junrao

@apoorvmittal10 : Thanks for the updated PR. It seems the following tests are quarantined?

  SharePartitionTest > testAcquireSingleBatchWithBatchSizeAndEndOffsetLargerThanBatchFirstOffset()
  SharePartitionTest > testAcquireWithBatchSizeAndMaxFetchRecords()
  SharePartitionTest > testAcquireWithBatchSizeAndMultipleBatches()
  SharePartitionTest > testAcquireBatchSkipWithBatchSizeAndEndOffsetLargerThanFirstBatch()
  SharePartitionTest > testAcquireWithBatchSizeAndEndOffsetLargerThanBatchFirstOffset()
  SharePartitionTest > testAcquireWithBatchSizeAndSingleBatch()

apoorvmittal10 · 2025-01-15T09:15:19Z

@apoorvmittal10 : Thanks for the updated PR. It seems the following tests are quarantined?

  SharePartitionTest > testAcquireSingleBatchWithBatchSizeAndEndOffsetLargerThanBatchFirstOffset()
  SharePartitionTest > testAcquireWithBatchSizeAndMaxFetchRecords()
  SharePartitionTest > testAcquireWithBatchSizeAndMultipleBatches()
  SharePartitionTest > testAcquireBatchSkipWithBatchSizeAndEndOffsetLargerThanFirstBatch()
  SharePartitionTest > testAcquireWithBatchSizeAndEndOffsetLargerThanBatchFirstOffset()
  SharePartitionTest > testAcquireWithBatchSizeAndSingleBatch()

Ohh strange, looking into.

apoorvmittal10 · 2025-01-15T13:59:52Z

@apoorvmittal10 : Thanks for the updated PR. It seems the following tests are quarantined?

  SharePartitionTest > testAcquireSingleBatchWithBatchSizeAndEndOffsetLargerThanBatchFirstOffset()
  SharePartitionTest > testAcquireWithBatchSizeAndMaxFetchRecords()
  SharePartitionTest > testAcquireWithBatchSizeAndMultipleBatches()
  SharePartitionTest > testAcquireBatchSkipWithBatchSizeAndEndOffsetLargerThanFirstBatch()
  SharePartitionTest > testAcquireWithBatchSizeAndEndOffsetLargerThanBatchFirstOffset()
  SharePartitionTest > testAcquireWithBatchSizeAndSingleBatch()

@junrao I spoke to @mumrah and he mentioned that any new tests will be auto quarantined for 7 days. As these tests are added in this PR hence marked as quarantined.

junrao

@apoorvmittal10 : Thanks for the explanation. LGTM

The PR implements a way to divide acquired batches into batch size as desirable by client. The BatchSize is the soft limit and should align the batches in response and cached state in broker at the log batch boundaries. Reviewers: Andrew Schofield <[email protected]>, Abhinav Dixit <[email protected]>, Jun Rao <[email protected]>

KAFKA-18452: Implemented batch size in acquired records

0bb5055

apoorvmittal10 requested review from AndrewJSchofield, chia7712 and junrao January 9, 2025 11:24

github-actions bot added triage PRs from the community core Kafka Broker KIP-932 Queues for Kafka labels Jan 9, 2025

adixitconfluent suggested changes Jan 10, 2025

View reviewed changes

github-actions bot removed the triage PRs from the community label Jan 11, 2025

apoorvmittal10 requested a review from adixitconfluent January 13, 2025 14:34

apoorvmittal10 added 3 commits January 13, 2025 14:35

Merge remote-tracking branch 'upstream/trunk' into KAFKA-18452

fa2ae1c

Correcting test name

9477327

Fixing build

16a615b

junrao reviewed Jan 14, 2025

View reviewed changes

AndrewJSchofield requested changes Jan 14, 2025

View reviewed changes

Addressing review comments

c30e8c9

apoorvmittal10 requested review from junrao and AndrewJSchofield January 14, 2025 16:39

junrao reviewed Jan 14, 2025

View reviewed changes

Added test, removed if condition and corrected javadoc

5fbd209

apoorvmittal10 requested a review from junrao January 14, 2025 19:10

Corrected test name and comment

d1d4f28

junrao reviewed Jan 14, 2025

View reviewed changes

apoorvmittal10 added 2 commits January 14, 2025 22:35

Correcting test names

dc7e56e

Merge remote-tracking branch 'upstream/trunk' into KAFKA-18452

0d97419

junrao reviewed Jan 15, 2025

View reviewed changes

apoorvmittal10 requested a review from junrao January 15, 2025 14:00

apoorvmittal10 added 2 commits January 15, 2025 14:01

Merge remote-tracking branch 'upstream/trunk' into KAFKA-18452

b2cdb49

Merge remote-tracking branch 'upstream/trunk' into KAFKA-18452

e33011a

junrao approved these changes Jan 15, 2025

View reviewed changes

junrao merged commit 0640017 into apache:trunk Jan 15, 2025
9 checks passed

KAFKA-18452: Implemented batch size in acquired records #18459

KAFKA-18452: Implemented batch size in acquired records #18459

Conversation

apoorvmittal10 commented Jan 9, 2025

Committer Checklist (excluded from commit message)

adixitconfluent left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

adixitconfluent left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

junrao left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

AndrewJSchofield left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

junrao left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

junrao left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

junrao left a comment

Choose a reason for hiding this comment

apoorvmittal10 commented Jan 15, 2025

apoorvmittal10 commented Jan 15, 2025

junrao left a comment

Choose a reason for hiding this comment