KAFKA-20864: Fix to avoid early batch close if extension fails by lianetm · Pull Request #23008 · apache/kafka

lianetm · 2026-07-31T19:57:41Z

This PR includes the following fixes to avoid closing the wrong batch
when a mid-batch extension acquire fails on an exhausted pool:

close the batch only if it's the same we were trying to extend. If
it's not, do not close it and continue the loop, so the next iteration
checks against whatever is open (will take the new-batch path if
nothing is open, or retry with the new open batch, sizing the extension
needed again)
bound the retries the above introduces (retry of non-blocking op). Fix
by always allowing a first attempt (it's non-blocking), but bound
retries against the remaining max.block.ms. Edge case to ensure send
does not exceed max.block.ms if the extension keeps failing with the
batch being replaced. Bound all retries, successful and failed ones.
track an absolute deadline, so that the non-blocking retries time also
count against max.block.ms (like the blocking ops do)

Reviewers: Jun Rao junrao@gmail.com

junrao

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

junrao · 2026-08-03T22:52:00Z

+        // Set once the extension acquire has failed on an exhausted pool. That acquire is non-blocking, so we
+        // always allow a first attempt (even with max.block.ms 0), and only check retries of it against the
+        // deadline (see throwIfExtensionBudgetSpent), to avoid retrying it continuously with no bound.
+        boolean extensionAcquireFailed = false;


In theory, we can be in an unbounded while loop even when the extension allocation never fails. The appender will keep doing successful extension and keep finding the batch has been replaced by some other appenders. Should we gate that case too? For example, we can simply check the deadline at the beginning of the loop for each iteration except for the first one.

yeap, good catch. Changed to check the deadline on all extension retries (failed or successful).

Added it after the tryAppend though (not at the top of the loop), just to allow an append that can fit into the open batch without extension. Also added a test for this testSucceedingExtensionRetriesStopOnceMaxBlockTimeIsUsedUp

junrao · 2026-08-03T22:52:45Z

+     * on re-entry only, after {@code tryAppend} has re-confirmed the record needs chunks the exhausted
+     * pool would not hand over: the extension acquire never blocks, so we enforce the max.block.ms here.
+     */
+    private void throwIfExtensionBudgetSpent(long deadlineMs, long maxTimeToBlock, String topic, int partition) {


throwIfExtensionBudgetExceeded?

junrao · 2026-08-03T22:55:46Z

+     * {@code batchToExtend} for appends so the record retries on the new-batch path (blocks for memory)
+     * <p>
+     * The acquire runs off the deque lock, so the open batch may no longer be the one the gap was
+     * sized against by the time this would close it: it can have been drained and replaced by a batch


it can have been => it could have been

junrao · 2026-08-03T23:13:31Z

+                    ChunkedRecordAccumulator accum = accumRef.get();
+                    Deque<ProducerBatch> dq = accum.getDeque(tp1);
+                    ProducerBatch drained;
+                    synchronized (dq) {


Could we reuse simulateConcurrentDrainAndReplace here?

junrao · 2026-08-03T23:16:01Z

+            accum.append(topic, partition1, 0L, key, new byte[100], Record.EMPTY_HEADERS, null,
+                    maxBlockTimeMs, time.milliseconds(), cluster);
+
+            // Needs an extension, which fails with the budget already spent.


with the budget => with the time budget

junrao · 2026-08-03T23:18:45Z

+            assertEquals(1.0, (double) exhausted.metricValue(),
+                    "giving up on the extension path must count the dropped record exactly once");
+
+            // Giving up must leave the batch it declined to close untouched.


leave the batch it declined to close untouched => leave the open batch untouched

junrao · 2026-08-03T23:21:42Z

+     * it has room to spare for the retried one.
+     */
+    private BufferPool poolFailingFirstExtensionAfterBatchReplaced(int chunkSize,
+                                                                  AtomicReference<ChunkedRecordAccumulator> accumRef,


junrao · 2026-08-03T23:24:03Z

+    }
+
+    /**
+     * The extension acquire fails with the batch it was sized against already replaced, so nothing is


The extension acquire fails with the batch it was sized against already replaced, => The extension acquire fails with the batch it was sized against, which is already replaced,

lianetm · 2026-08-04T19:34:09Z

Thanks for the review @junrao ! All comments addressed

junrao

@lianetm : Thanks for the updated PR. A couple of more comments.

junrao · 2026-08-04T21:58:05Z

+        // or succeeded without getting the record appended (not enough capacity by then, or no longer the batch it was
+        // sized against). Neither blocks, so a retry never times out on its own and is checked against the deadline
+        // based on this flag. A first attempt always runs, even with max.block.ms 0.
+        boolean extensionRetried = false;


In theory, the existing append could loop forever if the partition keep changing. Should we have a more general approach to handle that too?

uhm interesting point about the case of partition changing, agree with the gap, but that's an issue in the "full" strategy as well, right? (send without partition + default partitioning + no key + concurrent send with same default partitioning that makes the partition change. Then append just continues retrying, unbounded)

Here in RecordAccumulator (and again a few lines below):

kafka/clients/src/main/java/org/apache/kafka/clients/producer/internals/RecordAccumulator.java

Lines 322 to 324 in 6425799

// After taking the lock, validate that the partition hasn't changed and retry.

if (partitionChanged(topic, topicInfo, partitionInfo, dq, nowMs, cluster))

continue;

We probably need to consider enforcing a deadline for that partition-change case too (just as we did on the extension path here). A separate PR/jira better? given that the issue exists in both strategies (and the fix would go in the default path RecordAccumulator and the new Chunk one too I expect)

The partition changing case applies to both strategies. I am just wondering if we can have a single mechanism to handle all cases of potential infinite loop, instead of having special mechanisms for each individual case.

Agreed, added a single mechanism to track the deadline across retries (all retries).

The fix applies to ChunkedAccum only, but I defined the core funcs in the base RecordAccum so we can reuse to bound retries on the full strategy too (even though there it should be the partition change only, will file jira to review that separately as it would touch the default path).

lianetm · 2026-08-05T19:29:49Z

Thanks @junrao ! Comments addressed, main one open above

lianetm · 2026-08-11T21:29:52Z

Thanks @junrao , comment addressed : #23008 (comment)

junrao

@lianetm : Thanks for the updated PR. A few more comments. Also, could we rerun the perf test to make sure there is no perf degradation?

junrao · 2026-08-11T23:13:49Z

+     * <li>the first pass is always allowed — the deadline is not even read, so an append that completes in one
+     *     pass never depends on the clock;</li>
+     * <li>retries are allowed while there is time left before {@code deadlineMs};</li>
+     * <li>after the deadline, one more retry is allowed, since it may need no memory at all and


Intuitively, this feels weird. Why do we want to allow an extra retry after the deadline has passed? If max.block.ms > 0, this is not needed since it's covered by deadline. If max.block.ms=0, does one more retry guarantee success when the blocking part is never hit?

If max.block.ms=0, does one more retry guarantee success when the blocking part is never hit?

no, it doesn;t (this extra allowed was mainly thinking of the case where it could guarantee success. But agree that it's weird, changed it to the simper approach : first pass always allowed. Following passes (retries) allowed while there is time.

junrao · 2026-08-11T23:15:02Z

+        // Whether the non-blocking extension was denied memory on the pass that just ended (only
+        // memory exhaustion case a pass can survive because it's non-blocking, all others throw).
+        // Cleared once the next pass has read it, so it can only ever describe the pass immediately before.
+        boolean nonBlockingAllocationDeniedMemory = false;


nonBlockingAllocationDeniedMemory => nonBlockingMemoryAllocationDenied ?

junrao · 2026-08-11T23:47:17Z

+            TimeoutException timeout = assertThrows(TimeoutException.class,
+                    () -> accum.throwIfNoMoreRetriesAllowed(expired, spent, false, topic));
+            assertEquals(TimeoutException.class, timeout.getClass(), timeout.getMessage());
+            assertTrue(timeout.getMessage().contains("kept restarting"), timeout.getMessage());


Should we use retry consistently instead of having both restart and retry?

lianetm · 2026-08-13T16:58:24Z

Thanks @junrao ! Comments addressed.

Ran the ProducerAppendPathBenchmark we used on the initial PR (#22997), trunk vs this branch with the incremental (full/default not touched on this PR): no regressions identified with that

--update
If the single mechanism I introduced here to bound retries on the incremental makes sense, I will file a jira to review the full path separately for the same (retries on partition change unbounded)

lianetm added 2 commits July 31, 2026 11:44

fix & test

01164ae

enforce deadline on retries if non-blocking extension fails

55880a0

lianetm requested a review from junrao July 31, 2026 19:57

github-actions Bot added producer clients labels Jul 31, 2026

lianetm mentioned this pull request Aug 3, 2026

KAFKA-20578: Initial producer incremental allocation for uncompressed data #22654

Merged

junrao reviewed Aug 3, 2026

View reviewed changes

lianetm added 3 commits August 4, 2026 13:15

renames & reuse func

31dd67e

fix & test

90dc331

Merge branch 'trunk' into lm-fix-early-close-race

2ff701c

junrao reviewed Aug 4, 2026

View reviewed changes

lianetm added 2 commits August 5, 2026 12:28

firstExtension flag

3ffd967

fix checkstyle

11f8a2a

lianetm added 4 commits August 11, 2026 08:41

Merge branch 'trunk' into lm-fix-early-close-race

2730a72

bound all retries

5ec0b12

tests & docs

43c3fe8

Merge branch 'trunk' into lm-fix-early-close-race

e4a7212

junrao reviewed Aug 11, 2026

View reviewed changes

lianetm added 5 commits August 12, 2026 12:14

renames

3118517

enforce deadline consistently

861ee73

tests & clean ups

9b9ef7c

tests

465d180

Merge branch 'trunk' into lm-fix-early-close-race

218b062

	// After taking the lock, validate that the partition hasn't changed and retry.
	if (partitionChanged(topic, topicInfo, partitionInfo, dq, nowMs, cluster))
	continue;

Uh oh!

Conversation

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

Uh oh!

Uh oh!

junrao left a comment

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!

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!

lianetm commented Aug 4, 2026

Uh oh!

junrao left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

lianetm Aug 5, 2026 • 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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

lianetm commented Aug 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

lianetm commented Aug 11, 2026

Uh oh!

junrao left a comment

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!

lianetm commented Aug 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

lianetm commented Jul 31, 2026 •

edited by github-actions Bot

Loading

lianetm Aug 5, 2026 •

edited

Loading

lianetm commented Aug 5, 2026 •

edited

Loading

lianetm commented Aug 13, 2026 •

edited

Loading