You are viewing a plain text version of this content. The canonical link for it is here.
Posted to jira@kafka.apache.org by "junrao (via GitHub)" <gi...@apache.org> on 2026/08/03 23:48:18 UTC

Re: [PR] KAFKA-20864: Fix to avoid early batch close if extension fails [kafka]

junrao commented on code in PR #23008:
URL: https://github.com/apache/kafka/pull/23008#discussion_r3708309084


##########
clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java:
##########
@@ -480,6 +485,258 @@ public void closeForRecordAppends() {
         }
     }
 
+    /**
+     * The extension acquire runs off the deque lock, so the open batch can be replaced while it is in
+     * flight: the sender drains the batch the gap was sized against and a concurrent appender creates a
+     * new one with the memory that drain just freed. On exhaustion the append must then leave that new
+     * batch open.
+     */
+    @Test
+    public void testExhaustedExtensionLeavesAReplacementBatchOpen() throws Exception {
+        int chunkSize = 256;
+        AtomicBoolean injected = new AtomicBoolean();
+        AtomicInteger closeForAppendsCalls = new AtomicInteger();
+        AtomicReference<ChunkedRecordAccumulator> accumRef = new AtomicReference<>();
+        AtomicReference<ProducerBatch> drainedRef = new AtomicReference<>();
+
+        BufferPool pool = new BufferPool(16L * chunkSize, chunkSize, metrics, time, "producer-metrics", BufferPool.AllocationMode.INCREMENTAL) {
+            @Override
+            public List<ByteBuffer> allocateChunks(int totalSize, long maxTimeToBlockMs) throws InterruptedException {
+                // Only the first non-blocking (extension) acquire is intercepted; the deque lock is not
+                // held here, which is exactly what lets the open batch change under the appender.
+                if (maxTimeToBlockMs == 0L && injected.compareAndSet(false, true)) {
+                    ChunkedRecordAccumulator accum = accumRef.get();
+                    Deque<ProducerBatch> dq = accum.getDeque(tp1);
+                    ProducerBatch drained;
+                    synchronized (dq) {

Review Comment:
   Could we reuse `simulateConcurrentDrainAndReplace` here?



##########
clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java:
##########
@@ -266,25 +284,51 @@ public RecordAppendResult append(String topic,
         }
     }
 
+    /**
+     * Give up on the extension path once this append's share of {@code max.block.ms} is spent. Called
+     * 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) {

Review Comment:
   throwIfExtensionBudgetExceeded?



##########
clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java:
##########
@@ -266,25 +284,51 @@ public RecordAppendResult append(String topic,
         }
     }
 
+    /**
+     * Give up on the extension path once this append's share of {@code max.block.ms} is spent. Called
+     * 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) {
+        if (time.milliseconds() < deadlineMs)
+            return;
+        chunkedFree.recordBufferExhausted();
+        throw new BufferExhaustedException("Failed to extend the open batch for topic " + topic + " partition "
+                + partition + " within the remaining " + ProducerConfig.MAX_BLOCK_MS_CONFIG + " of "
+                + maxTimeToBlock + " ms. Total memory: " + chunkedFree.totalMemory()
+                + " bytes. Available memory: " + chunkedFree.availableMemory() + " bytes.");
+    }
+
     /**
      * Mid-batch extension: the open batch can still take this record so grow it in place. The
-     * acquire is non-blocking and fails fast when the pool is exhausted, closing the open batch for
-     * appends if it is still writable so the record retries on the new-batch path.
+     * acquire is non-blocking and fails fast when the pool is exhausted, closing
+     * {@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

Review Comment:
   it can have been => it could have been 



##########
clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java:
##########
@@ -122,11 +123,17 @@ public RecordAppendResult append(String topic,
         // with that size. Set and cleared together across retries; null when none is held.
         NewBatchBuffer newBatch = null;
         List<ByteBuffer> extensionChunks = null;
-        // Budget shared by every blocking acquisition this append makes, so the total blocking time
-        // stays within maxTimeToBlock. The full strategy holds its one buffer across retries and so
-        // blocks at most once; this loop can release the chunks it acquired (when a concurrent
-        // appender created a batch to extend instead) and block again on a later iteration.
-        long remainingTimeToBlock = maxTimeToBlock;
+
+        // Bounds how long this append waits for memory, so it stays within maxTimeToBlock, even across retries.
+        // E.g., this loop may release the chunks it acquired (when a concurrent appender created a batch to extend instead)
+        // and block again on a later iteration, so the blocking acquire is given whatever is left of the deadline.
+        long deadlineMs = maxTimeToBlock > Long.MAX_VALUE - nowMs ? Long.MAX_VALUE : nowMs + maxTimeToBlock;
+
+        // 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;

Review Comment:
   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.



##########
clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java:
##########
@@ -480,6 +485,258 @@ public void closeForRecordAppends() {
         }
     }
 
+    /**
+     * The extension acquire runs off the deque lock, so the open batch can be replaced while it is in
+     * flight: the sender drains the batch the gap was sized against and a concurrent appender creates a
+     * new one with the memory that drain just freed. On exhaustion the append must then leave that new
+     * batch open.
+     */
+    @Test
+    public void testExhaustedExtensionLeavesAReplacementBatchOpen() throws Exception {
+        int chunkSize = 256;
+        AtomicBoolean injected = new AtomicBoolean();
+        AtomicInteger closeForAppendsCalls = new AtomicInteger();
+        AtomicReference<ChunkedRecordAccumulator> accumRef = new AtomicReference<>();
+        AtomicReference<ProducerBatch> drainedRef = new AtomicReference<>();
+
+        BufferPool pool = new BufferPool(16L * chunkSize, chunkSize, metrics, time, "producer-metrics", BufferPool.AllocationMode.INCREMENTAL) {
+            @Override
+            public List<ByteBuffer> allocateChunks(int totalSize, long maxTimeToBlockMs) throws InterruptedException {
+                // Only the first non-blocking (extension) acquire is intercepted; the deque lock is not
+                // held here, which is exactly what lets the open batch change under the appender.
+                if (maxTimeToBlockMs == 0L && injected.compareAndSet(false, true)) {
+                    ChunkedRecordAccumulator accum = accumRef.get();
+                    Deque<ProducerBatch> dq = accum.getDeque(tp1);
+                    ProducerBatch drained;
+                    synchronized (dq) {
+                        drained = dq.pollFirst();
+                    }
+                    // Simulate the sender draining the sized batch, returning its chunks to the pool...
+                    drainedRef.set(drained);
+                    accum.deallocate(drained);
+                    // Simulate a concurrent appender claiming that memory for a new batch on the same
+                    // partition. This replaces the batch, so from here on dq.peekLast() is
+                    // no longer the batch the gap was sized against.
+                    accum.append(topic, partition1, 0L, key, new byte[100], Record.EMPTY_HEADERS, null,
+                            maxBlockTimeMs, time.milliseconds(), cluster);
+                    throw new BufferExhaustedException("injected: pool exhausted");
+                }
+                return super.allocateChunks(totalSize, maxTimeToBlockMs);
+            }
+        };
+        ChunkedRecordAccumulator accum = new ChunkedRecordAccumulator(logContext, 8192, Compression.NONE,
+                /* lingerMs */ 0, /* retryBackoffMs */ 0L, /* retryBackoffMaxMs */ 0L,
+                /* deliveryTimeoutMs */ 3200, metrics, "producer-metrics", time,
+                /* transactionManager */ null, pool) {
+            @Override
+            protected ProducerBatch createProducerBatch(TopicPartition tp, MemoryRecordsBuilder recordsBuilder, long nowMs) {
+                return new ChunkedProducerBatch(tp, recordsBuilder, nowMs) {
+                    @Override
+                    public void closeForRecordAppends() {
+                        closeForAppendsCalls.incrementAndGet();
+                        super.closeForRecordAppends();
+                    }
+                };
+            }
+        };
+        accumRef.set(accum);
+        try {
+            // First record establishes the open batch the extension gap will be sized against.
+            accum.append(topic, partition1, 0L, key, new byte[100], Record.EMPTY_HEADERS, null,
+                    maxBlockTimeMs, time.milliseconds(), cluster);
+
+            // Second record overflows that batch's chunk, so it needs an extension. The injected
+            // exhaustion fires after the batch has been replaced by the nested append's batch.
+            accum.append(topic, partition1, 0L, key, new byte[100], Record.EMPTY_HEADERS, null,
+                    maxBlockTimeMs, time.milliseconds(), cluster);
+
+            assertTrue(injected.get(), "the extension acquire must have been intercepted");
+            assertNotNull(drainedRef.get(), "the sized batch must have been drained by the injection");
+            assertEquals(0, closeForAppendsCalls.get(),
+                    "the failed extension must not close a batch it did not size the gap against");
+
+            Deque<ProducerBatch> dq = batchesFor(accum, tp1);
+            assertEquals(1, dq.size(), "only the replacement batch is expected");
+            ProducerBatch replacement = dq.peekLast();
+            // Far below its writeLimit, so isFull() can only be true via a closed append stream.
+            assertFalse(replacement.isFull(), "the replacement batch must stay open for appends");
+            assertEquals(2, replacement.recordCount,
+                    "the retried record must land in the replacement batch, extending it");
+        } finally {
+            accum.close();
+        }
+    }
+
+    /**
+     * Simulates the concurrent activity that can move the deque while an extension acquire runs off
+     * the deque lock: the sender drains the open batch, returning its chunks to the pool, and another
+     * appender claims that memory for a fresh batch in its place.
+     */
+    private void simulateConcurrentDrainAndReplace(ChunkedRecordAccumulator accum) throws InterruptedException {
+        Deque<ProducerBatch> dq = batchesFor(accum, tp1);
+        ProducerBatch drained;
+        synchronized (dq) {
+            drained = dq.pollFirst();
+        }
+        assertNotNull(drained, "there must be an open batch to drain");
+        accum.deallocate(drained);
+        accum.append(topic, partition1, 0L, key, new byte[100], Record.EMPTY_HEADERS, null,
+                maxBlockTimeMs, time.milliseconds(), cluster);
+    }
+
+    /**
+     * Pool that fails the first extension acquire (non-blocking), after the batch it was attempting
+     * to extend gets replaced, and max.block.ms has run out.
+     * The failed extension finds a batch it never sized against, so it closes nothing, and the append
+     * retries with no budget left. Only that first acquire is intercepted; later ones delegate to the
+     * real pool and succeed, so a missing deadline check surfaces as a failed assertion, not as a hang.
+     * <p>
+     * The replacement batch is pre-sized for its own single record, so {@code chunkSize} decides whether
+     * it has room to spare for the retried one.
+     */
+    private BufferPool poolFailingFirstExtensionAfterBatchReplaced(int chunkSize,
+                                                                  AtomicReference<ChunkedRecordAccumulator> accumRef,
+                                                                  AtomicBoolean injected) {
+        return new BufferPool(16L * chunkSize, chunkSize, metrics, time, "producer-metrics", BufferPool.AllocationMode.INCREMENTAL) {
+            @Override
+            public List<ByteBuffer> allocateChunks(int totalSize, long maxTimeToBlockMs) throws InterruptedException {
+                // The extension is the only acquire that does not block.
+                boolean isExtensionPath = maxTimeToBlockMs == 0L;
+                if (isExtensionPath && injected.compareAndSet(false, true)) {
+                    simulateConcurrentDrainAndReplace(accumRef.get());
+                    // Leave the append with no budget left, so the deadline lands on its retry. Stands in
+                    // for time spent earlier in the append (a prior blocking acquire, or the metadata wait).
+                    time.sleep(maxBlockTimeMs + 1);
+                    throw new BufferExhaustedException("injected: pool exhausted");
+                }
+                return super.allocateChunks(totalSize, maxTimeToBlockMs);
+            }
+        };
+    }
+
+    /**
+     * The extension acquire fails with the batch it was sized against already replaced, so nothing is

Review Comment:
   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,



##########
clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java:
##########
@@ -480,6 +485,258 @@ public void closeForRecordAppends() {
         }
     }
 
+    /**
+     * The extension acquire runs off the deque lock, so the open batch can be replaced while it is in
+     * flight: the sender drains the batch the gap was sized against and a concurrent appender creates a
+     * new one with the memory that drain just freed. On exhaustion the append must then leave that new
+     * batch open.
+     */
+    @Test
+    public void testExhaustedExtensionLeavesAReplacementBatchOpen() throws Exception {
+        int chunkSize = 256;
+        AtomicBoolean injected = new AtomicBoolean();
+        AtomicInteger closeForAppendsCalls = new AtomicInteger();
+        AtomicReference<ChunkedRecordAccumulator> accumRef = new AtomicReference<>();
+        AtomicReference<ProducerBatch> drainedRef = new AtomicReference<>();
+
+        BufferPool pool = new BufferPool(16L * chunkSize, chunkSize, metrics, time, "producer-metrics", BufferPool.AllocationMode.INCREMENTAL) {
+            @Override
+            public List<ByteBuffer> allocateChunks(int totalSize, long maxTimeToBlockMs) throws InterruptedException {
+                // Only the first non-blocking (extension) acquire is intercepted; the deque lock is not
+                // held here, which is exactly what lets the open batch change under the appender.
+                if (maxTimeToBlockMs == 0L && injected.compareAndSet(false, true)) {
+                    ChunkedRecordAccumulator accum = accumRef.get();
+                    Deque<ProducerBatch> dq = accum.getDeque(tp1);
+                    ProducerBatch drained;
+                    synchronized (dq) {
+                        drained = dq.pollFirst();
+                    }
+                    // Simulate the sender draining the sized batch, returning its chunks to the pool...
+                    drainedRef.set(drained);
+                    accum.deallocate(drained);
+                    // Simulate a concurrent appender claiming that memory for a new batch on the same
+                    // partition. This replaces the batch, so from here on dq.peekLast() is
+                    // no longer the batch the gap was sized against.
+                    accum.append(topic, partition1, 0L, key, new byte[100], Record.EMPTY_HEADERS, null,
+                            maxBlockTimeMs, time.milliseconds(), cluster);
+                    throw new BufferExhaustedException("injected: pool exhausted");
+                }
+                return super.allocateChunks(totalSize, maxTimeToBlockMs);
+            }
+        };
+        ChunkedRecordAccumulator accum = new ChunkedRecordAccumulator(logContext, 8192, Compression.NONE,
+                /* lingerMs */ 0, /* retryBackoffMs */ 0L, /* retryBackoffMaxMs */ 0L,
+                /* deliveryTimeoutMs */ 3200, metrics, "producer-metrics", time,
+                /* transactionManager */ null, pool) {
+            @Override
+            protected ProducerBatch createProducerBatch(TopicPartition tp, MemoryRecordsBuilder recordsBuilder, long nowMs) {
+                return new ChunkedProducerBatch(tp, recordsBuilder, nowMs) {
+                    @Override
+                    public void closeForRecordAppends() {
+                        closeForAppendsCalls.incrementAndGet();
+                        super.closeForRecordAppends();
+                    }
+                };
+            }
+        };
+        accumRef.set(accum);
+        try {
+            // First record establishes the open batch the extension gap will be sized against.
+            accum.append(topic, partition1, 0L, key, new byte[100], Record.EMPTY_HEADERS, null,
+                    maxBlockTimeMs, time.milliseconds(), cluster);
+
+            // Second record overflows that batch's chunk, so it needs an extension. The injected
+            // exhaustion fires after the batch has been replaced by the nested append's batch.
+            accum.append(topic, partition1, 0L, key, new byte[100], Record.EMPTY_HEADERS, null,
+                    maxBlockTimeMs, time.milliseconds(), cluster);
+
+            assertTrue(injected.get(), "the extension acquire must have been intercepted");
+            assertNotNull(drainedRef.get(), "the sized batch must have been drained by the injection");
+            assertEquals(0, closeForAppendsCalls.get(),
+                    "the failed extension must not close a batch it did not size the gap against");
+
+            Deque<ProducerBatch> dq = batchesFor(accum, tp1);
+            assertEquals(1, dq.size(), "only the replacement batch is expected");
+            ProducerBatch replacement = dq.peekLast();
+            // Far below its writeLimit, so isFull() can only be true via a closed append stream.
+            assertFalse(replacement.isFull(), "the replacement batch must stay open for appends");
+            assertEquals(2, replacement.recordCount,
+                    "the retried record must land in the replacement batch, extending it");
+        } finally {
+            accum.close();
+        }
+    }
+
+    /**
+     * Simulates the concurrent activity that can move the deque while an extension acquire runs off
+     * the deque lock: the sender drains the open batch, returning its chunks to the pool, and another
+     * appender claims that memory for a fresh batch in its place.
+     */
+    private void simulateConcurrentDrainAndReplace(ChunkedRecordAccumulator accum) throws InterruptedException {
+        Deque<ProducerBatch> dq = batchesFor(accum, tp1);
+        ProducerBatch drained;
+        synchronized (dq) {
+            drained = dq.pollFirst();
+        }
+        assertNotNull(drained, "there must be an open batch to drain");
+        accum.deallocate(drained);
+        accum.append(topic, partition1, 0L, key, new byte[100], Record.EMPTY_HEADERS, null,
+                maxBlockTimeMs, time.milliseconds(), cluster);
+    }
+
+    /**
+     * Pool that fails the first extension acquire (non-blocking), after the batch it was attempting
+     * to extend gets replaced, and max.block.ms has run out.
+     * The failed extension finds a batch it never sized against, so it closes nothing, and the append
+     * retries with no budget left. Only that first acquire is intercepted; later ones delegate to the
+     * real pool and succeed, so a missing deadline check surfaces as a failed assertion, not as a hang.
+     * <p>
+     * The replacement batch is pre-sized for its own single record, so {@code chunkSize} decides whether
+     * it has room to spare for the retried one.
+     */
+    private BufferPool poolFailingFirstExtensionAfterBatchReplaced(int chunkSize,
+                                                                  AtomicReference<ChunkedRecordAccumulator> accumRef,
+                                                                  AtomicBoolean injected) {
+        return new BufferPool(16L * chunkSize, chunkSize, metrics, time, "producer-metrics", BufferPool.AllocationMode.INCREMENTAL) {
+            @Override
+            public List<ByteBuffer> allocateChunks(int totalSize, long maxTimeToBlockMs) throws InterruptedException {
+                // The extension is the only acquire that does not block.
+                boolean isExtensionPath = maxTimeToBlockMs == 0L;
+                if (isExtensionPath && injected.compareAndSet(false, true)) {
+                    simulateConcurrentDrainAndReplace(accumRef.get());
+                    // Leave the append with no budget left, so the deadline lands on its retry. Stands in
+                    // for time spent earlier in the append (a prior blocking acquire, or the metadata wait).
+                    time.sleep(maxBlockTimeMs + 1);
+                    throw new BufferExhaustedException("injected: pool exhausted");
+                }
+                return super.allocateChunks(totalSize, maxTimeToBlockMs);
+            }
+        };
+    }
+
+    /**
+     * The extension acquire fails with the batch it was sized against already replaced, so nothing is
+     * closed, the budget is spent, and the replacement needs memory too. The append must give up
+     * rather than come back to the same non-blocking acquire indefinitely.
+     */
+    @Test
+    public void testExtensionRetryStopsOnceMaxBlockTimeIsUsedUp() throws Exception {
+        AtomicBoolean injected = new AtomicBoolean();
+        AtomicReference<ChunkedRecordAccumulator> accumRef = new AtomicReference<>();
+        // A chunk barely bigger than a record, so the replacement batch cannot take the retried one
+        // either and the append has to give up rather than recover.
+        BufferPool pool = poolFailingFirstExtensionAfterBatchReplaced(256, accumRef, injected);
+        ChunkedRecordAccumulator accum = newAccumulator(8192, Compression.NONE, pool);
+        accumRef.set(accum);
+        try {
+            KafkaMetric exhausted = metrics.metric(metrics.metricName("buffer-exhausted-total", "producer-metrics"));
+
+            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.

Review Comment:
   with the budget  => with the time budget 



##########
clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java:
##########
@@ -480,6 +485,258 @@ public void closeForRecordAppends() {
         }
     }
 
+    /**
+     * The extension acquire runs off the deque lock, so the open batch can be replaced while it is in
+     * flight: the sender drains the batch the gap was sized against and a concurrent appender creates a
+     * new one with the memory that drain just freed. On exhaustion the append must then leave that new
+     * batch open.
+     */
+    @Test
+    public void testExhaustedExtensionLeavesAReplacementBatchOpen() throws Exception {
+        int chunkSize = 256;
+        AtomicBoolean injected = new AtomicBoolean();
+        AtomicInteger closeForAppendsCalls = new AtomicInteger();
+        AtomicReference<ChunkedRecordAccumulator> accumRef = new AtomicReference<>();
+        AtomicReference<ProducerBatch> drainedRef = new AtomicReference<>();
+
+        BufferPool pool = new BufferPool(16L * chunkSize, chunkSize, metrics, time, "producer-metrics", BufferPool.AllocationMode.INCREMENTAL) {
+            @Override
+            public List<ByteBuffer> allocateChunks(int totalSize, long maxTimeToBlockMs) throws InterruptedException {
+                // Only the first non-blocking (extension) acquire is intercepted; the deque lock is not
+                // held here, which is exactly what lets the open batch change under the appender.
+                if (maxTimeToBlockMs == 0L && injected.compareAndSet(false, true)) {
+                    ChunkedRecordAccumulator accum = accumRef.get();
+                    Deque<ProducerBatch> dq = accum.getDeque(tp1);
+                    ProducerBatch drained;
+                    synchronized (dq) {
+                        drained = dq.pollFirst();
+                    }
+                    // Simulate the sender draining the sized batch, returning its chunks to the pool...
+                    drainedRef.set(drained);
+                    accum.deallocate(drained);
+                    // Simulate a concurrent appender claiming that memory for a new batch on the same
+                    // partition. This replaces the batch, so from here on dq.peekLast() is
+                    // no longer the batch the gap was sized against.
+                    accum.append(topic, partition1, 0L, key, new byte[100], Record.EMPTY_HEADERS, null,
+                            maxBlockTimeMs, time.milliseconds(), cluster);
+                    throw new BufferExhaustedException("injected: pool exhausted");
+                }
+                return super.allocateChunks(totalSize, maxTimeToBlockMs);
+            }
+        };
+        ChunkedRecordAccumulator accum = new ChunkedRecordAccumulator(logContext, 8192, Compression.NONE,
+                /* lingerMs */ 0, /* retryBackoffMs */ 0L, /* retryBackoffMaxMs */ 0L,
+                /* deliveryTimeoutMs */ 3200, metrics, "producer-metrics", time,
+                /* transactionManager */ null, pool) {
+            @Override
+            protected ProducerBatch createProducerBatch(TopicPartition tp, MemoryRecordsBuilder recordsBuilder, long nowMs) {
+                return new ChunkedProducerBatch(tp, recordsBuilder, nowMs) {
+                    @Override
+                    public void closeForRecordAppends() {
+                        closeForAppendsCalls.incrementAndGet();
+                        super.closeForRecordAppends();
+                    }
+                };
+            }
+        };
+        accumRef.set(accum);
+        try {
+            // First record establishes the open batch the extension gap will be sized against.
+            accum.append(topic, partition1, 0L, key, new byte[100], Record.EMPTY_HEADERS, null,
+                    maxBlockTimeMs, time.milliseconds(), cluster);
+
+            // Second record overflows that batch's chunk, so it needs an extension. The injected
+            // exhaustion fires after the batch has been replaced by the nested append's batch.
+            accum.append(topic, partition1, 0L, key, new byte[100], Record.EMPTY_HEADERS, null,
+                    maxBlockTimeMs, time.milliseconds(), cluster);
+
+            assertTrue(injected.get(), "the extension acquire must have been intercepted");
+            assertNotNull(drainedRef.get(), "the sized batch must have been drained by the injection");
+            assertEquals(0, closeForAppendsCalls.get(),
+                    "the failed extension must not close a batch it did not size the gap against");
+
+            Deque<ProducerBatch> dq = batchesFor(accum, tp1);
+            assertEquals(1, dq.size(), "only the replacement batch is expected");
+            ProducerBatch replacement = dq.peekLast();
+            // Far below its writeLimit, so isFull() can only be true via a closed append stream.
+            assertFalse(replacement.isFull(), "the replacement batch must stay open for appends");
+            assertEquals(2, replacement.recordCount,
+                    "the retried record must land in the replacement batch, extending it");
+        } finally {
+            accum.close();
+        }
+    }
+
+    /**
+     * Simulates the concurrent activity that can move the deque while an extension acquire runs off
+     * the deque lock: the sender drains the open batch, returning its chunks to the pool, and another
+     * appender claims that memory for a fresh batch in its place.
+     */
+    private void simulateConcurrentDrainAndReplace(ChunkedRecordAccumulator accum) throws InterruptedException {
+        Deque<ProducerBatch> dq = batchesFor(accum, tp1);
+        ProducerBatch drained;
+        synchronized (dq) {
+            drained = dq.pollFirst();
+        }
+        assertNotNull(drained, "there must be an open batch to drain");
+        accum.deallocate(drained);
+        accum.append(topic, partition1, 0L, key, new byte[100], Record.EMPTY_HEADERS, null,
+                maxBlockTimeMs, time.milliseconds(), cluster);
+    }
+
+    /**
+     * Pool that fails the first extension acquire (non-blocking), after the batch it was attempting
+     * to extend gets replaced, and max.block.ms has run out.
+     * The failed extension finds a batch it never sized against, so it closes nothing, and the append
+     * retries with no budget left. Only that first acquire is intercepted; later ones delegate to the
+     * real pool and succeed, so a missing deadline check surfaces as a failed assertion, not as a hang.
+     * <p>
+     * The replacement batch is pre-sized for its own single record, so {@code chunkSize} decides whether
+     * it has room to spare for the retried one.
+     */
+    private BufferPool poolFailingFirstExtensionAfterBatchReplaced(int chunkSize,
+                                                                  AtomicReference<ChunkedRecordAccumulator> accumRef,
+                                                                  AtomicBoolean injected) {
+        return new BufferPool(16L * chunkSize, chunkSize, metrics, time, "producer-metrics", BufferPool.AllocationMode.INCREMENTAL) {
+            @Override
+            public List<ByteBuffer> allocateChunks(int totalSize, long maxTimeToBlockMs) throws InterruptedException {
+                // The extension is the only acquire that does not block.
+                boolean isExtensionPath = maxTimeToBlockMs == 0L;
+                if (isExtensionPath && injected.compareAndSet(false, true)) {
+                    simulateConcurrentDrainAndReplace(accumRef.get());
+                    // Leave the append with no budget left, so the deadline lands on its retry. Stands in
+                    // for time spent earlier in the append (a prior blocking acquire, or the metadata wait).
+                    time.sleep(maxBlockTimeMs + 1);
+                    throw new BufferExhaustedException("injected: pool exhausted");
+                }
+                return super.allocateChunks(totalSize, maxTimeToBlockMs);
+            }
+        };
+    }
+
+    /**
+     * The extension acquire fails with the batch it was sized against already replaced, so nothing is
+     * closed, the budget is spent, and the replacement needs memory too. The append must give up
+     * rather than come back to the same non-blocking acquire indefinitely.
+     */
+    @Test
+    public void testExtensionRetryStopsOnceMaxBlockTimeIsUsedUp() throws Exception {
+        AtomicBoolean injected = new AtomicBoolean();
+        AtomicReference<ChunkedRecordAccumulator> accumRef = new AtomicReference<>();
+        // A chunk barely bigger than a record, so the replacement batch cannot take the retried one
+        // either and the append has to give up rather than recover.
+        BufferPool pool = poolFailingFirstExtensionAfterBatchReplaced(256, accumRef, injected);
+        ChunkedRecordAccumulator accum = newAccumulator(8192, Compression.NONE, pool);
+        accumRef.set(accum);
+        try {
+            KafkaMetric exhausted = metrics.metric(metrics.metricName("buffer-exhausted-total", "producer-metrics"));
+
+            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.
+            BufferExhaustedException e = assertThrows(BufferExhaustedException.class,
+                    () -> accum.append(topic, partition1, 0L, key, new byte[100], Record.EMPTY_HEADERS, null,
+                            maxBlockTimeMs, time.milliseconds(), cluster));
+            // Must come from the extension path's deadline check, not from the blocking new-batch
+            // acquire — BufferPool's own exhaustion message says nothing about extending a batch.
+            assertTrue(e.getMessage().contains("Failed to extend the open batch"), e.getMessage());
+            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.

Review Comment:
   leave the batch it declined to close untouched => leave the open batch untouched



##########
clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java:
##########
@@ -480,6 +485,258 @@ public void closeForRecordAppends() {
         }
     }
 
+    /**
+     * The extension acquire runs off the deque lock, so the open batch can be replaced while it is in
+     * flight: the sender drains the batch the gap was sized against and a concurrent appender creates a
+     * new one with the memory that drain just freed. On exhaustion the append must then leave that new
+     * batch open.
+     */
+    @Test
+    public void testExhaustedExtensionLeavesAReplacementBatchOpen() throws Exception {
+        int chunkSize = 256;
+        AtomicBoolean injected = new AtomicBoolean();
+        AtomicInteger closeForAppendsCalls = new AtomicInteger();
+        AtomicReference<ChunkedRecordAccumulator> accumRef = new AtomicReference<>();
+        AtomicReference<ProducerBatch> drainedRef = new AtomicReference<>();
+
+        BufferPool pool = new BufferPool(16L * chunkSize, chunkSize, metrics, time, "producer-metrics", BufferPool.AllocationMode.INCREMENTAL) {
+            @Override
+            public List<ByteBuffer> allocateChunks(int totalSize, long maxTimeToBlockMs) throws InterruptedException {
+                // Only the first non-blocking (extension) acquire is intercepted; the deque lock is not
+                // held here, which is exactly what lets the open batch change under the appender.
+                if (maxTimeToBlockMs == 0L && injected.compareAndSet(false, true)) {
+                    ChunkedRecordAccumulator accum = accumRef.get();
+                    Deque<ProducerBatch> dq = accum.getDeque(tp1);
+                    ProducerBatch drained;
+                    synchronized (dq) {
+                        drained = dq.pollFirst();
+                    }
+                    // Simulate the sender draining the sized batch, returning its chunks to the pool...
+                    drainedRef.set(drained);
+                    accum.deallocate(drained);
+                    // Simulate a concurrent appender claiming that memory for a new batch on the same
+                    // partition. This replaces the batch, so from here on dq.peekLast() is
+                    // no longer the batch the gap was sized against.
+                    accum.append(topic, partition1, 0L, key, new byte[100], Record.EMPTY_HEADERS, null,
+                            maxBlockTimeMs, time.milliseconds(), cluster);
+                    throw new BufferExhaustedException("injected: pool exhausted");
+                }
+                return super.allocateChunks(totalSize, maxTimeToBlockMs);
+            }
+        };
+        ChunkedRecordAccumulator accum = new ChunkedRecordAccumulator(logContext, 8192, Compression.NONE,
+                /* lingerMs */ 0, /* retryBackoffMs */ 0L, /* retryBackoffMaxMs */ 0L,
+                /* deliveryTimeoutMs */ 3200, metrics, "producer-metrics", time,
+                /* transactionManager */ null, pool) {
+            @Override
+            protected ProducerBatch createProducerBatch(TopicPartition tp, MemoryRecordsBuilder recordsBuilder, long nowMs) {
+                return new ChunkedProducerBatch(tp, recordsBuilder, nowMs) {
+                    @Override
+                    public void closeForRecordAppends() {
+                        closeForAppendsCalls.incrementAndGet();
+                        super.closeForRecordAppends();
+                    }
+                };
+            }
+        };
+        accumRef.set(accum);
+        try {
+            // First record establishes the open batch the extension gap will be sized against.
+            accum.append(topic, partition1, 0L, key, new byte[100], Record.EMPTY_HEADERS, null,
+                    maxBlockTimeMs, time.milliseconds(), cluster);
+
+            // Second record overflows that batch's chunk, so it needs an extension. The injected
+            // exhaustion fires after the batch has been replaced by the nested append's batch.
+            accum.append(topic, partition1, 0L, key, new byte[100], Record.EMPTY_HEADERS, null,
+                    maxBlockTimeMs, time.milliseconds(), cluster);
+
+            assertTrue(injected.get(), "the extension acquire must have been intercepted");
+            assertNotNull(drainedRef.get(), "the sized batch must have been drained by the injection");
+            assertEquals(0, closeForAppendsCalls.get(),
+                    "the failed extension must not close a batch it did not size the gap against");
+
+            Deque<ProducerBatch> dq = batchesFor(accum, tp1);
+            assertEquals(1, dq.size(), "only the replacement batch is expected");
+            ProducerBatch replacement = dq.peekLast();
+            // Far below its writeLimit, so isFull() can only be true via a closed append stream.
+            assertFalse(replacement.isFull(), "the replacement batch must stay open for appends");
+            assertEquals(2, replacement.recordCount,
+                    "the retried record must land in the replacement batch, extending it");
+        } finally {
+            accum.close();
+        }
+    }
+
+    /**
+     * Simulates the concurrent activity that can move the deque while an extension acquire runs off
+     * the deque lock: the sender drains the open batch, returning its chunks to the pool, and another
+     * appender claims that memory for a fresh batch in its place.
+     */
+    private void simulateConcurrentDrainAndReplace(ChunkedRecordAccumulator accum) throws InterruptedException {
+        Deque<ProducerBatch> dq = batchesFor(accum, tp1);
+        ProducerBatch drained;
+        synchronized (dq) {
+            drained = dq.pollFirst();
+        }
+        assertNotNull(drained, "there must be an open batch to drain");
+        accum.deallocate(drained);
+        accum.append(topic, partition1, 0L, key, new byte[100], Record.EMPTY_HEADERS, null,
+                maxBlockTimeMs, time.milliseconds(), cluster);
+    }
+
+    /**
+     * Pool that fails the first extension acquire (non-blocking), after the batch it was attempting
+     * to extend gets replaced, and max.block.ms has run out.
+     * The failed extension finds a batch it never sized against, so it closes nothing, and the append
+     * retries with no budget left. Only that first acquire is intercepted; later ones delegate to the
+     * real pool and succeed, so a missing deadline check surfaces as a failed assertion, not as a hang.
+     * <p>
+     * The replacement batch is pre-sized for its own single record, so {@code chunkSize} decides whether
+     * it has room to spare for the retried one.
+     */
+    private BufferPool poolFailingFirstExtensionAfterBatchReplaced(int chunkSize,
+                                                                  AtomicReference<ChunkedRecordAccumulator> accumRef,

Review Comment:
   identation



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: jira-unsubscribe@kafka.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org