From 4ec8d0467b65d322845c45f34e673ecdaa0b6aaa Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Fri, 5 Jun 2026 12:09:20 -0400 Subject: [PATCH 01/61] support dyn for uncompressed --- checkstyle/suppressions.xml | 2 +- .../producer/internals/BufferPool.java | 62 ++-- .../producer/internals/ChunkedBufferPool.java | 200 +++++++++++ .../ChunkedByteBufferOutputStream.java | 246 +++++++++++++ .../internals/ChunkedRecordAccumulator.java | 336 ++++++++++++++++++ .../producer/internals/ProducerBatch.java | 33 ++ .../producer/internals/RecordAccumulator.java | 61 ++-- .../record/internal/MemoryRecordsBuilder.java | 9 + gradle/spotbugs-exclude.xml | 12 + 9 files changed, 914 insertions(+), 47 deletions(-) create mode 100644 clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java create mode 100644 clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java create mode 100644 clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java diff --git a/checkstyle/suppressions.xml b/checkstyle/suppressions.xml index 1881f95ead5ea..a537e61d0217e 100644 --- a/checkstyle/suppressions.xml +++ b/checkstyle/suppressions.xml @@ -71,7 +71,7 @@ files="(Utils|Topic|Lz4BlockOutputStream|JoinGroupRequest).java"/> + files="(AbstractFetch|ChunkedBufferPool|ClientTelemetryReporter|ConsumerCoordinator|CommitRequestManager|FetchCollector|OffsetFetcherUtils|KafkaProducer|Sender|ConfigDef|KerberosLogin|AbstractRequest|AbstractResponse|Selector|SslFactory|SslTransportLayer|SaslClientAuthenticator|SaslClientCallbackHandler|SaslServerAuthenticator|AbstractCoordinator|TransactionManager|AbstractStickyAssignor|DefaultSslEngineFactory|Authorizer|RecordAccumulator|MemoryRecords|FetchSessionHandler|MockAdminClient).java"/> diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/BufferPool.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/BufferPool.java index 517a2bd9ca7a6..f376c0fe2f4a3 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/BufferPool.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/BufferPool.java @@ -48,15 +48,19 @@ public class BufferPool { private final long totalMemory; private final int poolableSize; - private final ReentrantLock lock; - private final Deque free; - private final Deque waiters; - /** Total available memory is the sum of nonPooledAvailableMemory and the number of byte buffers in free * poolableSize. */ - private long nonPooledAvailableMemory; + /** Lock held for any read or write of {@link #free}, {@link #waiters}, {@link #nonPooledAvailableMemory}, or {@link #closed}. */ + protected final ReentrantLock lock; + /** Pooled buffers of capacity {@link #poolableSize}, available for reuse. Guarded by {@link #lock}. */ + protected final Deque free; + /** FIFO queue of pending allocation requests; the longest-waiting thread is woken first. Guarded by {@link #lock}. */ + protected final Deque waiters; + /** Total available memory is the sum of nonPooledAvailableMemory and the number of byte buffers in free * poolableSize. Guarded by {@link #lock}. */ + protected long nonPooledAvailableMemory; private final Metrics metrics; - private final Time time; + protected final Time time; private final Sensor waitTime; - private boolean closed; + /** True once {@link #close()} has been invoked. */ + protected boolean closed; /** * Create a new buffer pool @@ -192,8 +196,7 @@ public ByteBuffer allocate(int size, long maxTimeToBlockMs) throws InterruptedEx // signal any additional waiters if there is more memory left // over for them try { - if (!(this.nonPooledAvailableMemory == 0 && this.free.isEmpty()) && !this.waiters.isEmpty()) - this.waiters.peekFirst().signal(); + signalNextWaiterIfMemoryAvailable(); } finally { // Another finally... otherwise find bugs complains lock.unlock(); @@ -211,6 +214,15 @@ protected void recordWaitTime(long timeNs) { this.waitTime.record(timeNs, time.milliseconds()); } + /** + * Wake the longest-waiting thread if any memory (pooled or non-pooled) is available. + * Must be called with {@link #lock} held. No-op if no waiters or no memory is free. + */ + protected void signalNextWaiterIfMemoryAvailable() { + if (!(this.nonPooledAvailableMemory == 0 && this.free.isEmpty()) && !this.waiters.isEmpty()) + this.waiters.peekFirst().signal(); + } + /** * Allocate a buffer. If buffer allocation fails (e.g. because of OOM) then return the size count back to * available memory and signal the next waiter if it exists. @@ -222,16 +234,24 @@ private ByteBuffer safeAllocateByteBuffer(int size) { error = false; return buffer; } finally { - if (error) { - this.lock.lock(); - try { - this.nonPooledAvailableMemory += size; - if (!this.waiters.isEmpty()) - this.waiters.peekFirst().signal(); - } finally { - this.lock.unlock(); - } - } + if (error) + releaseReservedBytes(size); + } + } + + /** + * Return previously-reserved non-pooled bytes to the pool and signal the next + * waiter. Acquires {@link #lock} internally. Used by callers + * that reserve memory and then need to roll back the reservation (e.g., upon errors). + */ + protected void releaseReservedBytes(long bytes) { + this.lock.lock(); + try { + this.nonPooledAvailableMemory += bytes; + if (!this.waiters.isEmpty()) + this.waiters.peekFirst().signal(); + } finally { + this.lock.unlock(); } } @@ -242,9 +262,9 @@ protected ByteBuffer allocateByteBuffer(int size) { /** * Attempt to ensure we have at least the requested number of bytes of memory for allocation by deallocating pooled - * buffers (if needed) + * buffers (if needed). Must be called with {@link #lock} held. */ - private void freeUp(int size) { + protected void freeUp(int size) { while (!this.free.isEmpty() && this.nonPooledAvailableMemory < size) this.nonPooledAvailableMemory += this.free.pollLast().capacity(); } diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java new file mode 100644 index 0000000000000..96a4e6e0b36e4 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java @@ -0,0 +1,200 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.kafka.clients.producer.internals; + +import org.apache.kafka.clients.producer.BufferExhaustedException; +import org.apache.kafka.common.KafkaException; +import org.apache.kafka.common.metrics.Metrics; +import org.apache.kafka.common.utils.Time; + +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.locks.Condition; + +/** + * A {@link BufferPool} dedicated to chunk-sized buffer reuse. The chunk size is passed to the + * parent as its {@link #poolableSize()}. + *

+ * Adds {@link #allocateChunks(int, long)} for callers that need to acquire multiple + * chunks at once. The implementation is atomic across the K-chunk request: + * a single lock acquisition for the whole reservation, a single {@link Condition} + * added to {@link #waiters} so the request is FIFO-fair against single-chunk acquirers, and + * any failure (timeout, close, OOM) refunds the entire reservation atomically before the + * exception propagates. No partial chunk holds are visible outside the lock during the wait. + */ +public class ChunkedBufferPool extends BufferPool { + + public ChunkedBufferPool(long memory, int chunkSize, Metrics metrics, Time time, String metricGrpName) { + super(memory, chunkSize, metrics, time, metricGrpName); + } + + /** + * Allocate {@code ceil(totalSize / chunkSize)} chunk-sized buffers atomically. + * This follows the same principle as the allocate in the parent BufferPool class: if there is + * memory available, allocate the requested number of chunks. If not, wait for memory. + * If it needs to wait, it blocks up to {@code maxTimeToBlockMs} for the whole request, + * queued on {@link #waiters} alongside single-chunk acquirers (FIFO). + *

+ * On any failure path every byte reserved during the call is returned to the pool + * and the next waiter is signaled. No partial chunk holds are visible to other threads during the wait. + * The reservation is tracked as bytes against {@link #nonPooledAvailableMemory} plus chunks polled from {@link #free}. + * + * @param totalSize minimum total bytes of capacity required across the returned chunks + * @param maxTimeToBlockMs maximum time in milliseconds to block waiting for memory + * @return list of {@code ceil(totalSize / chunkSize)} {@code ByteBuffer}s, each of capacity + * {@code chunkSize} + * @throws InterruptedException if interrupted while waiting + * @throws IllegalArgumentException if {@code totalSize <= 0} or {@code totalSize > totalMemory()} + * @throws BufferExhaustedException if the request can't be satisfied within {@code maxTimeToBlockMs} + * @throws KafkaException if the pool is closed during the wait + */ + public List allocateChunks(int totalSize, long maxTimeToBlockMs) throws InterruptedException { + if (totalSize <= 0) + throw new IllegalArgumentException("totalSize must be positive: " + totalSize); + if (totalSize > totalMemory()) + throw new IllegalArgumentException("Attempt to allocate " + totalSize + + " bytes across chunks, but there is a hard limit of " + + totalMemory() + " on memory allocations."); + + int chunkSize = poolableSize(); + int numChunks = (int) (((long) totalSize + chunkSize - 1L) / chunkSize); + long memoryRequired = (long) numChunks * chunkSize; + + // Chunks pulled from the free list during the reservation. Remaining bytes + // (memoryRequired - pooled.size() * chunkSize) are reserved against + // nonPooledAvailableMemory and materialized as raw allocations after the lock is + // released. + List pooled = new ArrayList<>(numChunks); + + lock.lock(); + if (this.closed) { + lock.unlock(); + throw new KafkaException("Producer closed while allocating memory"); + } + try { + long freeListBytes = (long) free.size() * chunkSize; + if (this.nonPooledAvailableMemory + freeListBytes >= memoryRequired) { + // Enough memory available to allocate the chunks needed + while (pooled.size() < numChunks && !free.isEmpty()) + pooled.add(free.pollFirst()); + long remainingBytes = memoryRequired - (long) pooled.size() * chunkSize; + if (remainingBytes > 0) { + // remainingBytes is bounded by memoryRequired ≤ totalMemory; the entry check + // rejects requests larger than that, so the int cast is safe in practice. + freeUp((int) remainingBytes); + this.nonPooledAvailableMemory -= remainingBytes; + } + } else { + // Not enough memory available to allocate the chunks needed, so we need to wait for memory. + // Same as in BufferPool.allocate, but wait to acquire the memory needed for all the chunks. + // A single Condition is added to the waiter's list to ensure FIFO fairness at the request level. + // + // `accumulated` tracks bytes drawn from nonPooledAvailableMemory only (pool chunks + // already taken live in `pooled`). Matches BufferPool.allocate's semantics: on + // failure, `accumulated` is exactly the amount to refund; on success it is reset to 0. + long accumulated = 0; + Condition moreMemory = lock.newCondition(); + try { + long remainingTimeToBlockNs = TimeUnit.MILLISECONDS.toNanos(maxTimeToBlockMs); + waiters.addLast(moreMemory); + while ((long) pooled.size() * chunkSize + accumulated < memoryRequired) { + long startWaitNs = time.nanoseconds(); + long timeNs; + boolean waitingTimeElapsed; + try { + waitingTimeElapsed = !moreMemory.await(remainingTimeToBlockNs, TimeUnit.NANOSECONDS); + } finally { + long endWaitNs = time.nanoseconds(); + timeNs = Math.max(0L, endWaitNs - startWaitNs); + recordWaitTime(timeNs); + } + + if (this.closed) + throw new KafkaException("Producer closed while allocating memory"); + + if (waitingTimeElapsed) { + throw new BufferExhaustedException("Failed to allocate " + memoryRequired + + " bytes (" + numChunks + " chunks of " + chunkSize + + ") within the configured max blocking time " + maxTimeToBlockMs + + " ms. Total memory: " + totalMemory() + " bytes. Available memory: " + + availableMemory() + " bytes."); + } + + remainingTimeToBlockNs -= timeNs; + + // Take pool chunks first, then reserve non-pool bytes for the remainder. + // Pool chunks don't bump `accumulated` because their refund path is + // `free.addFirst` in the outer catch, not nonPool. + while (pooled.size() < numChunks + && (long) (pooled.size() + 1) * chunkSize + accumulated <= memoryRequired + && !free.isEmpty()) { + pooled.add(free.pollFirst()); + } + long stillNeeded = memoryRequired - (long) pooled.size() * chunkSize - accumulated; + if (stillNeeded > 0) { + freeUp((int) stillNeeded); + long got = Math.min(stillNeeded, this.nonPooledAvailableMemory); + this.nonPooledAvailableMemory -= got; + accumulated += got; + } + } + // Loop terminated normally — clear the rollback marker (mirrors BufferPool.allocate). + accumulated = 0; + } finally { + // On failure (timeout / close / interrupt), refund the non-pool bytes taken. + // Pool chunks already in `pooled` are returned to `free` separately by the + // outer catch. + this.nonPooledAvailableMemory += accumulated; + waiters.remove(moreMemory); + } + } + } catch (RuntimeException | InterruptedException e) { + // Any chunks taken from the pool must be returned. + for (ByteBuffer chunk : pooled) + free.addFirst(chunk); + throw e; + } finally { + try { + signalNextWaiterIfMemoryAvailable(); + } finally { + lock.unlock(); + } + } + + // Allocate raw chunks for the reserved non-pooled portion outside the lock. On OOM, + // refund all reserved bytes (memoryRequired) and let the next waiter try. + int chunksStillNeeded = numChunks - pooled.size(); + List result = new ArrayList<>(numChunks); + result.addAll(pooled); + boolean error = true; + try { + for (int i = 0; i < chunksStillNeeded; i++) + result.add(allocateByteBuffer(chunkSize)); + error = false; + return result; + } finally { + if (error) { + // The pooled buffers we already drained are also lost on this path; mirror + // BufferPool.safeAllocateByteBuffer's behaviour (the bytes return, the ByteBuffer + // instances become garbage). + releaseReservedBytes(memoryRequired); + } + } + } +} diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java new file mode 100644 index 0000000000000..ccf1d3d7374c9 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java @@ -0,0 +1,246 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.kafka.clients.producer.internals; + +import org.apache.kafka.common.utils.ByteBufferOutputStream; + +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.List; + +/** + * A {@link ByteBufferOutputStream} backed by a linked list of fixed-size chunks instead of a + * single re-allocated buffer. + *

+ * This stream does not grow on its own. Chunks are supplied by the caller — initial + * chunks via the constructor, additional chunks via {@link #addBuffers(List)}. When a write's + * size exceeds the stream's {@link #remaining()} (sum of free bytes across all attached + * chunks), {@link IllegalStateException} is thrown. The caller is responsible for attaching + * additional chunks before any such write. + *

+ * Automatic mid-write growth (allocating a chunk from inside {@code write()}) will be added + * together with compression support in a follow-up: compressor buffering can cause a record's + * write to exceed its reservation, which the caller can't predict in advance. Until then, the + * write path stays allocation-free and deterministic. + *

+ * When {@link #buffer()} is called (typically at batch close), all chunks are flattened into a + * single contiguous ByteBuffer — exactly one copy operation. + */ +public class ChunkedByteBufferOutputStream extends ByteBufferOutputStream { + + private final List chunks; + private final int chunkSize; + private final BufferPool pool; + private ByteBuffer currentChunk; + private int currentChunkIndex; + private ByteBuffer flattenedBuffer; + private boolean dirty; + + /** + * Constructs a chunked output stream backed by the given pre-allocated chunks. Ownership of + * {@code initialChunks} transfers to this stream — they will be returned to the pool via + * {@link #deallocate()}. + * + * @param initialChunks pre-allocated chunks; must be non-empty; each chunk's capacity must + * equal {@code chunkSize} + * @param chunkSize the size of each chunk in bytes + * @param pool the buffer pool used for deallocation + */ + @SuppressWarnings("this-escape") + public ChunkedByteBufferOutputStream(List initialChunks, int chunkSize, BufferPool pool) { + super(initialChunks.get(0)); + this.chunkSize = chunkSize; + this.pool = pool; + this.chunks = new ArrayList<>(initialChunks); + this.currentChunk = this.chunks.get(0); + this.currentChunkIndex = 0; + this.dirty = true; + } + + @Override + public void write(int b) { + ensureChunkCapacity(1); + currentChunk.put((byte) b); + dirty = true; + } + + @Override + public void write(byte[] bytes, int off, int len) { + while (len > 0) { + ensureChunkCapacity(1); + int toWrite = Math.min(len, currentChunk.remaining()); + currentChunk.put(bytes, off, toWrite); + off += toWrite; + len -= toWrite; + } + dirty = true; + } + + @Override + public void write(ByteBuffer sourceBuffer) { + while (sourceBuffer.hasRemaining()) { + ensureChunkCapacity(1); + int toWrite = Math.min(sourceBuffer.remaining(), currentChunk.remaining()); + int oldLimit = sourceBuffer.limit(); + sourceBuffer.limit(sourceBuffer.position() + toWrite); + currentChunk.put(sourceBuffer); + sourceBuffer.limit(oldLimit); + } + dirty = true; + } + + private void ensureChunkCapacity(int needed) { + if (currentChunk.remaining() < needed) { + advanceToNextChunk(); + } + } + + /** + * Advances {@code currentChunk} to the next pre-supplied chunk. The caller is responsible + * for obtaining additional chunks (e.g. {@code ChunkedBufferPool.allocateChunks}) and + * attaching them via {@link #addBuffers(List)} before any write that would exceed + * {@link #remaining()}. The KIP's mid-record growth path (non-blocking pool then direct + * heap) lands together with compression support. + */ + private void advanceToNextChunk() { + if (currentChunkIndex + 1 >= chunks.size()) { + throw new IllegalStateException( + "No more chunks available; the write exceeded the stream's remaining capacity. " + + "Caller must obtain additional chunks (e.g. from ChunkedBufferPool) and attach them " + + "via addBuffers before any write whose size exceeds remaining()"); + } + currentChunkIndex++; + currentChunk = chunks.get(currentChunkIndex); + } + + /** + * Appends pre-allocated chunks to this stream. Ownership of {@code newChunks} transfers to + * the stream; they will be returned to the pool via {@link #deallocate()}. + */ + public void addBuffers(List newChunks) { + chunks.addAll(newChunks); + } + + @Override + public ByteBuffer buffer() { + if (flattenedBuffer != null && !dirty) { + return flattenedBuffer; + } + int totalSize = 0; + for (ByteBuffer chunk : chunks) { + totalSize += chunk.position(); + } + flattenedBuffer = ByteBuffer.allocate(totalSize); + for (ByteBuffer chunk : chunks) { + int chunkPos = chunk.position(); + chunk.flip(); + flattenedBuffer.put(chunk); + chunk.limit(chunk.capacity()); + chunk.position(chunkPos); + } + dirty = false; + return flattenedBuffer; + } + + /** + * Total bytes written across all chunks. + */ + @Override + public int position() { + int total = 0; + for (ByteBuffer chunk : chunks) { + total += chunk.position(); + } + return total; + } + + /** + * Sets the write position, walking across pre-supplied chunks if the requested position + * exceeds the first chunk's capacity. Only valid before any write. + */ + @Override + public void position(int position) { + if (currentChunkIndex != 0 || currentChunk.position() != 0) { + throw new IllegalStateException("position() can only be called before any writes"); + } + int remaining = position; + int idx = 0; + while (remaining > 0 && idx < chunks.size()) { + ByteBuffer chunk = chunks.get(idx); + int take = Math.min(remaining, chunk.capacity()); + chunk.position(take); + remaining -= take; + if (remaining > 0) + idx++; + } + if (remaining > 0) { + throw new IllegalArgumentException("position " + position + + " exceeds total pre-allocated capacity"); + } + currentChunkIndex = idx; + currentChunk = chunks.get(idx); + dirty = true; + } + + /** + * Total bytes available across the current chunk and every queued (not-yet-active) chunk. + */ + @Override + public int remaining() { + int total = currentChunk.remaining(); + for (int i = currentChunkIndex + 1; i < chunks.size(); i++) + total += chunks.get(i).remaining(); + return total; + } + + @Override + public int limit() { + return Integer.MAX_VALUE; + } + + @Override + public int initialCapacity() { + return chunks.isEmpty() ? chunkSize : chunks.get(0).capacity(); + } + + @Override + public void ensureRemaining(int remainingBytesRequired) { + // The stream advances one chunk at a time — the most any single ensureChunkCapacity + // call can guarantee is `chunkSize` of contiguous-by-chunk space. Callers needing more + // than that are expected to attach chunks via addBuffers before writing; write(byte[]) + // itself loops across chunks so contiguous capacity isn't required. + ensureChunkCapacity(Math.min(remainingBytesRequired, chunkSize)); + } + + /** + * Returns all pool-allocated chunks to the buffer pool. + */ + public void deallocate(BufferPool pool) { + if (pool != null) { + for (ByteBuffer chunk : chunks) { + pool.deallocate(chunk); + } + } + chunks.clear(); + currentChunk = null; + flattenedBuffer = null; + } + + public void deallocate() { + deallocate(pool); + } +} diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java new file mode 100644 index 0000000000000..cac18c608d195 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java @@ -0,0 +1,336 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.kafka.clients.producer.internals; + +import org.apache.kafka.clients.producer.BufferExhaustedException; +import org.apache.kafka.clients.producer.RecordMetadata; +import org.apache.kafka.common.Cluster; +import org.apache.kafka.common.compress.Compression; +import org.apache.kafka.common.header.Header; +import org.apache.kafka.common.metrics.Metrics; +import org.apache.kafka.common.record.TimestampType; +import org.apache.kafka.common.record.internal.AbstractRecords; +import org.apache.kafka.common.record.internal.CompressionType; +import org.apache.kafka.common.record.internal.MemoryRecordsBuilder; +import org.apache.kafka.common.record.internal.Record; +import org.apache.kafka.common.record.internal.RecordBatch; +import org.apache.kafka.common.utils.Time; +import org.apache.kafka.common.utils.internals.LogContext; + +import java.nio.ByteBuffer; +import java.util.ArrayDeque; +import java.util.Deque; +import java.util.List; + +/** + * A {@link RecordAccumulator} variant that uses chunked buffer allocation: each new batch + * pre-allocates {@code ceil(recordSize / chunkSize)} chunks from a {@link ChunkedBufferPool}, + * and grows mid-batch by attaching additional chunks via + * {@link ProducerBatch#addBuffers(List)} when later records would overflow. + *

+ * This is the "dynamic" buffer.memory allocation strategy: memory consumption scales with + * actual data buffered, not with {@code active_partition_count × batch.size}. + *

+ * PR1 scope: uncompressed only. Batch creation throws + * {@link UnsupportedOperationException} if compression is requested. The mid-batch flow + * matches the KIP: try a non-blocking pool acquire for the extension chunks; on + * {@link BufferExhaustedException}, close the existing batch (so the sender can drain it) + * and let the new record flow through the first-record path, which blocks on the pool for + * a fresh batch's chunks up to {@code max.block.ms}. + *

+ * TODO: add the new buffer.memory.allocation.strategy config and instantiate this class only when set to "dynamic" + * TODO: add support for compressed data. Throws UnsupportedOperationException for now, + * and IllegalStateException if overflow detected. Follow-ups will support compressed data and growth. + */ +public class ChunkedRecordAccumulator extends RecordAccumulator { + + private final ChunkedBufferPool chunkedFree; + + public ChunkedRecordAccumulator(LogContext logContext, + int batchSize, + Compression compression, + int lingerMs, + long retryBackoffMs, + long retryBackoffMaxMs, + int deliveryTimeoutMs, + PartitionerConfig partitionerConfig, + Metrics metrics, + String metricGrpName, + Time time, + TransactionManager transactionManager, + ChunkedBufferPool bufferPool) { + super(logContext, batchSize, compression, lingerMs, retryBackoffMs, retryBackoffMaxMs, + deliveryTimeoutMs, partitionerConfig, metrics, metricGrpName, time, transactionManager, bufferPool); + this.chunkedFree = bufferPool; + } + + public ChunkedRecordAccumulator(LogContext logContext, + int batchSize, + Compression compression, + int lingerMs, + long retryBackoffMs, + long retryBackoffMaxMs, + int deliveryTimeoutMs, + Metrics metrics, + String metricGrpName, + Time time, + TransactionManager transactionManager, + ChunkedBufferPool bufferPool) { + this(logContext, batchSize, compression, lingerMs, retryBackoffMs, retryBackoffMaxMs, + deliveryTimeoutMs, new PartitionerConfig(), metrics, metricGrpName, time, transactionManager, + bufferPool); + } + + @Override + public RecordAppendResult append(String topic, + int partition, + long timestamp, + byte[] key, + byte[] value, + Header[] headers, + AppendCallbacks callbacks, + long maxTimeToBlock, + long nowMs, + Cluster cluster) throws InterruptedException { + TopicInfo topicInfo = topicInfoMap.computeIfAbsent(topic, + k -> new TopicInfo(createBuiltInPartitioner(logContext, k, batchSize, partitionerRackAware, rack))); + + appendsInProgress.incrementAndGet(); + ChunkedByteBufferOutputStream bufferStream = null; + List extensionChunks = null; + int extensionBytes = 0; + if (headers == null) headers = Record.EMPTY_HEADERS; + try { + while (true) { + final BuiltInPartitioner.StickyPartitionInfo partitionInfo; + final int effectivePartition; + if (partition == RecordMetadata.UNKNOWN_PARTITION) { + partitionInfo = topicInfo.builtInPartitioner.peekCurrentPartitionInfo(cluster); + effectivePartition = partitionInfo.partition(); + } else { + partitionInfo = null; + effectivePartition = partition; + } + setPartition(callbacks, effectivePartition); + + Deque dq = topicInfo.batches.computeIfAbsent(effectivePartition, k -> new ArrayDeque<>()); + synchronized (dq) { + if (partitionChanged(topic, topicInfo, partitionInfo, dq, nowMs, cluster)) + continue; + + // Probe the tail directly. If the existing batch has logical room but the + // underlying stream lacks physical capacity, capture the gap and fall through + // to allocate the extension outside the deque lock. Otherwise, attempt the + // append against the existing batch. + ProducerBatch tail = dq.peekLast(); + int gap = tail == null ? 0 : tail.extensionBytesNeeded(timestamp, key, value, headers); + if (gap > 0) { + extensionBytes = gap; + } else { + extensionBytes = 0; // clear any stale value carried from a prior iteration + RecordAppendResult appendResult = tryAppend(timestamp, key, value, headers, callbacks, dq, nowMs); + if (appendResult != null) { + boolean enableSwitch = allBatchesFull(dq); + topicInfo.builtInPartitioner.updatePartitionInfo(partitionInfo, appendResult.appendedBytes, cluster, enableSwitch); + return appendResult; + } + } + } + + if (extensionBytes > 0 && extensionChunks == null) { + // Mid-batch extension: non-blocking attempt only. The producer thread already + // holds chunks for the open batch; blocking here risks deadlock with the Sender + // thread (which frees pool memory by completing batches). If the pool is + // exhausted, we close the existing batch (so it becomes drainable) and fall + // through to the first-record blocking path on the next loop iteration. + try { + extensionChunks = chunkedFree.allocateChunks(extensionBytes, 0L); + } catch (BufferExhaustedException e) { + log.trace("Pool exhausted while extending batch for topic {} partition {}; closing existing batch", + topic, effectivePartition); + synchronized (dq) { + ProducerBatch last = dq.peekLast(); + if (last != null && last.isWritable()) { + last.closeForRecordAppends(); + } + } + extensionBytes = 0; + continue; + } + nowMs = time.milliseconds(); + } else if (extensionBytes == 0 && bufferStream == null) { + // First-record path: block on the pool for enough chunks to fit this record. + // Temporary while compressed data not supported. + // TODO: drop this throw and route through the + // compressed path (which adds the mid-record growth fallback for compressor overshoot). + if (compression.type() != CompressionType.NONE) { + throw new UnsupportedOperationException( + "Compression is not implemented yet for the dynamic buffer.memory allocation strategy"); + } + int size = AbstractRecords.estimateSizeInBytesUpperBound( + RecordBatch.CURRENT_MAGIC_VALUE, compression.type(), key, value, headers); + log.trace("Allocating {} byte chunked buffer ({} byte chunks) for topic {} partition {} with remaining timeout {}ms", + size, chunkedFree.poolableSize(), topic, effectivePartition, maxTimeToBlock); + List initialChunks = chunkedFree.allocateChunks(size, maxTimeToBlock); + nowMs = time.milliseconds(); + bufferStream = new ChunkedByteBufferOutputStream(initialChunks, chunkedFree.poolableSize(), chunkedFree); + } + + synchronized (dq) { + if (partitionChanged(topic, topicInfo, partitionInfo, dq, nowMs, cluster)) + continue; + + if (extensionChunks != null) { + ProducerBatch last = dq.peekLast(); + // The off-lock allocateChunks window allows the original tail batch to be + // drained and a split batch (plain ByteBufferOutputStream) to be addFirst'd + // into an emptied deque — see splitAndReenqueue. The instanceof guard below + // is load-bearing: addBuffers is only defined on ChunkedByteBufferOutputStream; + // the else branch returns the chunks to the pool to avoid leaking them. + if (last != null && last.isWritable() + && last.bufferStream() instanceof ChunkedByteBufferOutputStream) { + ((ChunkedByteBufferOutputStream) last.bufferStream()).addBuffers(extensionChunks); + extensionChunks = null; + extensionBytes = 0; + // Re-probe after attach: between our first-lock probe and now we were off-lock, + // and another appender (or a tail-replacement) may have consumed capacity that + // our extension allocation was sized against. If a gap remains, loop to acquire + // more chunks against the now-current state. Without this re-probe, a + // tryAppend that passes hasRoomFor (writeLimit) but exceeds the stream's + // physical remaining() trips IllegalStateException in + // ChunkedByteBufferOutputStream.advanceToNextChunk. Termination: writeLimit is + // bounded and fixed; once hasRoomFor turns false, extensionBytesNeeded returns 0 + // and we fall through to new-batch creation. Regression test: + // testConcurrentExtensionRaceDoesNotOverflowChunkedStream. + int recheck = last.extensionBytesNeeded(timestamp, key, value, headers); + if (recheck > 0) { + extensionBytes = recheck; + continue; + } + RecordAppendResult retryResult = tryAppend(timestamp, key, value, headers, callbacks, dq, nowMs); + if (retryResult != null) { + boolean enableSwitch = allBatchesFull(dq); + topicInfo.builtInPartitioner.updatePartitionInfo(partitionInfo, retryResult.appendedBytes, cluster, enableSwitch); + return retryResult; + } + // Batch became full via writeLimit while we were off-lock — fall through + // to new-batch creation on the next iteration. + continue; + } + // Tail is gone, closed, or non-chunked (e.g., a split batch). Return chunks to pool. + for (ByteBuffer chunk : extensionChunks) + chunkedFree.deallocate(chunk); + extensionChunks = null; + extensionBytes = 0; + continue; + } + + // bufferStream is non-null here (first-record path). Before creating a new + // batch with it, re-probe the tail: a concurrent appender may have created a + // chunked batch we should extend instead. If so, release the oversized + // bufferStream and let the next iteration take the extension path. + ProducerBatch tail = dq.peekLast(); + if (tail != null && tail.isWritable() + && tail.bufferStream() instanceof ChunkedByteBufferOutputStream + && tail.extensionBytesNeeded(timestamp, key, value, headers) > 0) { + bufferStream.deallocate(); + bufferStream = null; + continue; + } + + int firstRecordSize = AbstractRecords.estimateSizeInBytesUpperBound( + RecordBatch.CURRENT_MAGIC_VALUE, compression.type(), key, value, headers); + final ChunkedByteBufferOutputStream batchStream = bufferStream; + RecordAppendResult appendResult = appendNewBatch(topic, effectivePartition, dq, timestamp, key, value, headers, callbacks, + () -> chunkedRecordsBuilder(batchStream, firstRecordSize), nowMs); + if (appendResult.newBatchCreated) + bufferStream = null; + boolean enableSwitch = allBatchesFull(dq); + topicInfo.builtInPartitioner.updatePartitionInfo(partitionInfo, appendResult.appendedBytes, cluster, enableSwitch); + return appendResult; + } + } + } finally { + if (bufferStream != null) + bufferStream.deallocate(); + if (extensionChunks != null) { + for (ByteBuffer chunk : extensionChunks) + chunkedFree.deallocate(chunk); + } + appendsInProgress.decrementAndGet(); + } + } + + /** + * Build a {@link MemoryRecordsBuilder} backed by the chunked stream. + * {@code writeLimit} is set to {@code max(batchSize, firstRecordSize)} — the same logical + * cap the legacy path uses. Physical chunk capacity grows on demand via + * {@link ChunkedByteBufferOutputStream#addBuffers(List)}; {@code writeLimit} bounds when a + * batch is considered logically full and a new one must be started. + */ + private MemoryRecordsBuilder chunkedRecordsBuilder(ChunkedByteBufferOutputStream bufferStream, + int firstRecordSize) { + int writeLimit = Math.max(batchSize, firstRecordSize); + return new MemoryRecordsBuilder(bufferStream, RecordBatch.CURRENT_MAGIC_VALUE, compression, + TimestampType.CREATE_TIME, 0L, RecordBatch.NO_TIMESTAMP, RecordBatch.NO_PRODUCER_ID, + RecordBatch.NO_PRODUCER_EPOCH, RecordBatch.NO_SEQUENCE, false, false, + RecordBatch.NO_PARTITION_LEADER_EPOCH, writeLimit); + } + + /** + * Chunked-aware deallocation for {@link ProducerBatch}. + *

+ * Intercepts the {@link #isInflight()} branch of the parent's + * {@link RecordAccumulator#deallocate(ProducerBatch)} when the batch is backed by a + * {@link ChunkedByteBufferOutputStream}. The parent's defensive path credits only + * {@code initialCapacity} (one {@code chunkSize}) back to the pool and throws — that's the + * KAFKA-19012-aware safety net for the legacy single-buffer path, where the in-flight + * buffer is the pooled buffer. A chunked batch holds K chunks but the in-flight + * network bytes live in a separate flattened heap buffer (see + * {@link ChunkedByteBufferOutputStream#buffer()}), so it is safe to return all K chunks + * to the pool here. The {@code IllegalStateException} is still propagated to preserve the + * contract that deallocating an inflight batch is unexpected upstream. + */ + @Override + public void deallocate(ProducerBatch batch) { + if (!batch.isSplitBatch() + && !batch.isBufferDeallocated() + && batch.isInflight() + && batch.bufferStream() instanceof ChunkedByteBufferOutputStream) { + batch.markBufferDeallocated(); + deallocateBatchBuffer(batch); + throw new IllegalStateException("Attempting to deallocate a batch that is inflight. Batch is " + batch); + } + super.deallocate(batch); + } + + /** + * Route chunked batch deallocation through the stream rather than the legacy + * single-buffer pool path. The flattened buffer returned by {@code batch.buffer()} is a + * fresh heap allocation and must NOT be passed to the pool — instead, the stream owns + * the chunk-sized buffers and returns them to the pool. + */ + @Override + protected void deallocateBatchBuffer(ProducerBatch batch) { + if (batch.bufferStream() instanceof ChunkedByteBufferOutputStream) { + ((ChunkedByteBufferOutputStream) batch.bufferStream()).deallocate(chunkedFree); + } else { + // Split batches and any non-chunked batches go through the default deallocation. + super.deallocateBatchBuffer(batch); + } + } +} diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ProducerBatch.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ProducerBatch.java index 14d0d6ef6334d..a7e47583d310b 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ProducerBatch.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ProducerBatch.java @@ -31,6 +31,7 @@ import org.apache.kafka.common.record.internal.Record; import org.apache.kafka.common.record.internal.RecordBatch; import org.apache.kafka.common.requests.ProduceResponse; +import org.apache.kafka.common.utils.ByteBufferOutputStream; import org.apache.kafka.common.utils.ProducerIdAndEpoch; import org.apache.kafka.common.utils.Time; @@ -139,6 +140,38 @@ boolean hasLeaderChangedForTheOngoingRetry() { } + /** + * Bytes of physical buffer this batch needs before {@code tryAppend} could accept the given + * record. Returns 0 when the record fits (either logically full, or stream has room). + * Positive when {@code hasRoomFor} allows the record but the underlying stream lacks + * physical capacity — used by the dynamic strategy to allocate exactly the missing bytes + * before retrying. + *

+ * For batches with the default {@link ByteBufferOutputStream} (legacy path), the stream's + * {@code remaining()} is large enough that this almost always returns 0. + */ + int extensionBytesNeeded(long timestamp, byte[] key, byte[] value, Header[] headers) { + if (recordCount == 0) + return 0; + if (!recordsBuilder.hasRoomFor(timestamp, key, value, headers)) + return 0; + // Upper-bound estimate is fine: over-counting by header bytes only over-allocates by a + // small fraction of one chunk. Under-estimation would be unsafe. + int recordSize = AbstractRecords.estimateSizeInBytesUpperBound( + magic(), recordsBuilder.compression().type(), key, value, headers); + int gap = recordSize - recordsBuilder.bufferStream().remaining(); + return Math.max(0, gap); + } + + /** + * The batch's underlying output stream. Exposed for the dynamic strategy to attach + * extension chunks and route deallocation through the chunked stream rather than the + * standard {@code BufferPool.deallocate(ByteBuffer, int)} path. + */ + ByteBufferOutputStream bufferStream() { + return recordsBuilder.bufferStream(); + } + /** * Append the record to the current record set and return the relative offset within that record set * diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/RecordAccumulator.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/RecordAccumulator.java index b1eced17e4f25..29ddd07a489bc 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/RecordAccumulator.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/RecordAccumulator.java @@ -58,6 +58,7 @@ import java.util.Set; import java.util.concurrent.ConcurrentMap; import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.Supplier; /** * This class acts as a queue that accumulates records into {@link MemoryRecords} @@ -68,23 +69,23 @@ */ public class RecordAccumulator { - private final LogContext logContext; - private final Logger log; - private volatile boolean closed; + protected final LogContext logContext; + protected final Logger log; + protected volatile boolean closed; private final AtomicInteger flushesInProgress; - private final AtomicInteger appendsInProgress; - private final int batchSize; - private final Compression compression; + protected final AtomicInteger appendsInProgress; + protected final int batchSize; + protected final Compression compression; private final int lingerMs; private final ExponentialBackoff retryBackoff; private final int deliveryTimeoutMs; private final long partitionAvailabilityTimeoutMs; // latency threshold for marking partition temporary unavailable - private final boolean partitionerRackAware; - private final String rack; + protected final boolean partitionerRackAware; + protected final String rack; private final boolean enableAdaptivePartitioning; private final BufferPool free; - private final Time time; - private final ConcurrentMap topicInfoMap = new CopyOnWriteMap<>(); + protected final Time time; + protected final ConcurrentMap topicInfoMap = new CopyOnWriteMap<>(); private final ConcurrentMap nodeStats = new CopyOnWriteMap<>(); private final IncompleteBatches incomplete; // The following variables are only accessed by the sender thread, so we don't need to protect them. @@ -217,7 +218,7 @@ private void registerMetrics(Metrics metrics, String metricGrpName) { (config, now) -> free.availableMemory()); } - private void setPartition(AppendCallbacks callbacks, int partition) { + protected void setPartition(AppendCallbacks callbacks, int partition) { if (callbacks != null) callbacks.setPartition(partition); } @@ -234,7 +235,7 @@ private void setPartition(AppendCallbacks callbacks, int partition) { * @return 'true' if partition changed and we need to get new partition info and retry, * 'false' otherwise */ - private boolean partitionChanged(String topic, + protected boolean partitionChanged(String topic, TopicInfo topicInfo, BuiltInPartitioner.StickyPartitionInfo partitionInfo, Deque deque, long nowMs, @@ -347,7 +348,10 @@ public RecordAppendResult append(String topic, if (partitionChanged(topic, topicInfo, partitionInfo, dq, nowMs, cluster)) continue; - RecordAppendResult appendResult = appendNewBatch(topic, effectivePartition, dq, timestamp, key, value, headers, callbacks, buffer, nowMs); + final ByteBuffer batchBuffer = buffer; + RecordAppendResult appendResult = appendNewBatch(topic, effectivePartition, dq, timestamp, key, value, headers, callbacks, + () -> MemoryRecords.builder(batchBuffer, RecordBatch.CURRENT_MAGIC_VALUE, compression, TimestampType.CREATE_TIME, 0L), + nowMs); // Set buffer to null, so that deallocate doesn't return it back to free pool, since it's used in the batch. if (appendResult.newBatchCreated) buffer = null; @@ -374,10 +378,13 @@ public RecordAppendResult append(String topic, * @param value The value for the record * @param headers the Headers for the record * @param callbacks The callbacks to execute - * @param buffer The buffer for the new batch + * @param recordsBuilderSupplier Supplies the {@link MemoryRecordsBuilder} for the new + * batch. Invoked lazily, only when this method is committed to creating a new + * batch (i.e. after the in-lock concurrent-append race check has lost). The + * chunked subclass passes a supplier that produces a stream-backed builder. * @param nowMs The current time, in milliseconds */ - private RecordAppendResult appendNewBatch(String topic, + protected RecordAppendResult appendNewBatch(String topic, int partition, Deque dq, long timestamp, @@ -385,7 +392,7 @@ private RecordAppendResult appendNewBatch(String topic, byte[] value, Header[] headers, AppendCallbacks callbacks, - ByteBuffer buffer, + Supplier recordsBuilderSupplier, long nowMs) { assert partition != RecordMetadata.UNKNOWN_PARTITION; @@ -395,7 +402,7 @@ private RecordAppendResult appendNewBatch(String topic, return appendResult; } - MemoryRecordsBuilder recordsBuilder = recordsBuilder(buffer); + MemoryRecordsBuilder recordsBuilder = recordsBuilderSupplier.get(); ProducerBatch batch = new ProducerBatch(new TopicPartition(topic, partition), recordsBuilder, nowMs); FutureRecordMetadata future = Objects.requireNonNull(batch.tryAppend(timestamp, key, value, headers, callbacks, nowMs)); @@ -406,14 +413,10 @@ private RecordAppendResult appendNewBatch(String topic, return new RecordAppendResult(future, dq.size() > 1 || batch.isFull(), true, batch.estimatedSizeInBytes()); } - private MemoryRecordsBuilder recordsBuilder(ByteBuffer buffer) { - return MemoryRecords.builder(buffer, RecordBatch.CURRENT_MAGIC_VALUE, compression, TimestampType.CREATE_TIME, 0L); - } - /** * Check if all batches in the queue are full. */ - private boolean allBatchesFull(Deque deque) { + protected boolean allBatchesFull(Deque deque) { // Only the last batch may be incomplete, so we just check that. ProducerBatch last = deque.peekLast(); return last == null || last.isFull(); @@ -427,7 +430,7 @@ private boolean allBatchesFull(Deque deque) { * and memory records built) in one of the following cases (whichever comes first): right before send, * if it is expired, or when the producer is closed. */ - private RecordAppendResult tryAppend(long timestamp, byte[] key, byte[] value, Header[] headers, + protected RecordAppendResult tryAppend(long timestamp, byte[] key, byte[] value, Header[] headers, Callback callback, Deque deque, long nowMs) { if (closed) throw new KafkaException("Producer closed while send in progress"); @@ -1058,11 +1061,19 @@ public void deallocate(ProducerBatch batch) { free.deallocate(ByteBuffer.allocate(batch.initialCapacity())); throw new IllegalStateException("Attempting to deallocate a batch that is inflight. Batch is " + batch); } - free.deallocate(batch.buffer(), batch.initialCapacity()); + deallocateBatchBuffer(batch); } } } + /** + * Return the batch's underlying buffer to the pool. + * This default implementation returns the buffer at its initial capacity (single buffer). + */ + protected void deallocateBatchBuffer(ProducerBatch batch) { + free.deallocate(batch.buffer(), batch.initialCapacity()); + } + /** * Remove from the incomplete list but do not free memory yet */ @@ -1298,7 +1309,7 @@ public ReadyCheckResult(Set readyNodes, long nextReadyCheckDelayMs, Set> batches = new CopyOnWriteMap<>(); public final BuiltInPartitioner builtInPartitioner; diff --git a/clients/src/main/java/org/apache/kafka/common/record/internal/MemoryRecordsBuilder.java b/clients/src/main/java/org/apache/kafka/common/record/internal/MemoryRecordsBuilder.java index 3fd89deaba135..47a955cfcf18e 100644 --- a/clients/src/main/java/org/apache/kafka/common/record/internal/MemoryRecordsBuilder.java +++ b/clients/src/main/java/org/apache/kafka/common/record/internal/MemoryRecordsBuilder.java @@ -211,6 +211,15 @@ public int initialCapacity() { return bufferStream.initialCapacity(); } + /** + * The underlying output stream. Exposed so the dynamic-strategy callers can attach + * pre-allocated chunks to a {@code ChunkedByteBufferOutputStream} and route deallocation + * polymorphically. + */ + public ByteBufferOutputStream bufferStream() { + return bufferStream; + } + public double compressionRatio() { return actualCompressionRatio; } diff --git a/gradle/spotbugs-exclude.xml b/gradle/spotbugs-exclude.xml index 6d18691973a45..1a56dc05037c9 100644 --- a/gradle/spotbugs-exclude.xml +++ b/gradle/spotbugs-exclude.xml @@ -127,6 +127,17 @@ For a detailed description of spotbugs bug categories, see https://spotbugs.read + + + + + + @@ -546,6 +557,7 @@ For a detailed description of spotbugs bug categories, see https://spotbugs.read --> + From 203a410dac3ec480aae2a483b26852d378df8513 Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Fri, 5 Jun 2026 12:09:30 -0400 Subject: [PATCH 02/61] tests --- .../internals/ChunkedBufferPoolTest.java | 422 ++++++++++++++++++ .../ChunkedByteBufferOutputStreamTest.java | 171 +++++++ .../ChunkedRecordAccumulatorTest.java | 350 +++++++++++++++ 3 files changed, 943 insertions(+) create mode 100644 clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPoolTest.java create mode 100644 clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java create mode 100644 clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPoolTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPoolTest.java new file mode 100644 index 0000000000000..e0de8ae19e4dc --- /dev/null +++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPoolTest.java @@ -0,0 +1,422 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.kafka.clients.producer.internals; + +import org.apache.kafka.clients.producer.BufferExhaustedException; +import org.apache.kafka.common.metrics.Metrics; +import org.apache.kafka.common.utils.MockTime; + +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; + +import org.apache.kafka.common.KafkaException; + +import java.nio.ByteBuffer; +import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +public class ChunkedBufferPoolTest { + + private final MockTime time = new MockTime(); + private final Metrics metrics = new Metrics(time); + private final String metricGroup = "producer-metrics"; + + @AfterEach + public void teardown() { + metrics.close(); + } + + private ChunkedBufferPool pool(long totalMemory, int chunkSize) { + return new ChunkedBufferPool(totalMemory, chunkSize, metrics, time, metricGroup); + } + + /** Single-chunk request returns a list of one buffer at the chunk size. */ + @Test + public void testAllocateOneChunk() throws Exception { + int chunkSize = 64; + ChunkedBufferPool p = pool(1024, chunkSize); + List chunks = p.allocateChunks(chunkSize, 100); + assertEquals(1, chunks.size()); + assertEquals(chunkSize, chunks.get(0).capacity()); + } + + /** Total size that's not a multiple of chunk size rounds up; returned chunks cover it. */ + @Test + public void testAllocateRoundsUpToChunkBoundary() throws Exception { + int chunkSize = 64; + ChunkedBufferPool p = pool(1024, chunkSize); + // 65 bytes requested → 2 chunks (128 bytes total). + List chunks = p.allocateChunks(65, 100); + assertEquals(2, chunks.size()); + for (ByteBuffer chunk : chunks) + assertEquals(chunkSize, chunk.capacity()); + } + + /** Multi-chunk request returns ceil(totalSize / chunkSize) chunks. */ + @Test + public void testAllocateMultipleChunks() throws Exception { + int chunkSize = 64; + ChunkedBufferPool p = pool(1024, chunkSize); + List chunks = p.allocateChunks(4 * chunkSize, 100); + assertEquals(4, chunks.size()); + } + + /** Pool memory accounting: after allocation, the unallocated portion shrinks by the request. */ + @Test + public void testAvailableMemoryAfterAllocation() throws Exception { + int chunkSize = 64; + long total = 256; + ChunkedBufferPool p = pool(total, chunkSize); + p.allocateChunks(3 * chunkSize, 100); + assertEquals(total - 3L * chunkSize, p.availableMemory()); + } + + /** Returning chunks via deallocate restores pool memory. */ + @Test + public void testDeallocationRestoresMemory() throws Exception { + int chunkSize = 64; + long total = 256; + ChunkedBufferPool p = pool(total, chunkSize); + List chunks = p.allocateChunks(2 * chunkSize, 100); + for (ByteBuffer chunk : chunks) + p.deallocate(chunk); + assertEquals(total, p.availableMemory()); + } + + /** Requesting more than total memory throws IllegalArgumentException. */ + @Test + public void testRejectsRequestExceedingTotalMemory() { + ChunkedBufferPool p = pool(128, 64); + assertThrows(IllegalArgumentException.class, () -> p.allocateChunks(129, 100)); + } + + /** Non-positive totalSize is rejected. */ + @Test + public void testRejectsNonPositiveRequest() { + ChunkedBufferPool p = pool(128, 64); + assertThrows(IllegalArgumentException.class, () -> p.allocateChunks(0, 100)); + assertThrows(IllegalArgumentException.class, () -> p.allocateChunks(-1, 100)); + } + + /** + * When a multi-chunk request can't be fully satisfied within the deadline, chunks already + * acquired during the call must be returned to the pool — the caller must not leak memory. + */ + @Test + public void testRollbackOnPartialFailure() throws Exception { + int chunkSize = 64; + long total = 2 * chunkSize; // only 2 chunks worth of memory + ChunkedBufferPool p = pool(total, chunkSize); + // Reserve one chunk so the pool has only 1 left. + ByteBuffer held = p.allocate(chunkSize, 100); + + // Request 2 chunks with a zero deadline — first chunk fits, second can't, must throw. + assertThrows(BufferExhaustedException.class, () -> p.allocateChunks(2 * chunkSize, 0)); + + // Available memory must reflect only the chunk we deliberately hold; the request's + // first-chunk acquisition was rolled back. + assertEquals(total - chunkSize, p.availableMemory()); + + p.deallocate(held); + assertEquals(total, p.availableMemory()); + } + + /** + * No partial chunk holds during the wait. While a multi-chunk request blocks, the pool's + * {@code availableMemory()} must reflect bytes the waiter has not yet "earned" — i.e., the + * reservation is via the accumulator and not via chunks pulled out of the pool's accounting. + * Distinguishes the atomic implementation from the older loop-with-rollback (which held + * partials off the accounting between iterations). + */ + @Test + public void testNoPartialHoldsDuringWait() throws Exception { + int chunkSize = 64; + long total = 4L * chunkSize; + ChunkedBufferPool p = pool(total, chunkSize); + // Drain the pool down to 1 chunk's worth so a 2-chunk request can't be satisfied immediately. + ByteBuffer h1 = p.allocate(chunkSize, 100); + ByteBuffer h2 = p.allocate(chunkSize, 100); + ByteBuffer h3 = p.allocate(chunkSize, 100); + long available = p.availableMemory(); + assertEquals(chunkSize, available, "pool should have exactly 1 chunk free"); + + // Background thread requests 2 chunks; will block on the 2nd. + AtomicReference err = new AtomicReference<>(); + Thread t = new Thread(() -> { + try { + p.allocateChunks(2 * chunkSize, 5_000); + } catch (Throwable th) { + err.set(th); + } + }, "chunked-waiter"); + t.start(); + // Wait until the thread has joined the waiters queue. + long deadlineMs = System.currentTimeMillis() + 2_000; + while (p.queued() == 0 && System.currentTimeMillis() < deadlineMs) + Thread.sleep(5); + assertEquals(1, p.queued(), "waiter should be parked on the pool's queue"); + + // While the waiter is parked, the pool's available memory must still report the + // 1 chunk's worth — the waiter has NOT consumed any of it (no partial hold). + assertEquals(chunkSize, p.availableMemory(), + "pool memory must reflect no partial holds during the atomic wait"); + + // Free enough chunks for the waiter to complete. + p.deallocate(h1); + p.deallocate(h2); + t.join(5_000); + assertNull(err.get(), "waiter unexpectedly threw: " + err.get()); + p.deallocate(h3); + } + + /** + * FIFO fairness across the K-chunk request. A multi-chunk request that joins the wait queue + * before a single-chunk request must complete first when memory becomes available — the + * K-chunk request occupies a single waiter slot, not K of them. + */ + @Test + public void testFifoFairnessAcrossChunkedAndSingleChunkRequests() throws Exception { + int chunkSize = 64; + long total = 4L * chunkSize; + ChunkedBufferPool p = pool(total, chunkSize); + // Drain the pool entirely so any new request must wait. + ByteBuffer h1 = p.allocate(chunkSize, 100); + ByteBuffer h2 = p.allocate(chunkSize, 100); + ByteBuffer h3 = p.allocate(chunkSize, 100); + ByteBuffer h4 = p.allocate(chunkSize, 100); + assertEquals(0, p.availableMemory()); + + // T_multi enters first, requesting 2 chunks; T_single enters second, requesting 1 chunk. + AtomicReference multiCompletionMs = new AtomicReference<>(); + AtomicReference singleCompletionMs = new AtomicReference<>(); + CountDownLatch multiStarted = new CountDownLatch(1); + Thread tMulti = new Thread(() -> { + try { + multiStarted.countDown(); + List got = p.allocateChunks(2 * chunkSize, 10_000); + multiCompletionMs.set(System.nanoTime()); + for (ByteBuffer b : got) p.deallocate(b); + } catch (Throwable th) { + multiCompletionMs.set(-1L); + } + }, "multi"); + Thread tSingle = new Thread(() -> { + try { + ByteBuffer got = p.allocate(chunkSize, 10_000); + singleCompletionMs.set(System.nanoTime()); + p.deallocate(got); + } catch (Throwable th) { + singleCompletionMs.set(-1L); + } + }, "single"); + + tMulti.start(); + assertTrue(multiStarted.await(2, TimeUnit.SECONDS)); + // Wait for tMulti to be parked before starting tSingle, so the FIFO order is deterministic. + long deadlineMs = System.currentTimeMillis() + 2_000; + while (p.queued() == 0 && System.currentTimeMillis() < deadlineMs) + Thread.sleep(5); + assertEquals(1, p.queued(), "multi-chunk waiter should be the only one queued"); + tSingle.start(); + // Wait for tSingle to also be parked. + deadlineMs = System.currentTimeMillis() + 2_000; + while (p.queued() < 2 && System.currentTimeMillis() < deadlineMs) + Thread.sleep(5); + assertEquals(2, p.queued(), "single-chunk waiter joined after the multi-chunk one"); + + // Now free chunks one at a time, allowing the multi-chunk request to accumulate. + p.deallocate(h1); + Thread.sleep(50); + p.deallocate(h2); + // Both freed — the FIFO leader (multi) should claim both before the single-chunk waiter + // gets any. The multi completes first. + tMulti.join(5_000); + // Now free another chunk for the single-chunk waiter. + p.deallocate(h3); + tSingle.join(5_000); + + assertNotNull(multiCompletionMs.get(), "multi did not complete"); + assertNotNull(singleCompletionMs.get(), "single did not complete"); + assertTrue(multiCompletionMs.get() != -1L && singleCompletionMs.get() != -1L, + "both threads must complete without exception"); + assertTrue(multiCompletionMs.get() <= singleCompletionMs.get(), + "FIFO violated: single (joined later) completed at " + + singleCompletionMs.get() + " before multi at " + multiCompletionMs.get()); + p.deallocate(h4); + } + + /** + * Slow-path rollback must not double-refund pooled chunks. When the waiter polls some + * chunks from the free list, then times out on a subsequent iteration, the inner finally + * refunds {@code accumulated} bytes to {@code nonPooledAvailableMemory} and the outer + * catch re-inserts the polled chunks into {@code free}. If {@code accumulated} includes + * bytes that came from polled chunks, the same memory is refunded twice — the pool + * over-reports {@code availableMemory()} and subsequent allocations can exceed the + * configured {@code buffer.memory} limit. + */ + @Test + public void testSlowPathDoesNotDoubleRefundPolledChunksOnTimeout() throws Exception { + int chunkSize = 64; + long total = 3L * chunkSize; + ChunkedBufferPool p = pool(total, chunkSize); + + // Drain so the next allocateChunks must wait. + ByteBuffer h1 = p.allocate(chunkSize, 100); + ByteBuffer h2 = p.allocate(chunkSize, 100); + ByteBuffer h3 = p.allocate(chunkSize, 100); + assertEquals(0, p.availableMemory()); + + AtomicReference err = new AtomicReference<>(); + Thread t = new Thread(() -> { + try { + // Asks for 3 chunks with a finite deadline. After we deallocate 2 chunks below + // the waiter will poll them, find itself still 1 short, and time out on the + // next await iteration — with 2 chunks held in `pooled` and 2*chunkSize in + // `accumulated`. + p.allocateChunks(3 * chunkSize, 500); + } catch (Throwable th) { + err.set(th); + } + }, "partial-holder"); + t.start(); + + long deadline = System.currentTimeMillis() + 2_000; + while (p.queued() == 0 && System.currentTimeMillis() < deadline) + Thread.sleep(5); + assertEquals(1, p.queued()); + + // Hand the waiter 2 chunks — it polls them into `pooled` then awaits again for the 3rd. + p.deallocate(h1); + p.deallocate(h2); + + // Give the waiter a moment to wake, poll, and re-enter await before the timeout fires. + Thread.sleep(100); + + t.join(5_000); + assertTrue(err.get() instanceof BufferExhaustedException, + "expected BufferExhaustedException, got " + err.get()); + + // After the throw, exactly 2 chunkSizes of physical memory should be back in the pool + // (h1 + h2). h3 is still held outside the pool. The bug double-refunds: it credits + // 2*chunkSize to `nonPooledAvailableMemory` (inner finally) AND re-inserts h1, h2 into + // `free` (outer catch), so availableMemory() reports 4*chunkSize. + assertEquals(2L * chunkSize, p.availableMemory(), + "pool over-reports availableMemory due to double-refund of pooled chunks on rollback"); + + p.deallocate(h3); + } + + /** + * Slow-path success must not corrupt the pool's accounting. When the waiter polls chunks + * from the free list to satisfy its request, the finally block runs with the inner loop's + * "success" reset still in effect — it must not refund or decrement {@code nonPool}. + * Pool chunks transferred to the caller via {@code pooled} are accounted for by removing + * them from {@code free}, not by deducting their bytes from {@code nonPool}. + */ + @Test + public void testSlowPathSuccessDoesNotCorruptAccounting() throws Exception { + int chunkSize = 64; + long total = 3L * chunkSize; + ChunkedBufferPool p = pool(total, chunkSize); + + // Drain so a 2-chunk request must wait. + ByteBuffer h1 = p.allocate(chunkSize, 100); + ByteBuffer h2 = p.allocate(chunkSize, 100); + ByteBuffer h3 = p.allocate(chunkSize, 100); + assertEquals(0, p.availableMemory()); + + AtomicReference err = new AtomicReference<>(); + AtomicReference> got = new AtomicReference<>(); + Thread t = new Thread(() -> { + try { + got.set(p.allocateChunks(2 * chunkSize, 10_000)); + } catch (Throwable th) { + err.set(th); + } + }, "slow-success"); + t.start(); + + long deadline = System.currentTimeMillis() + 2_000; + while (p.queued() == 0 && System.currentTimeMillis() < deadline) + Thread.sleep(5); + assertEquals(1, p.queued()); + + // Deallocate 2 chunks → waiter wakes, polls both, accumulated reaches memoryRequired, + // exits normally. The success path's finally must NOT refund or decrement nonPool. + p.deallocate(h1); + p.deallocate(h2); + t.join(5_000); + assertNull(err.get(), "waiter unexpectedly threw: " + err.get()); + + // After success: waiter owns the 2 chunks (returned via `got`), the test still holds h3. + // The pool has lent out everything it had — availableMemory must be exactly 0. + assertEquals(0, p.availableMemory(), + "slow-path success corrupted accounting (likely a phantom nonPool decrement in the finally)"); + + // Returning all the held buffers must fully restore the pool. + for (ByteBuffer chunk : got.get()) + p.deallocate(chunk); + p.deallocate(h3); + assertEquals(total, p.availableMemory(), + "pool not fully restored after returning all chunks"); + } + + /** + * Closing the pool while a multi-chunk request is parked must surface a + * {@link KafkaException} and refund all reserved bytes. + */ + @Test + public void testCloseDuringAtomicWait() throws Exception { + int chunkSize = 64; + long total = 2L * chunkSize; + ChunkedBufferPool p = pool(total, chunkSize); + // Drain so the next request must wait. + ByteBuffer h1 = p.allocate(chunkSize, 100); + ByteBuffer h2 = p.allocate(chunkSize, 100); + assertEquals(0, p.availableMemory()); + + AtomicReference err = new AtomicReference<>(); + Thread t = new Thread(() -> { + try { + p.allocateChunks(2 * chunkSize, 10_000); + } catch (Throwable th) { + err.set(th); + } + }, "close-during-wait"); + t.start(); + long deadlineMs = System.currentTimeMillis() + 2_000; + while (p.queued() == 0 && System.currentTimeMillis() < deadlineMs) + Thread.sleep(5); + assertEquals(1, p.queued()); + + // Close the pool: signals all waiters; the waiter must throw KafkaException. + p.close(); + t.join(5_000); + assertTrue(err.get() instanceof KafkaException, + "expected KafkaException, got " + err.get()); + p.deallocate(h1); + p.deallocate(h2); + } +} diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java new file mode 100644 index 0000000000000..569c3bbcfcbdc --- /dev/null +++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java @@ -0,0 +1,171 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.kafka.clients.producer.internals; + +import org.apache.kafka.common.metrics.Metrics; +import org.apache.kafka.common.utils.MockTime; + +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; + +import java.nio.ByteBuffer; +import java.util.Collections; +import java.util.List; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; + +public class ChunkedByteBufferOutputStreamTest { + + private final MockTime time = new MockTime(); + private final Metrics metrics = new Metrics(time); + private final String metricGroup = "test"; + + @AfterEach + public void teardown() { + metrics.close(); + } + + private ChunkedBufferPool pool(long total, int chunkSize) { + return new ChunkedBufferPool(total, chunkSize, metrics, time, metricGroup); + } + + private List chunks(ChunkedBufferPool pool, int chunkSize, int count) throws InterruptedException { + return pool.allocateChunks(chunkSize * count, 100); + } + + /** Writes that fit in the initial chunk: stream produces the bytes back via buffer(). */ + @Test + public void testSingleChunkWriteRoundtrip() throws Exception { + int chunkSize = 16; + ChunkedBufferPool p = pool(64, chunkSize); + ChunkedByteBufferOutputStream stream = new ChunkedByteBufferOutputStream(chunks(p, chunkSize, 1), chunkSize, p); + + byte[] payload = new byte[]{1, 2, 3, 4, 5}; + stream.write(payload, 0, payload.length); + + ByteBuffer flat = stream.buffer(); + flat.flip(); + byte[] out = new byte[flat.remaining()]; + flat.get(out); + assertArrayEquals(payload, out); + + stream.deallocate(); + } + + /** Writes that span multiple pre-supplied chunks land in subsequent chunks in order. */ + @Test + public void testWriteSpansChunks() throws Exception { + int chunkSize = 8; + ChunkedBufferPool p = pool(64, chunkSize); + ChunkedByteBufferOutputStream stream = new ChunkedByteBufferOutputStream(chunks(p, chunkSize, 3), chunkSize, p); + + byte[] payload = new byte[20]; + for (int i = 0; i < payload.length; i++) payload[i] = (byte) i; + stream.write(payload, 0, payload.length); + + ByteBuffer flat = stream.buffer(); + flat.flip(); + byte[] out = new byte[flat.remaining()]; + flat.get(out); + assertArrayEquals(payload, out); + + stream.deallocate(); + } + + /** remaining() sums current + queued chunk capacities and shrinks as writes consume bytes. */ + @Test + public void testRemainingSumsChunks() throws Exception { + int chunkSize = 8; + ChunkedBufferPool p = pool(64, chunkSize); + ChunkedByteBufferOutputStream stream = new ChunkedByteBufferOutputStream(chunks(p, chunkSize, 2), chunkSize, p); + + assertEquals(2 * chunkSize, stream.remaining()); + stream.write(new byte[3], 0, 3); + assertEquals(2 * chunkSize - 3, stream.remaining()); + stream.write(new byte[chunkSize], 0, chunkSize); // crosses into chunk 2 + assertEquals(chunkSize - 3, stream.remaining()); + + stream.deallocate(); + } + + /** addBuffers extends capacity so a write that would have overflowed now fits. */ + @Test + public void testAddBuffersExtendsStream() throws Exception { + int chunkSize = 8; + ChunkedBufferPool p = pool(64, chunkSize); + ChunkedByteBufferOutputStream stream = new ChunkedByteBufferOutputStream(chunks(p, chunkSize, 1), chunkSize, p); + + // Fill the initial chunk. + stream.write(new byte[chunkSize], 0, chunkSize); + assertEquals(0, stream.remaining()); + + // Extend and write more — must land in the new chunk. + stream.addBuffers(Collections.singletonList(p.allocate(chunkSize, 100))); + assertEquals(chunkSize, stream.remaining()); + + byte[] more = new byte[]{9, 9, 9}; + stream.write(more, 0, more.length); + assertEquals(chunkSize - 3, stream.remaining()); + + stream.deallocate(); + } + + /** Writing past all pre-supplied (and added) chunks throws — uncompressed-only contract. */ + @Test + public void testOverflowThrows() throws Exception { + int chunkSize = 4; + ChunkedBufferPool p = pool(64, chunkSize); + ChunkedByteBufferOutputStream stream = new ChunkedByteBufferOutputStream(chunks(p, chunkSize, 1), chunkSize, p); + + stream.write(new byte[chunkSize], 0, chunkSize); + assertThrows(IllegalStateException.class, () -> stream.write((byte) 1)); + + stream.deallocate(); + } + + /** position(int) at construction walks across pre-supplied chunks. */ + @Test + public void testPositionWalksAcrossChunks() throws Exception { + int chunkSize = 4; + ChunkedBufferPool p = pool(32, chunkSize); + ChunkedByteBufferOutputStream stream = new ChunkedByteBufferOutputStream(chunks(p, chunkSize, 3), chunkSize, p); + + stream.position(6); // straddles chunk 0 (4 bytes) and chunk 1 (2 bytes) + assertEquals(6, stream.position()); + assertEquals(6, stream.remaining()); + + stream.deallocate(); + } + + /** deallocate returns all chunks to the pool. */ + @Test + public void testDeallocateReturnsAllChunks() throws Exception { + int chunkSize = 8; + long total = 64; + ChunkedBufferPool p = pool(total, chunkSize); + List initial = chunks(p, chunkSize, 3); + ChunkedByteBufferOutputStream stream = new ChunkedByteBufferOutputStream(initial, chunkSize, p); + // Also attach one extra chunk so deallocate must handle the added-buffer case too. + stream.addBuffers(Collections.singletonList(p.allocate(chunkSize, 100))); + assertEquals(total - 4L * chunkSize, p.availableMemory()); + + stream.deallocate(); + assertEquals(total, p.availableMemory()); + } +} diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java new file mode 100644 index 0000000000000..b191aadb4645b --- /dev/null +++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java @@ -0,0 +1,350 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.kafka.clients.producer.internals; + +import org.apache.kafka.clients.MetadataSnapshot; +import org.apache.kafka.common.Cluster; +import org.apache.kafka.common.Node; +import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.compress.Compression; +import org.apache.kafka.common.metrics.Metrics; +import org.apache.kafka.common.protocol.Errors; +import org.apache.kafka.common.record.internal.MemoryRecords; +import org.apache.kafka.common.record.internal.Record; +import org.apache.kafka.common.record.internal.RecordBatch; +import org.apache.kafka.common.requests.MetadataResponse.PartitionMetadata; +import org.apache.kafka.common.utils.MockTime; +import org.apache.kafka.common.utils.internals.LogContext; + +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; + +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.Deque; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; +import java.util.stream.Collectors; +import java.util.stream.Stream; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +public class ChunkedRecordAccumulatorTest { + + private final String topic = "test"; + private final int partition1 = 0; + private final Node node1 = new Node(0, "localhost", 1111); + private final TopicPartition tp1 = new TopicPartition(topic, partition1); + + private final PartitionMetadata partMetadata1 = + new PartitionMetadata(Errors.NONE, tp1, Optional.of(node1.id()), Optional.empty(), null, null, null); + private final List partMetadatas = new ArrayList<>(Arrays.asList(partMetadata1)); + private final Map nodes = + Stream.of(node1).collect(Collectors.toMap(Node::id, java.util.function.Function.identity())); + private final MetadataSnapshot metadataCache = new MetadataSnapshot(null, nodes, partMetadatas, + Collections.emptySet(), Collections.emptySet(), Collections.emptySet(), null, Collections.emptyMap()); + private final Cluster cluster = metadataCache.cluster(); + + private final MockTime time = new MockTime(); + private final Metrics metrics = new Metrics(time); + private final LogContext logContext = new LogContext(); + private final long maxBlockTimeMs = 1000; + private final byte[] key = "k".getBytes(); + + @AfterEach + public void teardown() { + metrics.close(); + } + + private ChunkedRecordAccumulator newAccumulator(int batchSize, int chunkSize, long totalMemory, Compression compression) { + ChunkedBufferPool pool = new ChunkedBufferPool(totalMemory, chunkSize, metrics, time, "producer-metrics"); + return new ChunkedRecordAccumulator(logContext, batchSize, compression, + /* lingerMs */ 0, /* retryBackoffMs */ 0L, /* retryBackoffMaxMs */ 0L, + /* deliveryTimeoutMs */ 3200, metrics, "producer-metrics", time, + /* transactionManager */ null, pool); + } + + /** First record creates a batch backed by chunked allocation. */ + @Test + public void testFirstRecordCreatesChunkedBatch() throws Exception { + int chunkSize = 256; + ChunkedRecordAccumulator accum = newAccumulator(1024, chunkSize, 16L * chunkSize, Compression.NONE); + + byte[] value = new byte[64]; + accum.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null, + maxBlockTimeMs, time.milliseconds(), cluster); + + Deque dq = batchesFor(accum, tp1); + assertEquals(1, dq.size()); + ProducerBatch batch = dq.peekFirst(); + assertNotNull(batch); + assertEquals(1, batch.recordCount); + accum.close(); + } + + /** + * A small follow-up record fits in the existing batch's pre-allocated chunks — no extension + * needed. + */ + @Test + public void testFollowupRecordFitsWithoutExtension() throws Exception { + int chunkSize = 256; + ChunkedRecordAccumulator accum = newAccumulator(1024, chunkSize, 16L * chunkSize, Compression.NONE); + + accum.append(topic, partition1, 0L, key, new byte[16], Record.EMPTY_HEADERS, null, + maxBlockTimeMs, time.milliseconds(), cluster); + accum.append(topic, partition1, 0L, key, new byte[16], Record.EMPTY_HEADERS, null, + maxBlockTimeMs, time.milliseconds(), cluster); + + Deque dq = batchesFor(accum, tp1); + assertEquals(1, dq.size(), "Both records should land in the same batch"); + assertEquals(2, dq.peekFirst().recordCount); + accum.close(); + } + + /** + * A follow-up record that would overflow the existing chunks triggers mid-batch extension: + * additional chunks are attached and the record lands in the same batch. + */ + @Test + public void testMidBatchExtensionGrowsExistingBatch() throws Exception { + int chunkSize = 128; + ChunkedRecordAccumulator accum = newAccumulator(8192, chunkSize, 16L * chunkSize, Compression.NONE); + + // First record fits in 1 chunk (~64+overhead). + accum.append(topic, partition1, 0L, key, new byte[32], Record.EMPTY_HEADERS, null, + maxBlockTimeMs, time.milliseconds(), cluster); + + Deque dq = batchesFor(accum, tp1); + assertEquals(1, dq.size()); + int initialRecordCount = dq.peekFirst().recordCount; + assertEquals(1, initialRecordCount); + + // Second record big enough to require an extra chunk. + accum.append(topic, partition1, 0L, key, new byte[200], Record.EMPTY_HEADERS, null, + maxBlockTimeMs, time.milliseconds(), cluster); + + // Same batch, now with the second record. + assertEquals(1, dq.size(), "Should still be the same batch — extension should not roll"); + assertEquals(2, dq.peekFirst().recordCount, + "Second record should land in the extended batch"); + accum.close(); + } + + /** Compression is rejected at batch creation in PR1 with a clear message. */ + @Test + public void testCompressionThrowsAtBatchCreation() { + int chunkSize = 256; + ChunkedRecordAccumulator accum = newAccumulator(1024, chunkSize, 16L * chunkSize, Compression.gzip().build()); + + UnsupportedOperationException ex = assertThrows(UnsupportedOperationException.class, () -> + accum.append(topic, partition1, 0L, key, new byte[32], Record.EMPTY_HEADERS, null, + maxBlockTimeMs, time.milliseconds(), cluster)); + assertTrue(ex.getMessage().contains("not implemented yet"), + "Expected 'not implemented yet' in message, got: " + ex.getMessage()); + accum.close(); + } + + /** + * Race between two concurrent appenders on the same chunked tail. Both threads probe + * {@code extensionBytesNeeded} under the deque lock against the same physical remaining + * capacity, drop the lock, and allocate gap-sized chunks off-lock. When the first appender + * re-takes the lock and consumes most of the original remaining, the second appender's + * allocation — sized against the pre-race remaining — is short by the bytes the first + * appender consumed. Without a re-probe in the second-lock block, the second appender + * proceeds to write past the stream's physical capacity and {@code advanceToNextChunk} + * throws {@link IllegalStateException}. + *

+ * Coordination: a {@link ChunkedBufferPool} subclass blocks the FIRST {@code allocateChunks} + * call after the test arms a latch (i.e. thread A's). Thread B's allocate is the second call + * and proceeds normally — B re-takes the lock, attaches its chunks, appends its record, and + * returns. Only then is thread A released, so A attaches its chunks and tries to append + * against a tail whose remaining has already shrunk. + */ + @Test + public void testConcurrentExtensionRaceDoesNotOverflowChunkedStream() throws Exception { + final int chunkSize = 256; + final int recordValueSize = 350; // sized so gap ≈ chunkSize → minimal rounding cushion + + final AtomicBoolean armed = new AtomicBoolean(false); + final CountDownLatch aHoldsChunks = new CountDownLatch(1); + final CountDownLatch bDoneAttaching = new CountDownLatch(1); + + ChunkedBufferPool pool = new ChunkedBufferPool(64L * chunkSize, chunkSize, metrics, time, "producer-metrics") { + @Override + public List allocateChunks(int totalSize, long maxTimeToBlockMs) throws InterruptedException { + boolean blockThisCall = armed.compareAndSet(true, false); + List chunks = super.allocateChunks(totalSize, maxTimeToBlockMs); + if (blockThisCall) { + aHoldsChunks.countDown(); + if (!bDoneAttaching.await(10, TimeUnit.SECONDS)) { + throw new InterruptedException("test timeout waiting for B"); + } + } + return chunks; + } + }; + ChunkedRecordAccumulator accum = new ChunkedRecordAccumulator(logContext, 8192, Compression.NONE, + /* lingerMs */ 0, /* retryBackoffMs */ 0L, /* retryBackoffMaxMs */ 0L, + /* deliveryTimeoutMs */ 3200, metrics, "producer-metrics", time, + /* transactionManager */ null, pool); + try { + // Warmup: tiny first record establishes a chunked tail with a known remainingInStream. + accum.append(topic, partition1, 0L, key, new byte[1], Record.EMPTY_HEADERS, null, + maxBlockTimeMs, time.milliseconds(), cluster); + + // Arm the race: the next allocateChunks call (A's) will block until B has appended. + armed.set(true); + + AtomicReference aError = new AtomicReference<>(); + Thread tA = new Thread(() -> { + try { + accum.append(topic, partition1, 0L, key, new byte[recordValueSize], + Record.EMPTY_HEADERS, null, maxBlockTimeMs, time.milliseconds(), cluster); + } catch (Throwable t) { + aError.set(t); + } + }, "test-appender-A"); + tA.start(); + + // Wait for A to be parked inside allocateChunks holding its (pre-attach) chunks. + assertTrue(aHoldsChunks.await(10, TimeUnit.SECONDS), "A did not reach allocateChunks"); + + // B runs on this thread. Its allocateChunks is the second call — armed=false now, + // so it proceeds without blocking. B probes the same R as A (A hasn't attached yet), + // attaches its own chunks, appends its record. Now the tail's remaining has shrunk + // by B's actual record size, but A's still-off-lock allocation didn't know that. + accum.append(topic, partition1, 0L, key, new byte[recordValueSize], Record.EMPTY_HEADERS, null, + maxBlockTimeMs, time.milliseconds(), cluster); + + // Release A. Without the fix in the second-lock block, A's tryAppendToExisting + // will overflow the chunked stream and throw IllegalStateException. + bDoneAttaching.countDown(); + tA.join(10_000); + assertFalse(tA.isAlive(), "Thread A did not complete in time"); + + assertNull(aError.get(), + "Thread A failed: " + (aError.get() == null ? "" : aError.get().toString())); + } finally { + // Drain any state regardless of outcome. + accum.close(); + } + } + + /** + * Regression guard: an inflight chunked batch returns all K chunks to the pool on the + * expiration / broker-disconnect path, not just one {@code initialCapacity} chunkSize. + * The parent {@code RecordAccumulator.deallocate} for inflight batches credits only + * {@code initialCapacity} back to the pool and throws; {@link ChunkedRecordAccumulator} + * overrides {@code deallocate} to return all chunks before the throw. Removing that + * override would re-introduce a (K−1)×chunkSize leak and fail this test. + */ + @Test + public void testInflightExpirationReturnsAllChunksToPool() throws Exception { + int chunkSize = 128; + long totalMemory = 32L * chunkSize; + ChunkedBufferPool pool = new ChunkedBufferPool(totalMemory, chunkSize, metrics, time, "producer-metrics"); + ChunkedRecordAccumulator accum = new ChunkedRecordAccumulator(logContext, 8192, Compression.NONE, + /* lingerMs */ 0, /* retryBackoffMs */ 0L, /* retryBackoffMaxMs */ 0L, + /* deliveryTimeoutMs */ 3200, metrics, "producer-metrics", time, + /* transactionManager */ null, pool); + + // A record large enough that allocateChunks reserves multiple chunks (K > 1). + byte[] value = new byte[400]; + accum.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null, + maxBlockTimeMs, time.milliseconds(), cluster); + + Deque dq = batchesFor(accum, tp1); + assertEquals(1, dq.size()); + ProducerBatch batch = dq.peekFirst(); + + // Confirm the batch actually consumed multiple chunks. estimateSizeInBytesUpperBound for + // a 400-byte value with v2 framing is well over chunkSize, so K should be >= 2. + long heldBeforeDeallocate = totalMemory - pool.availableMemory(); + assertTrue(heldBeforeDeallocate >= 2L * chunkSize, + "test setup expects K >= 2; pool held " + heldBeforeDeallocate + " bytes"); + + // Mark the batch as if the Sender had drained it. + batch.setInflight(true); + + // The inflight branch throws after deallocating. The chunked override must return + // all K chunks (not just initialCapacity) before propagating the exception. + assertThrows(IllegalStateException.class, () -> accum.deallocate(batch)); + + assertEquals(totalMemory, pool.availableMemory(), + "pool should be fully restored after inflight-expiration deallocate; " + + "any K-1 unsurrendered chunks indicate the chunked-leak regression"); + + accum.close(); + } + + private Deque batchesFor(RecordAccumulator accum, TopicPartition tp) { + return accum.getDeque(tp); + } + + /** + * Sender's drain (and any other caller of {@code batch.close()}) cascades through + * {@code MemoryRecordsBuilder.closeForRecordAppends()} → + * {@code appendStream.close()} → {@code bufferStream.close()}. The chunked stream's + * buffer must remain readable after that cascade so the builder can flatten chunks into + * the heap buffer it sends over the network. If the chunked stream's {@code close()} + * prematurely returns chunks to the pool, the flattened buffer is empty and the header + * write fails (or produces an empty record set). + */ + @Test + public void testBatchCloseDoesNotDeallocateChunksPrematurely() throws Exception { + int chunkSize = 256; + ChunkedRecordAccumulator accum = newAccumulator(8192, chunkSize, 32L * chunkSize, Compression.NONE); + + byte[] value = new byte[64]; + accum.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null, + maxBlockTimeMs, time.milliseconds(), cluster); + + Deque dq = batchesFor(accum, tp1); + ProducerBatch batch = dq.peekFirst(); + assertNotNull(batch); + + // Matches the call sites in RecordAccumulator.drain (line 941) and Sender. + batch.close(); + MemoryRecords records = batch.records(); + assertTrue(records.sizeInBytes() > 0, + "batch must produce a non-empty record set after close; chunks were deallocated prematurely"); + + int count = 0; + for (RecordBatch rb : records.batches()) { + for (Record r : rb) { + count++; + assertNotNull(r.value()); + } + } + assertEquals(1, count, "expected exactly 1 record after close"); + + accum.close(); + } +} From 64147f6b492933e343d983f3a168edcdc427976b Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Thu, 18 Jun 2026 15:27:38 -0400 Subject: [PATCH 03/61] config & ratio-aware chunk sizing --- .../kafka/clients/producer/KafkaProducer.java | 58 ++++- .../clients/producer/ProducerConfig.java | 21 ++ .../producer/internals/ChunkedBufferPool.java | 47 ++-- .../ChunkedByteBufferOutputStream.java | 36 +-- .../internals/ChunkedProducerBatch.java | 94 +++++++ .../internals/ChunkedRecordAccumulator.java | 236 ++++++++---------- .../producer/internals/ProducerBatch.java | 38 +-- .../producer/internals/RecordAccumulator.java | 54 +++- .../record/internal/AbstractRecords.java | 25 ++ .../record/internal/MemoryRecordsBuilder.java | 48 +++- .../clients/producer/ProducerConfigTest.java | 35 +++ .../ChunkedRecordAccumulatorTest.java | 91 +++++-- 12 files changed, 544 insertions(+), 239 deletions(-) create mode 100644 clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedProducerBatch.java diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/KafkaProducer.java b/clients/src/main/java/org/apache/kafka/clients/producer/KafkaProducer.java index dfa72f2adad71..0990a68439171 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/KafkaProducer.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/KafkaProducer.java @@ -28,6 +28,8 @@ import org.apache.kafka.clients.consumer.OffsetCommitCallback; import org.apache.kafka.clients.producer.internals.BufferPool; import org.apache.kafka.clients.producer.internals.BuiltInPartitioner; +import org.apache.kafka.clients.producer.internals.ChunkedBufferPool; +import org.apache.kafka.clients.producer.internals.ChunkedRecordAccumulator; import org.apache.kafka.clients.producer.internals.KafkaProducerMetrics; import org.apache.kafka.clients.producer.internals.ProducerInterceptors; import org.apache.kafka.clients.producer.internals.ProducerMetadata; @@ -90,6 +92,7 @@ import java.util.Arrays; import java.util.Collections; import java.util.List; +import java.util.Locale; import java.util.Map; import java.util.Objects; import java.util.Optional; @@ -456,19 +459,48 @@ public KafkaProducer(Properties properties, Serializer keySerializer, Seriali // As per Kafka producer configuration documentation batch.size may be set to 0 to explicitly disable // batching which in practice actually means using a batch size of 1. int batchSize = Math.max(1, config.getInt(ProducerConfig.BATCH_SIZE_CONFIG)); - this.accumulator = new RecordAccumulator(logContext, - batchSize, - compression, - lingerMs(config), - retryBackoffMs, - retryBackoffMaxMs, - deliveryTimeoutMs, - partitionerConfig, - metrics, - PRODUCER_METRIC_GROUP_NAME, - time, - transactionManager, - new BufferPool(this.totalMemorySize, batchSize, metrics, time, PRODUCER_METRIC_GROUP_NAME)); + String allocationStrategy = config.getString(ProducerConfig.BUFFER_MEMORY_ALLOCATION_STRATEGY_CONFIG) + .toLowerCase(Locale.ROOT); + boolean incremental = ProducerConfig.BUFFER_MEMORY_ALLOCATION_STRATEGY_INCREMENTAL.equals(allocationStrategy); + if (incremental && compression.type() != CompressionType.NONE) { + throw new ConfigException("The " + ProducerConfig.BUFFER_MEMORY_ALLOCATION_STRATEGY_INCREMENTAL + + " " + ProducerConfig.BUFFER_MEMORY_ALLOCATION_STRATEGY_CONFIG + + " does not support compression yet. " + ProducerConfig.COMPRESSION_TYPE_CONFIG + + " must be set to none."); + } + // Use the chunked path only when a batch is at least one full chunk + // (batch.size >= CHUNK_SIZE). Below that, a batch can't fill even one chunk, so chunking + // would reserve a whole chunk for a batch capped under it (over-reservation); the full + // strategy (reserving up to batch.size per batch) is both leaner and identical to trunk. + if (incremental && batchSize >= ChunkedRecordAccumulator.CHUNK_SIZE) { + this.accumulator = new ChunkedRecordAccumulator(logContext, + batchSize, + compression, + lingerMs(config), + retryBackoffMs, + retryBackoffMaxMs, + deliveryTimeoutMs, + partitionerConfig, + metrics, + PRODUCER_METRIC_GROUP_NAME, + time, + transactionManager, + new ChunkedBufferPool(this.totalMemorySize, ChunkedRecordAccumulator.CHUNK_SIZE, metrics, time, PRODUCER_METRIC_GROUP_NAME)); + } else { + this.accumulator = new RecordAccumulator(logContext, + batchSize, + compression, + lingerMs(config), + retryBackoffMs, + retryBackoffMaxMs, + deliveryTimeoutMs, + partitionerConfig, + metrics, + PRODUCER_METRIC_GROUP_NAME, + time, + transactionManager, + new BufferPool(this.totalMemorySize, batchSize, metrics, time, PRODUCER_METRIC_GROUP_NAME)); + } this.errors = this.metrics.sensor("errors"); this.sender = newSender(logContext, kafkaClient, this.metadata); diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/ProducerConfig.java b/clients/src/main/java/org/apache/kafka/clients/producer/ProducerConfig.java index 38a007b6d212b..286f3e37dc823 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/ProducerConfig.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/ProducerConfig.java @@ -228,6 +228,20 @@ public class ProducerConfig extends AbstractConfig { + "not all memory the producer uses is used for buffering. Some additional memory will be used for compression (if " + "compression is enabled) as well as for maintaining in-flight requests."; + /** buffer.memory.allocation.strategy */ + public static final String BUFFER_MEMORY_ALLOCATION_STRATEGY_CONFIG = "buffer.memory.allocation.strategy"; + public static final String BUFFER_MEMORY_ALLOCATION_STRATEGY_FULL = "full"; + public static final String BUFFER_MEMORY_ALLOCATION_STRATEGY_INCREMENTAL = "incremental"; + private static final String BUFFER_MEMORY_ALLOCATION_STRATEGY_DOC = "Controls how the producer allocates memory from " + BUFFER_MEMORY_CONFIG + " for record batches. The following values are supported: " + + "

    " + + "
  • " + BUFFER_MEMORY_ALLOCATION_STRATEGY_FULL + ": reserve " + BATCH_SIZE_CONFIG + " bytes up front for each in-progress batch, " + + "regardless of the actual size of the records appended to it.
    • " + + "
  • " + BUFFER_MEMORY_ALLOCATION_STRATEGY_INCREMENTAL + ": allocate memory on demand as records are appended, growing each batch " + + "up to " + BATCH_SIZE_CONFIG + " bytes. This keeps memory usage proportional to the data actually buffered rather than to the number of " + + "active partitions, which allows larger " + BATCH_SIZE_CONFIG + " values (e.g. for high-latency clusters) without reserving " + + "" + BATCH_SIZE_CONFIG + " bytes for every active partition.
    • " + + "
"; + /** retry.backoff.ms */ public static final String RETRY_BACKOFF_MS_CONFIG = CommonClientConfigs.RETRY_BACKOFF_MS_CONFIG; @@ -399,6 +413,13 @@ public class ProducerConfig extends AbstractConfig { Importance.MEDIUM, CommonClientConfigs.CLIENT_DNS_LOOKUP_DOC) .define(BUFFER_MEMORY_CONFIG, Type.LONG, 32 * 1024 * 1024L, atLeast(0L), Importance.HIGH, BUFFER_MEMORY_DOC) + .define(BUFFER_MEMORY_ALLOCATION_STRATEGY_CONFIG, + Type.STRING, + BUFFER_MEMORY_ALLOCATION_STRATEGY_FULL, + ConfigDef.CaseInsensitiveValidString + .in(BUFFER_MEMORY_ALLOCATION_STRATEGY_FULL, BUFFER_MEMORY_ALLOCATION_STRATEGY_INCREMENTAL), + Importance.MEDIUM, + BUFFER_MEMORY_ALLOCATION_STRATEGY_DOC) .define(RETRIES_CONFIG, Type.INT, Integer.MAX_VALUE, between(0, Integer.MAX_VALUE), Importance.HIGH, RETRIES_DOC) .define(ACKS_CONFIG, Type.STRING, diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java index 96a4e6e0b36e4..9fa367d80e889 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java @@ -28,15 +28,12 @@ import java.util.concurrent.locks.Condition; /** - * A {@link BufferPool} dedicated to chunk-sized buffer reuse. The chunk size is passed to the - * parent as its {@link #poolableSize()}. + * A {@link BufferPool} dedicated to chunk-sized buffer reuse (chunk size = {@link #poolableSize()}). *

- * Adds {@link #allocateChunks(int, long)} for callers that need to acquire multiple - * chunks at once. The implementation is atomic across the K-chunk request: - * a single lock acquisition for the whole reservation, a single {@link Condition} - * added to {@link #waiters} so the request is FIFO-fair against single-chunk acquirers, and - * any failure (timeout, close, OOM) refunds the entire reservation atomically before the - * exception propagates. No partial chunk holds are visible outside the lock during the wait. + * Adds {@link #allocateChunks(int, long)} to acquire multiple chunks atomically: one lock + * acquisition and a single {@link Condition} on {@link #waiters} (FIFO-fair against single-chunk + * acquirers), with any failure (timeout, close, OOM) refunding the whole reservation before the + * exception propagates. No partial holds are visible during the wait. */ public class ChunkedBufferPool extends BufferPool { @@ -45,22 +42,19 @@ public ChunkedBufferPool(long memory, int chunkSize, Metrics metrics, Time time, } /** - * Allocate {@code ceil(totalSize / chunkSize)} chunk-sized buffers atomically. - * This follows the same principle as the allocate in the parent BufferPool class: if there is - * memory available, allocate the requested number of chunks. If not, wait for memory. - * If it needs to wait, it blocks up to {@code maxTimeToBlockMs} for the whole request, - * queued on {@link #waiters} alongside single-chunk acquirers (FIFO). - *

- * On any failure path every byte reserved during the call is returned to the pool - * and the next waiter is signaled. No partial chunk holds are visible to other threads during the wait. - * The reservation is tracked as bytes against {@link #nonPooledAvailableMemory} plus chunks polled from {@link #free}. + * Allocate {@code ceil(totalSize / chunkSize)} chunk-sized buffers atomically, mirroring + * {@link BufferPool#allocate}: satisfied immediately if memory is available, else blocks up to + * {@code maxTimeToBlockMs} for the whole request (FIFO on {@link #waiters}). The reservation is + * tracked as bytes against {@link #nonPooledAvailableMemory} plus chunks polled from + * {@link #free}; on any failure path every reserved byte is returned and the next waiter signaled. * * @param totalSize minimum total bytes of capacity required across the returned chunks * @param maxTimeToBlockMs maximum time in milliseconds to block waiting for memory * @return list of {@code ceil(totalSize / chunkSize)} {@code ByteBuffer}s, each of capacity * {@code chunkSize} * @throws InterruptedException if interrupted while waiting - * @throws IllegalArgumentException if {@code totalSize <= 0} or {@code totalSize > totalMemory()} + * @throws IllegalArgumentException if {@code totalSize <= 0}, or if the request rounded up to + * whole chunks exceeds {@code totalMemory()} * @throws BufferExhaustedException if the request can't be satisfied within {@code maxTimeToBlockMs} * @throws KafkaException if the pool is closed during the wait */ @@ -75,11 +69,17 @@ public List allocateChunks(int totalSize, long maxTimeToBlockMs) thr int chunkSize = poolableSize(); int numChunks = (int) (((long) totalSize + chunkSize - 1L) / chunkSize); long memoryRequired = (long) numChunks * chunkSize; + // Rounding totalSize up to whole chunks can require more pool memory than buffer.memory + // holds even when totalSize alone fits (the single-buffer "full" path reserves exactly the + // record size, so it never hits this). Reject it here to fail fast, instead of blocking for + // max.block.ms and then throwing BufferExhausted for a request that can never be satisfied. + if (memoryRequired > totalMemory()) + throw new IllegalArgumentException("Attempt to allocate " + numChunks + " chunks of " + + chunkSize + " bytes (" + memoryRequired + " bytes total), but there is a hard limit of " + + totalMemory() + " on memory allocations."); - // Chunks pulled from the free list during the reservation. Remaining bytes - // (memoryRequired - pooled.size() * chunkSize) are reserved against - // nonPooledAvailableMemory and materialized as raw allocations after the lock is - // released. + // Chunks pulled from the free list; the remaining bytes are reserved against + // nonPooledAvailableMemory and materialized as raw allocations after the lock is released. List pooled = new ArrayList<>(numChunks); lock.lock(); @@ -95,8 +95,7 @@ public List allocateChunks(int totalSize, long maxTimeToBlockMs) thr pooled.add(free.pollFirst()); long remainingBytes = memoryRequired - (long) pooled.size() * chunkSize; if (remainingBytes > 0) { - // remainingBytes is bounded by memoryRequired ≤ totalMemory; the entry check - // rejects requests larger than that, so the int cast is safe in practice. + // remainingBytes <= memoryRequired <= totalMemory (validated above), so the int cast is safe. freeUp((int) remainingBytes); this.nonPooledAvailableMemory -= remainingBytes; } diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java index ccf1d3d7374c9..3b9a9ed42ff5e 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java @@ -32,10 +32,9 @@ * chunks), {@link IllegalStateException} is thrown. The caller is responsible for attaching * additional chunks before any such write. *

- * Automatic mid-write growth (allocating a chunk from inside {@code write()}) will be added - * together with compression support in a follow-up: compressor buffering can cause a record's - * write to exceed its reservation, which the caller can't predict in advance. Until then, the - * write path stays allocation-free and deterministic. + * Automatic mid-write growth (allocating from inside {@code write()}) is a follow-up tied to + * compression support, where compressor buffering can make a write exceed its reservation; until + * then the write path stays allocation-free and deterministic. *

* When {@link #buffer()} is called (typically at batch close), all chunks are flattened into a * single contiguous ByteBuffer — exactly one copy operation. @@ -110,11 +109,10 @@ private void ensureChunkCapacity(int needed) { } /** - * Advances {@code currentChunk} to the next pre-supplied chunk. The caller is responsible - * for obtaining additional chunks (e.g. {@code ChunkedBufferPool.allocateChunks}) and - * attaching them via {@link #addBuffers(List)} before any write that would exceed - * {@link #remaining()}. The KIP's mid-record growth path (non-blocking pool then direct - * heap) lands together with compression support. + * Advances {@code currentChunk} to the next pre-supplied chunk; throws if none is left. The + * caller must attach additional chunks via {@link #addBuffers(List)} before any write that + * exceeds {@link #remaining()}. (Mid-record growth — non-blocking pool then heap — is a + * follow-up that lands with compression support.) */ private void advanceToNextChunk() { if (currentChunkIndex + 1 >= chunks.size()) { @@ -140,6 +138,11 @@ public ByteBuffer buffer() { if (flattenedBuffer != null && !dirty) { return flattenedBuffer; } + // TODO: KAFKA-20687. This flatten runs at batch close, when the chunk set is final. + // Today all chunks (used and unused) are returned to the pool only when the batch + // completes (via deallocate(pool)). Consider releasing the fully-unused chunks + // early, at close, rather than holding them until completion — removing them from + // `chunks` here so the completion-time deallocate(pool) does not double-return them. int totalSize = 0; for (ByteBuffer chunk : chunks) { totalSize += chunk.position(); @@ -153,6 +156,12 @@ public ByteBuffer buffer() { chunk.position(chunkPos); } dirty = false; + // The bytes are now copied into flattenedBuffer, but we intentionally do not release the + // chunks until batch completion, so the in-flight data stays reserved against the + // buffer.memory budget (the pool's available memory reflects it), consistent with the + // "full" strategy. Releasing here would return that memory to the pool while the bytes are + // still in flight in the heap copy, letting the pool admit more than buffer.memory intends. + // This flattening is an initial approach and will be removed with KAFKA-20580. return flattenedBuffer; } @@ -219,15 +228,14 @@ public int initialCapacity() { @Override public void ensureRemaining(int remainingBytesRequired) { - // The stream advances one chunk at a time — the most any single ensureChunkCapacity - // call can guarantee is `chunkSize` of contiguous-by-chunk space. Callers needing more - // than that are expected to attach chunks via addBuffers before writing; write(byte[]) - // itself loops across chunks so contiguous capacity isn't required. + // A single call can guarantee at most `chunkSize` of space (the stream advances one chunk + // at a time); callers needing more attach chunks via addBuffers first. write(byte[]) loops + // across chunks, so contiguous capacity isn't required. ensureChunkCapacity(Math.min(remainingBytesRequired, chunkSize)); } /** - * Returns all pool-allocated chunks to the buffer pool. + * Returns all pool-allocated chunks to the buffer pool. Called at batch completion. */ public void deallocate(BufferPool pool) { if (pool != null) { diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedProducerBatch.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedProducerBatch.java new file mode 100644 index 0000000000000..e80ed7861826d --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedProducerBatch.java @@ -0,0 +1,94 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.kafka.clients.producer.internals; + +import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.header.Header; +import org.apache.kafka.common.record.internal.MemoryRecordsBuilder; +import org.apache.kafka.common.utils.ByteBufferOutputStream; + +import java.nio.ByteBuffer; +import java.util.List; + +/** + * A {@link ProducerBatch} for the incremental (chunked) buffer.memory allocation strategy, backed + * by a {@link MemoryRecordsBuilder} whose stream is a {@link ChunkedByteBufferOutputStream}. + * It adds mid-batch chunk extension support ({@link #extensionBytesNeeded} / + * {@link #addBuffers}) and overrides the pool deallocation hooks so all chunks are returned to + * the pool rather than a single buffer. + *

+ * This class is not thread safe and external synchronization must be used when modifying it. + */ +public class ChunkedProducerBatch extends ProducerBatch { + + // The parent's builder reference is private; keep our own to access stream capacity state. + private final MemoryRecordsBuilder recordsBuilder; + + public ChunkedProducerBatch(TopicPartition tp, MemoryRecordsBuilder recordsBuilder, long createdMs) { + super(tp, recordsBuilder, createdMs); + this.recordsBuilder = recordsBuilder; + } + + /** + * Bytes of physical buffer this batch needs before {@code tryAppend} could accept the given + * record. Returns 0 when the record fits (either logically full, or the stream's combined + * chunk capacity has room). Positive when {@code hasRoomFor} allows the record but the + * chunks lack physical capacity — the accumulator allocates exactly the missing bytes + * (rounded up to whole chunks) and attaches them via {@link #addBuffers} before retrying. + */ + int extensionBytesNeeded(long timestamp, byte[] key, byte[] value, Header[] headers) { + if (recordCount == 0) + return 0; + if (!recordsBuilder.hasRoomFor(timestamp, key, value, headers)) + return 0; + // Size against the batch's projected total output after this record (header counted once, + // ratio-adjusted when compressed), not per-record. Per-record sizing would over-count the + // header and miss the compressor's flush-accumulation behavior. + int target = recordsBuilder.estimatedBytesWrittenAfter(key, value, headers); + ByteBufferOutputStream stream = recordsBuilder.bufferStream(); + int totalAttachedCapacity = stream.position() + stream.remaining(); + return Math.max(0, target - totalAttachedCapacity); + } + + /** + * Attach pre-allocated chunks to this batch's stream so the next {@code tryAppend} can + * spill into them. Ownership of the chunks transfers to the stream. + */ + void addBuffers(List chunks) { + stream().addBuffers(chunks); + } + + @Override + protected void deallocateBuffer(BufferPool pool) { + stream().deallocate(pool); + } + + /** + * Unlike the single-buffer batch — which must donate a fresh buffer because the network + * layer may still be reading the pooled one — a chunked batch's inflight bytes live in the + * separate flattened buffer (see {@link ChunkedByteBufferOutputStream#buffer()}), so it is + * safe to return the actual chunks to the pool here. + */ + @Override + protected void deallocateInflightBuffer(BufferPool pool) { + stream().deallocate(pool); + } + + private ChunkedByteBufferOutputStream stream() { + return (ChunkedByteBufferOutputStream) recordsBuilder.bufferStream(); + } +} diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java index cac18c608d195..a587754d80c44 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java @@ -17,13 +17,17 @@ package org.apache.kafka.clients.producer.internals; import org.apache.kafka.clients.producer.BufferExhaustedException; +import org.apache.kafka.clients.producer.Callback; import org.apache.kafka.clients.producer.RecordMetadata; import org.apache.kafka.common.Cluster; +import org.apache.kafka.common.KafkaException; +import org.apache.kafka.common.TopicPartition; import org.apache.kafka.common.compress.Compression; import org.apache.kafka.common.header.Header; import org.apache.kafka.common.metrics.Metrics; import org.apache.kafka.common.record.TimestampType; import org.apache.kafka.common.record.internal.AbstractRecords; +import org.apache.kafka.common.record.internal.CompressionRatioEstimator; import org.apache.kafka.common.record.internal.CompressionType; import org.apache.kafka.common.record.internal.MemoryRecordsBuilder; import org.apache.kafka.common.record.internal.Record; @@ -40,24 +44,31 @@ * A {@link RecordAccumulator} variant that uses chunked buffer allocation: each new batch * pre-allocates {@code ceil(recordSize / chunkSize)} chunks from a {@link ChunkedBufferPool}, * and grows mid-batch by attaching additional chunks via - * {@link ProducerBatch#addBuffers(List)} when later records would overflow. + * {@link ChunkedProducerBatch#addBuffers(List)} when later records would overflow. *

- * This is the "dynamic" buffer.memory allocation strategy: memory consumption scales with + * This is the "incremental" buffer.memory allocation strategy: memory consumption scales with * actual data buffered, not with {@code active_partition_count × batch.size}. *

- * PR1 scope: uncompressed only. Batch creation throws - * {@link UnsupportedOperationException} if compression is requested. The mid-batch flow - * matches the KIP: try a non-blocking pool acquire for the extension chunks; on - * {@link BufferExhaustedException}, close the existing batch (so the sender can drain it) - * and let the new record flow through the first-record path, which blocks on the pool for - * a fresh batch's chunks up to {@code max.block.ms}. + * When a record arrives mid-batch (an open batch already exists) and the batch's chunks are + * full, the flow is: under the deque lock, probe that open batch — the tail of the partition's + * deque — via the {@link #tryAppend} override, which returns + * {@link RecordAppendResult#needsExtension(int)} when the batch has logical room but lacks + * physical chunk capacity, then attempt a non-blocking pool acquire for the extra chunks. If the pool is exhausted, close the current batch (so the sender can drain it) + * and route the record through the first-record path, which blocks on the pool up to + * {@code max.block.ms}. As a result, unlike the "full" strategy, {@code send()} may block (or + * fail with {@link BufferExhaustedException}) on records other than the first of a batch. *

- * TODO: add the new buffer.memory.allocation.strategy config and instantiate this class only when set to "dynamic" - * TODO: add support for compressed data. Throws UnsupportedOperationException for now, - * and IllegalStateException if overflow detected. Follow-ups will support compressed data and growth. + * TODO: support compressed data (with mid-record growth); the constructor rejects compression for now. */ public class ChunkedRecordAccumulator extends RecordAccumulator { + /** + * Fixed size of every chunk, independent of {@code batch.size}. The incremental strategy is + * only used when {@code batch.size >= CHUNK_SIZE} (see {@code KafkaProducer}); below it a batch + * is smaller than a single chunk, so the producer uses the full strategy instead. + */ + public static final int CHUNK_SIZE = 16 * 1024; + private final ChunkedBufferPool chunkedFree; public ChunkedRecordAccumulator(LogContext logContext, @@ -75,6 +86,11 @@ public ChunkedRecordAccumulator(LogContext logContext, ChunkedBufferPool bufferPool) { super(logContext, batchSize, compression, lingerMs, retryBackoffMs, retryBackoffMaxMs, deliveryTimeoutMs, partitionerConfig, metrics, metricGrpName, time, transactionManager, bufferPool); + // TODO: drop this once the incremental strategy supports compressed data (with the + // mid-record growth fallback for compressor overshoot). + if (compression.type() != CompressionType.NONE) + throw new UnsupportedOperationException( + "Compression is not yet supported with the incremental buffer.memory allocation strategy"); this.chunkedFree = bufferPool; } @@ -132,31 +148,25 @@ public RecordAppendResult append(String topic, if (partitionChanged(topic, topicInfo, partitionInfo, dq, nowMs, cluster)) continue; - // Probe the tail directly. If the existing batch has logical room but the - // underlying stream lacks physical capacity, capture the gap and fall through - // to allocate the extension outside the deque lock. Otherwise, attempt the - // append against the existing batch. - ProducerBatch tail = dq.peekLast(); - int gap = tail == null ? 0 : tail.extensionBytesNeeded(timestamp, key, value, headers); - if (gap > 0) { - extensionBytes = gap; - } else { - extensionBytes = 0; // clear any stale value carried from a prior iteration - RecordAppendResult appendResult = tryAppend(timestamp, key, value, headers, callbacks, dq, nowMs); - if (appendResult != null) { - boolean enableSwitch = allBatchesFull(dq); - topicInfo.builtInPartitioner.updatePartitionInfo(partitionInfo, appendResult.appendedBytes, cluster, enableSwitch); - return appendResult; - } + // The tryAppend override probes the physical capacity of the tail batch + // (dq.peekLast(), the partition's open batch): a needsBufferExtension result + // means it has logical room but its chunks lack space for this record — fall + // through to allocate the gap outside the deque lock. A null result means the + // tail batch is full or absent — fall through to the first-record (new batch) path. + RecordAppendResult appendResult = tryAppend(timestamp, key, value, headers, callbacks, dq, nowMs); + if (appendResult != null && !appendResult.needsBufferExtension) { + boolean enableSwitch = allBatchesFull(dq); + topicInfo.builtInPartitioner.updatePartitionInfo(partitionInfo, appendResult.appendedBytes, cluster, enableSwitch); + return appendResult; } + extensionBytes = appendResult == null ? 0 : appendResult.extensionBytesNeeded; } if (extensionBytes > 0 && extensionChunks == null) { - // Mid-batch extension: non-blocking attempt only. The producer thread already - // holds chunks for the open batch; blocking here risks deadlock with the Sender - // thread (which frees pool memory by completing batches). If the pool is - // exhausted, we close the existing batch (so it becomes drainable) and fall - // through to the first-record blocking path on the next loop iteration. + // Mid-batch extension: non-blocking only. The thread already holds the open + // batch's chunks, so blocking here could deadlock with the Sender (which frees + // pool memory by completing batches). On exhaustion, close the batch (making it + // drainable) and fall through to the blocking first-record path next iteration. try { extensionChunks = chunkedFree.allocateChunks(extensionBytes, 0L); } catch (BufferExhaustedException e) { @@ -168,21 +178,19 @@ public RecordAppendResult append(String topic, last.closeForRecordAppends(); } } - extensionBytes = 0; continue; } nowMs = time.milliseconds(); } else if (extensionBytes == 0 && bufferStream == null) { - // First-record path: block on the pool for enough chunks to fit this record. - // Temporary while compressed data not supported. - // TODO: drop this throw and route through the - // compressed path (which adds the mid-record growth fallback for compressor overshoot). - if (compression.type() != CompressionType.NONE) { - throw new UnsupportedOperationException( - "Compression is not implemented yet for the dynamic buffer.memory allocation strategy"); - } - int size = AbstractRecords.estimateSizeInBytesUpperBound( + // First-record path: block on the pool for enough chunks to fit this record, + // sized with the same cumulative estimator used mid-batch (header + record bytes + // for NONE, ratio-adjusted when compressed) so the two stay consistent. + int recordUncompressed = AbstractRecords.recordSizeUpperBound( RecordBatch.CURRENT_MAGIC_VALUE, compression.type(), key, value, headers); + int size = MemoryRecordsBuilder.estimatedBytesWritten( + RecordBatch.CURRENT_MAGIC_VALUE, compression.type(), + CompressionRatioEstimator.estimation(topic, compression.type()), + recordUncompressed); log.trace("Allocating {} byte chunked buffer ({} byte chunks) for topic {} partition {} with remaining timeout {}ms", size, chunkedFree.poolableSize(), topic, effectivePartition, maxTimeToBlock); List initialChunks = chunkedFree.allocateChunks(size, maxTimeToBlock); @@ -196,67 +204,48 @@ public RecordAppendResult append(String topic, if (extensionChunks != null) { ProducerBatch last = dq.peekLast(); - // The off-lock allocateChunks window allows the original tail batch to be - // drained and a split batch (plain ByteBufferOutputStream) to be addFirst'd - // into an emptied deque — see splitAndReenqueue. The instanceof guard below - // is load-bearing: addBuffers is only defined on ChunkedByteBufferOutputStream; - // the else branch returns the chunks to the pool to avoid leaking them. - if (last != null && last.isWritable() - && last.bufferStream() instanceof ChunkedByteBufferOutputStream) { - ((ChunkedByteBufferOutputStream) last.bufferStream()).addBuffers(extensionChunks); + // The off-lock allocateChunks window allows the probed tail batch to be + // drained and replaced — possibly by a split batch (a plain + // ProducerBatch), which can't take extension chunks. Only attach to a + // writable chunked tail; otherwise refund the chunks and re-evaluate. + if (last instanceof ChunkedProducerBatch && last.isWritable()) { + ((ChunkedProducerBatch) last).addBuffers(extensionChunks); extensionChunks = null; - extensionBytes = 0; - // Re-probe after attach: between our first-lock probe and now we were off-lock, - // and another appender (or a tail-replacement) may have consumed capacity that - // our extension allocation was sized against. If a gap remains, loop to acquire - // more chunks against the now-current state. Without this re-probe, a - // tryAppend that passes hasRoomFor (writeLimit) but exceeds the stream's - // physical remaining() trips IllegalStateException in - // ChunkedByteBufferOutputStream.advanceToNextChunk. Termination: writeLimit is - // bounded and fixed; once hasRoomFor turns false, extensionBytesNeeded returns 0 - // and we fall through to new-batch creation. Regression test: - // testConcurrentExtensionRaceDoesNotOverflowChunkedStream. - int recheck = last.extensionBytesNeeded(timestamp, key, value, headers); - if (recheck > 0) { - extensionBytes = recheck; - continue; - } RecordAppendResult retryResult = tryAppend(timestamp, key, value, headers, callbacks, dq, nowMs); - if (retryResult != null) { + if (retryResult != null && !retryResult.needsBufferExtension) { boolean enableSwitch = allBatchesFull(dq); topicInfo.builtInPartitioner.updatePartitionInfo(partitionInfo, retryResult.appendedBytes, cluster, enableSwitch); return retryResult; } - // Batch became full via writeLimit while we were off-lock — fall through - // to new-batch creation on the next iteration. + // needsBufferExtension: a concurrent appender consumed capacity our + // extension was sized against — loop to re-probe. null: batch became + // full — loop into new-batch creation. Terminates because writeLimit is + // fixed: once full, the probe stops requesting extension. Regression: + // testConcurrentExtensionRaceDoesNotOverflowChunkedStream. continue; } // Tail is gone, closed, or non-chunked (e.g., a split batch). Return chunks to pool. for (ByteBuffer chunk : extensionChunks) chunkedFree.deallocate(chunk); extensionChunks = null; - extensionBytes = 0; - continue; - } - - // bufferStream is non-null here (first-record path). Before creating a new - // batch with it, re-probe the tail: a concurrent appender may have created a - // chunked batch we should extend instead. If so, release the oversized - // bufferStream and let the next iteration take the extension path. - ProducerBatch tail = dq.peekLast(); - if (tail != null && tail.isWritable() - && tail.bufferStream() instanceof ChunkedByteBufferOutputStream - && tail.extensionBytesNeeded(timestamp, key, value, headers) > 0) { - bufferStream.deallocate(); - bufferStream = null; continue; } + // bufferStream is non-null here (first-record path). int firstRecordSize = AbstractRecords.estimateSizeInBytesUpperBound( RecordBatch.CURRENT_MAGIC_VALUE, compression.type(), key, value, headers); final ChunkedByteBufferOutputStream batchStream = bufferStream; RecordAppendResult appendResult = appendNewBatch(topic, effectivePartition, dq, timestamp, key, value, headers, callbacks, () -> chunkedRecordsBuilder(batchStream, firstRecordSize), nowMs); + if (appendResult.needsBufferExtension) { + // A concurrent appender created a tail batch we should extend rather + // than start a new one (detected by appendNewBatch's in-lock tryAppend). + // Our bufferStream was sized for a fresh batch — release it and loop so + // the extension path allocates exactly the gap-sized chunks. + bufferStream.deallocate(); + bufferStream = null; + continue; + } if (appendResult.newBatchCreated) bufferStream = null; boolean enableSwitch = allBatchesFull(dq); @@ -276,61 +265,46 @@ public RecordAppendResult append(String topic, } /** - * Build a {@link MemoryRecordsBuilder} backed by the chunked stream. - * {@code writeLimit} is set to {@code max(batchSize, firstRecordSize)} — the same logical - * cap the legacy path uses. Physical chunk capacity grows on demand via - * {@link ChunkedByteBufferOutputStream#addBuffers(List)}; {@code writeLimit} bounds when a - * batch is considered logically full and a new one must be started. - */ - private MemoryRecordsBuilder chunkedRecordsBuilder(ChunkedByteBufferOutputStream bufferStream, - int firstRecordSize) { - int writeLimit = Math.max(batchSize, firstRecordSize); - return new MemoryRecordsBuilder(bufferStream, RecordBatch.CURRENT_MAGIC_VALUE, compression, - TimestampType.CREATE_TIME, 0L, RecordBatch.NO_TIMESTAMP, RecordBatch.NO_PRODUCER_ID, - RecordBatch.NO_PRODUCER_EPOCH, RecordBatch.NO_SEQUENCE, false, false, - RecordBatch.NO_PARTITION_LEADER_EPOCH, writeLimit); - } - - /** - * Chunked-aware deallocation for {@link ProducerBatch}. + * Try to append to a ProducerBatch, with mid-batch chunk extension support. *

- * Intercepts the {@link #isInflight()} branch of the parent's - * {@link RecordAccumulator#deallocate(ProducerBatch)} when the batch is backed by a - * {@link ChunkedByteBufferOutputStream}. The parent's defensive path credits only - * {@code initialCapacity} (one {@code chunkSize}) back to the pool and throws — that's the - * KAFKA-19012-aware safety net for the legacy single-buffer path, where the in-flight - * buffer is the pooled buffer. A chunked batch holds K chunks but the in-flight - * network bytes live in a separate flattened heap buffer (see - * {@link ChunkedByteBufferOutputStream#buffer()}), so it is safe to return all K chunks - * to the pool here. The {@code IllegalStateException} is still propagated to preserve the - * contract that deallocating an inflight batch is unexpected upstream. + * If the tail batch has logical room (writeLimit-wise) but its chunked stream lacks + * physical capacity, returns {@link RecordAppendResult#needsExtension(int)} without + * attempting the append; the caller allocates chunks outside the deque lock, attaches + * them, and retries. Otherwise defers to the parent. */ @Override - public void deallocate(ProducerBatch batch) { - if (!batch.isSplitBatch() - && !batch.isBufferDeallocated() - && batch.isInflight() - && batch.bufferStream() instanceof ChunkedByteBufferOutputStream) { - batch.markBufferDeallocated(); - deallocateBatchBuffer(batch); - throw new IllegalStateException("Attempting to deallocate a batch that is inflight. Batch is " + batch); + protected RecordAppendResult tryAppend(long timestamp, byte[] key, byte[] value, Header[] headers, + Callback callback, Deque deque, long nowMs) { + if (closed) + throw new KafkaException("Producer closed while send in progress"); + ProducerBatch last = deque.peekLast(); + // Split batches in an incremental deque are plain ProducerBatch (heap-backed, grow-on-demand) + // and never need chunk extension, so the probe only applies to chunked batches. + if (last instanceof ChunkedProducerBatch) { + int extensionBytes = ((ChunkedProducerBatch) last).extensionBytesNeeded(timestamp, key, value, headers); + if (extensionBytes > 0) + return RecordAppendResult.needsExtension(extensionBytes); } - super.deallocate(batch); + return super.tryAppend(timestamp, key, value, headers, callback, deque, nowMs); + } + + @Override + protected ProducerBatch createProducerBatch(TopicPartition tp, MemoryRecordsBuilder recordsBuilder, long nowMs) { + return new ChunkedProducerBatch(tp, recordsBuilder, nowMs); } /** - * Route chunked batch deallocation through the stream rather than the legacy - * single-buffer pool path. The flattened buffer returned by {@code batch.buffer()} is a - * fresh heap allocation and must NOT be passed to the pool — instead, the stream owns - * the chunk-sized buffers and returns them to the pool. + * Build a {@link MemoryRecordsBuilder} backed by the chunked stream. {@code writeLimit} = + * {@code max(batchSize, firstRecordSize)} (the same logical cap as the full path) bounds when + * the batch is full; physical capacity grows on demand via + * {@link ChunkedByteBufferOutputStream#addBuffers(List)}. */ - @Override - protected void deallocateBatchBuffer(ProducerBatch batch) { - if (batch.bufferStream() instanceof ChunkedByteBufferOutputStream) { - ((ChunkedByteBufferOutputStream) batch.bufferStream()).deallocate(chunkedFree); - } else { - // Split batches and any non-chunked batches go through the default deallocation. - super.deallocateBatchBuffer(batch); - } + private MemoryRecordsBuilder chunkedRecordsBuilder(ChunkedByteBufferOutputStream bufferStream, + int firstRecordSize) { + int writeLimit = Math.max(batchSize, firstRecordSize); + return new MemoryRecordsBuilder(bufferStream, RecordBatch.CURRENT_MAGIC_VALUE, compression, + TimestampType.CREATE_TIME, 0L, RecordBatch.NO_TIMESTAMP, RecordBatch.NO_PRODUCER_ID, + RecordBatch.NO_PRODUCER_EPOCH, RecordBatch.NO_SEQUENCE, false, false, + RecordBatch.NO_PARTITION_LEADER_EPOCH, writeLimit); } } diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ProducerBatch.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ProducerBatch.java index a7e47583d310b..6adac7a24424d 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ProducerBatch.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ProducerBatch.java @@ -31,7 +31,6 @@ import org.apache.kafka.common.record.internal.Record; import org.apache.kafka.common.record.internal.RecordBatch; import org.apache.kafka.common.requests.ProduceResponse; -import org.apache.kafka.common.utils.ByteBufferOutputStream; import org.apache.kafka.common.utils.ProducerIdAndEpoch; import org.apache.kafka.common.utils.Time; @@ -58,7 +57,7 @@ * * This class is not thread safe and external synchronization must be used when modifying it */ -public final class ProducerBatch { +public class ProducerBatch { private static final Logger log = LoggerFactory.getLogger(ProducerBatch.class); @@ -141,35 +140,22 @@ boolean hasLeaderChangedForTheOngoingRetry() { /** - * Bytes of physical buffer this batch needs before {@code tryAppend} could accept the given - * record. Returns 0 when the record fits (either logically full, or stream has room). - * Positive when {@code hasRoomFor} allows the record but the underlying stream lacks - * physical capacity — used by the dynamic strategy to allocate exactly the missing bytes - * before retrying. - *

- * For batches with the default {@link ByteBufferOutputStream} (legacy path), the stream's - * {@code remaining()} is large enough that this almost always returns 0. + * Return this batch's buffer memory to the pool. The default single-buffer batch returns + * the pooled buffer at its initial capacity; {@link ChunkedProducerBatch} overrides this + * to return all of its chunks. */ - int extensionBytesNeeded(long timestamp, byte[] key, byte[] value, Header[] headers) { - if (recordCount == 0) - return 0; - if (!recordsBuilder.hasRoomFor(timestamp, key, value, headers)) - return 0; - // Upper-bound estimate is fine: over-counting by header bytes only over-allocates by a - // small fraction of one chunk. Under-estimation would be unsafe. - int recordSize = AbstractRecords.estimateSizeInBytesUpperBound( - magic(), recordsBuilder.compression().type(), key, value, headers); - int gap = recordSize - recordsBuilder.bufferStream().remaining(); - return Math.max(0, gap); + protected void deallocateBuffer(BufferPool pool) { + pool.deallocate(buffer(), initialCapacity()); } /** - * The batch's underlying output stream. Exposed for the dynamic strategy to attach - * extension chunks and route deallocation through the chunked stream rather than the - * standard {@code BufferPool.deallocate(ByteBuffer, int)} path. + * Credit this batch's memory back to the pool when it is unexpectedly still inflight + * (KAFKA-19012): the buffer can't be touched (the network may still be reading it), so the + * default donates a fresh same-capacity buffer. {@link ChunkedProducerBatch} instead returns + * the actual chunks (safe — inflight bytes live in the separate flattened buffer). */ - ByteBufferOutputStream bufferStream() { - return recordsBuilder.bufferStream(); + protected void deallocateInflightBuffer(BufferPool pool) { + pool.deallocate(ByteBuffer.allocate(initialCapacity())); } /** diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/RecordAccumulator.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/RecordAccumulator.java index 29ddd07a489bc..7bb4243e24f9d 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/RecordAccumulator.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/RecordAccumulator.java @@ -398,12 +398,15 @@ protected RecordAppendResult appendNewBatch(String topic, RecordAppendResult appendResult = tryAppend(timestamp, key, value, headers, callbacks, dq, nowMs); if (appendResult != null) { - // Somebody else found us a batch, return the one we waited for! Hopefully this doesn't happen often... + // Propagate without creating a new batch: either another thread already made us a batch + // (success), or — incremental strategy — a concurrent appender created an extendable tail + // (needsBufferExtension), so the caller releases its pre-allocated buffer and retries via + // the extension path. return appendResult; } MemoryRecordsBuilder recordsBuilder = recordsBuilderSupplier.get(); - ProducerBatch batch = new ProducerBatch(new TopicPartition(topic, partition), recordsBuilder, nowMs); + ProducerBatch batch = createProducerBatch(new TopicPartition(topic, partition), recordsBuilder, nowMs); FutureRecordMetadata future = Objects.requireNonNull(batch.tryAppend(timestamp, key, value, headers, callbacks, nowMs)); @@ -413,6 +416,14 @@ protected RecordAppendResult appendNewBatch(String topic, return new RecordAppendResult(future, dq.size() > 1 || batch.isFull(), true, batch.estimatedSizeInBytes()); } + /** + * Create the {@link ProducerBatch} for a new batch. The incremental strategy overrides this to + * create a {@link ChunkedProducerBatch}. + */ + protected ProducerBatch createProducerBatch(TopicPartition tp, MemoryRecordsBuilder recordsBuilder, long nowMs) { + return new ProducerBatch(tp, recordsBuilder, nowMs); + } + /** * Check if all batches in the queue are full. */ @@ -1057,23 +1068,16 @@ public void deallocate(ProducerBatch batch) { } else { batch.markBufferDeallocated(); if (batch.isInflight()) { - // Create a fresh ByteBuffer to give to BufferPool to reuse since we can't safely call deallocate with the ProduceBatch's buffer - free.deallocate(ByteBuffer.allocate(batch.initialCapacity())); + // We can't safely deallocate the buffer of an inflight batch; the batch credits + // the memory back to the pool in a way that suits its buffer type. + batch.deallocateInflightBuffer(free); throw new IllegalStateException("Attempting to deallocate a batch that is inflight. Batch is " + batch); } - deallocateBatchBuffer(batch); + batch.deallocateBuffer(free); } } } - /** - * Return the batch's underlying buffer to the pool. - * This default implementation returns the buffer at its initial capacity (single buffer). - */ - protected void deallocateBatchBuffer(ProducerBatch batch) { - free.deallocate(batch.buffer(), batch.initialCapacity()); - } - /** * Remove from the incomplete list but do not free memory yet */ @@ -1267,17 +1271,41 @@ public static final class RecordAppendResult { public final FutureRecordMetadata future; public final boolean batchIsFull; public final boolean newBatchCreated; + /** + * Signal (incremental strategy) that the tail batch has logical room (writeLimit-wise) but + * its chunks lack physical capacity. The append was NOT attempted; the caller allocates + * {@link #extensionBytesNeeded} bytes of chunk capacity, attaches them via + * {@link ChunkedProducerBatch#addBuffers}, and retries. When {@code true}, {@code future} is null. + */ + public final boolean needsBufferExtension; + public final int extensionBytesNeeded; public final int appendedBytes; public RecordAppendResult(FutureRecordMetadata future, boolean batchIsFull, boolean newBatchCreated, int appendedBytes) { + this(future, batchIsFull, newBatchCreated, false, 0, appendedBytes); + } + + private RecordAppendResult(FutureRecordMetadata future, + boolean batchIsFull, + boolean newBatchCreated, + boolean needsBufferExtension, + int extensionBytesNeeded, + int appendedBytes) { this.future = future; this.batchIsFull = batchIsFull; this.newBatchCreated = newBatchCreated; + this.needsBufferExtension = needsBufferExtension; + this.extensionBytesNeeded = extensionBytesNeeded; this.appendedBytes = appendedBytes; } + + /** A signal-only result indicating the caller must allocate more chunk capacity. */ + static RecordAppendResult needsExtension(int extensionBytesNeeded) { + return new RecordAppendResult(null, false, false, true, extensionBytesNeeded, 0); + } } /* diff --git a/clients/src/main/java/org/apache/kafka/common/record/internal/AbstractRecords.java b/clients/src/main/java/org/apache/kafka/common/record/internal/AbstractRecords.java index d0704d18ec37a..d8eb6c4dd0be7 100644 --- a/clients/src/main/java/org/apache/kafka/common/record/internal/AbstractRecords.java +++ b/clients/src/main/java/org/apache/kafka/common/record/internal/AbstractRecords.java @@ -143,6 +143,31 @@ else if (compressionType != CompressionType.NONE) return Records.LOG_OVERHEAD + LegacyRecord.recordSize(magic, key, value); } + /** + * Upper-bound size of one record's contribution to a batch, excluding the batch header. + * For magic v2 it's {@link #estimateSizeInBytesUpperBound} minus + * {@link #recordBatchHeaderSizeInBytes}; for older magic it equals + * {@link #estimateSizeInBytesUpperBound}. Used by the incremental strategy's cumulative sizing, + * where the header is counted once via {@link MemoryRecordsBuilder#estimatedBytesWritten}. + */ + public static int recordSizeUpperBound(byte magic, CompressionType compressionType, byte[] key, + byte[] value, Header[] headers) { + return recordSizeUpperBound(magic, compressionType, Utils.wrapNullable(key), Utils.wrapNullable(value), headers); + } + + /** + * @see #recordSizeUpperBound(byte, CompressionType, byte[], byte[], Header[]) + */ + private static int recordSizeUpperBound(byte magic, CompressionType compressionType, ByteBuffer key, + ByteBuffer value, Header[] headers) { + if (magic >= RecordBatch.MAGIC_VALUE_V2) + return DefaultRecord.recordSizeUpperBound(key, value, headers); + else if (compressionType != CompressionType.NONE) + return Records.LOG_OVERHEAD + LegacyRecord.recordOverhead(magic) + LegacyRecord.recordSize(magic, key, value); + else + return Records.LOG_OVERHEAD + LegacyRecord.recordSize(magic, key, value); + } + /** * Return the size of the record batch header. * diff --git a/clients/src/main/java/org/apache/kafka/common/record/internal/MemoryRecordsBuilder.java b/clients/src/main/java/org/apache/kafka/common/record/internal/MemoryRecordsBuilder.java index 47a955cfcf18e..2bf48518b2d6f 100644 --- a/clients/src/main/java/org/apache/kafka/common/record/internal/MemoryRecordsBuilder.java +++ b/clients/src/main/java/org/apache/kafka/common/record/internal/MemoryRecordsBuilder.java @@ -212,7 +212,7 @@ public int initialCapacity() { } /** - * The underlying output stream. Exposed so the dynamic-strategy callers can attach + * The underlying output stream. Exposed so the incremental-strategy callers can attach * pre-allocated chunks to a {@code ChunkedByteBufferOutputStream} and route deallocation * polymorphically. */ @@ -824,12 +824,52 @@ private void ensureOpenForRecordBatchWrite() { * @return The estimated number of bytes written */ private int estimatedBytesWritten() { - if (compression.type() == CompressionType.NONE) { + return estimatedBytesWritten(magic, compression.type(), estimatedCompressionRatio, uncompressedRecordsSizeInBytes); + } + + /** + * Static projection of the bytes a builder with the given magic, compression type, and ratio + * would have written for {@code uncompressedRecordsSizeInBytes} of records. Exact for + * uncompressed ({@code header + uncompressedBytes}); for compressed it applies the ratio and 5% + * safety multiplier to the whole total, like the per-builder {@link #estimatedBytesWritten()}. + *

+ * Used by the incremental strategy to size chunk reservations (first record, and mid-batch via + * {@link #estimatedBytesWrittenAfter}). This is the batch's projected total output, + * distinct from the {@link #hasRoomFor} writeLimit gate, which counts the incoming record at + * full uncompressed size; the two coincide for uncompressed data and diverge under compression. + */ + public static int estimatedBytesWritten(byte magic, CompressionType compressionType, + float compressionRatio, + int uncompressedRecordsSizeInBytes) { + int batchHeaderSizeInBytes = AbstractRecords.recordBatchHeaderSizeInBytes(magic, compressionType); + if (compressionType == CompressionType.NONE) { return batchHeaderSizeInBytes + uncompressedRecordsSizeInBytes; } else { - // estimate the written bytes to the underlying byte buffer based on uncompressed written bytes - return batchHeaderSizeInBytes + (int) (uncompressedRecordsSizeInBytes * estimatedCompressionRatio * COMPRESSION_RATE_ESTIMATION_FACTOR); + return batchHeaderSizeInBytes + (int) (uncompressedRecordsSizeInBytes * compressionRatio * COMPRESSION_RATE_ESTIMATION_FACTOR); + } + } + + /** + * Projected value of {@link #estimatedBytesWritten} after appending one more record with the + * given fields, using the record's worst-case (upper-bound) per-record size. Used by the + * incremental strategy to size mid-batch chunk extensions. + */ + public int estimatedBytesWrittenAfter(byte[] key, byte[] value, Header[] headers) { + return estimatedBytesWrittenAfter(wrapNullable(key), wrapNullable(value), headers); + } + + /** + * @see #estimatedBytesWrittenAfter(byte[], byte[], Header[]) + */ + private int estimatedBytesWrittenAfter(ByteBuffer key, ByteBuffer value, Header[] headers) { + final int recordSize; + if (magic < RecordBatch.MAGIC_VALUE_V2) { + recordSize = Records.LOG_OVERHEAD + LegacyRecord.recordSize(magic, key, value); + } else { + recordSize = DefaultRecord.recordSizeUpperBound(key, value, headers); } + return estimatedBytesWritten(magic, compression.type(), estimatedCompressionRatio, + uncompressedRecordsSizeInBytes + recordSize); } /** diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/ProducerConfigTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/ProducerConfigTest.java index 5fd9ab727e046..45615be25dce0 100644 --- a/clients/src/test/java/org/apache/kafka/clients/producer/ProducerConfigTest.java +++ b/clients/src/test/java/org/apache/kafka/clients/producer/ProducerConfigTest.java @@ -129,6 +129,41 @@ public void testInvalidMetadataRecoveryStrategy() { assertTrue(ce.getMessage().contains(CommonClientConfigs.METADATA_RECOVERY_STRATEGY_CONFIG)); } + @Test + public void testDefaultBufferMemoryAllocationStrategy() { + Map configs = new HashMap<>(); + configs.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, keySerializerClass); + configs.put(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, valueSerializerClass); + configs.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9092"); + final ProducerConfig producerConfig = new ProducerConfig(configs); + assertEquals(ProducerConfig.BUFFER_MEMORY_ALLOCATION_STRATEGY_FULL, + producerConfig.getString(ProducerConfig.BUFFER_MEMORY_ALLOCATION_STRATEGY_CONFIG)); + } + + @Test + public void testValidBufferMemoryAllocationStrategy() { + Map configs = new HashMap<>(); + configs.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, keySerializerClass); + configs.put(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, valueSerializerClass); + configs.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9092"); + configs.put(ProducerConfig.BUFFER_MEMORY_ALLOCATION_STRATEGY_CONFIG, + ProducerConfig.BUFFER_MEMORY_ALLOCATION_STRATEGY_INCREMENTAL); + final ProducerConfig producerConfig = new ProducerConfig(configs); + assertEquals(ProducerConfig.BUFFER_MEMORY_ALLOCATION_STRATEGY_INCREMENTAL, + producerConfig.getString(ProducerConfig.BUFFER_MEMORY_ALLOCATION_STRATEGY_CONFIG)); + } + + @Test + public void testInvalidBufferMemoryAllocationStrategy() { + Map configs = new HashMap<>(); + configs.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, keySerializerClass); + configs.put(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, valueSerializerClass); + configs.put(ProducerConfig.BUFFER_MEMORY_ALLOCATION_STRATEGY_CONFIG, "abc"); + configs.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9092"); + ConfigException ce = assertThrows(ConfigException.class, () -> new ProducerConfig(configs)); + assertTrue(ce.getMessage().contains(ProducerConfig.BUFFER_MEMORY_ALLOCATION_STRATEGY_CONFIG)); + } + @Test public void testCaseInsensitiveSecurityProtocol() { final String saslSslLowerCase = SecurityProtocol.SASL_SSL.name.toLowerCase(Locale.ROOT); diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java index b191aadb4645b..191224cd437d8 100644 --- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java +++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java @@ -157,20 +157,6 @@ public void testMidBatchExtensionGrowsExistingBatch() throws Exception { accum.close(); } - /** Compression is rejected at batch creation in PR1 with a clear message. */ - @Test - public void testCompressionThrowsAtBatchCreation() { - int chunkSize = 256; - ChunkedRecordAccumulator accum = newAccumulator(1024, chunkSize, 16L * chunkSize, Compression.gzip().build()); - - UnsupportedOperationException ex = assertThrows(UnsupportedOperationException.class, () -> - accum.append(topic, partition1, 0L, key, new byte[32], Record.EMPTY_HEADERS, null, - maxBlockTimeMs, time.milliseconds(), cluster)); - assertTrue(ex.getMessage().contains("not implemented yet"), - "Expected 'not implemented yet' in message, got: " + ex.getMessage()); - accum.close(); - } - /** * Race between two concurrent appenders on the same chunked tail. Both threads probe * {@code extensionBytesNeeded} under the deque lock against the same physical remaining @@ -347,4 +333,81 @@ public void testBatchCloseDoesNotDeallocateChunksPrematurely() throws Exception accum.close(); } + + /** + * Chunks attached grow with the batch's cumulative projected output, not per-record. With a + * chunkSize smaller than the batch's eventual content, a sequence of small records should + * extend the chunks as the running total (header + uncompressed bytes for NONE) crosses + * chunk boundaries — not allocate the first record's worth and then never extend until + * physical capacity is exhausted (which the per-record formula would do). + */ + @Test + public void testExtensionTracksCumulativeBatchSize() throws Exception { + int chunkSize = 64; + int batchSize = 512; + ChunkedRecordAccumulator accum = newAccumulator(batchSize, chunkSize, 64L * chunkSize, Compression.NONE); + + byte[] smallValue = new byte[24]; + for (int i = 0; i < 6; i++) { + accum.append(topic, partition1, 0L, key, smallValue, Record.EMPTY_HEADERS, null, + maxBlockTimeMs, time.milliseconds(), cluster); + } + + Deque dq = batchesFor(accum, tp1); + assertEquals(1, dq.size()); + ProducerBatch batch = dq.peekFirst(); + assertNotNull(batch); + assertEquals(6, batch.recordCount); + + // batch.close() flattens chunks and writes the header. If chunks under-allocated, this + // throws (insufficient capacity in the underlying buffer). + batch.close(); + MemoryRecords records = batch.records(); + int actualSize = records.sizeInBytes(); + // Under NONE compression, estimatedBytesWritten is exact: physical bytes ≈ header + sum + // of per-record bytes. The chunks attached must cover that, so actualSize must be > + // chunkSize for a multi-record batch with non-trivial content. + assertTrue(actualSize > chunkSize, + "batch should have grown beyond a single chunk; got " + actualSize); + + accum.close(); + } + + /** + * The cumulative sizing formula counts the batch header once. The previous per-record + * formula effectively included the batch header in each record's upper bound, which + * over-allocated as records accumulated. With the cumulative formula, chunk count for N + * small records of total uncompressed size {@code U} is {@code ceil((header + U) / chunkSize)}, + * not {@code N × ceil((header + recordSize) / chunkSize)}. + */ + @Test + public void testCumulativeAccountsForHeaderOnce() throws Exception { + int chunkSize = 256; + int batchSize = 8192; + long totalMemory = 64L * chunkSize; + ChunkedBufferPool pool = new ChunkedBufferPool(totalMemory, chunkSize, metrics, time, "producer-metrics"); + ChunkedRecordAccumulator accum = new ChunkedRecordAccumulator(logContext, batchSize, Compression.NONE, + /* lingerMs */ 0, /* retryBackoffMs */ 0L, /* retryBackoffMaxMs */ 0L, + /* deliveryTimeoutMs */ 3200, metrics, "producer-metrics", time, + /* transactionManager */ null, pool); + + long beforeAlloc = pool.availableMemory(); + // First record establishes the batch. Each subsequent small record contributes its + // uncompressed bytes to the cumulative target; the batch header is NOT re-counted. + byte[] smallValue = new byte[8]; + for (int i = 0; i < 4; i++) { + accum.append(topic, partition1, 0L, key, smallValue, Record.EMPTY_HEADERS, null, + maxBlockTimeMs, time.milliseconds(), cluster); + } + + // Each per-record append adds ~10-30 bytes (key + value + record overhead, V2). Cumulative + // total for 4 records is well below chunkSize=256, so only 1 chunk should ever be attached. + // The per-record formula (header counted once per record) would have allocated more. + long held = beforeAlloc - pool.availableMemory(); + assertEquals(chunkSize, held, + "cumulative formula should hold exactly one chunk for a small-record batch; " + + "header double-counting (per-record formula) would inflate this"); + + accum.close(); + } } From 8c82cd261c2c79d091697f35779bf7b5067cb635 Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Thu, 18 Jun 2026 15:27:50 -0400 Subject: [PATCH 04/61] integration tests --- .../kafka/api/BaseProducerSendTest.scala | 10 +- ...ncrementalAllocationProducerSendTest.scala | 94 +++++++++++++++++++ .../scala/unit/kafka/utils/TestUtils.scala | 4 +- 3 files changed, 106 insertions(+), 2 deletions(-) create mode 100644 core/src/test/scala/integration/kafka/api/IncrementalAllocationProducerSendTest.scala diff --git a/core/src/test/scala/integration/kafka/api/BaseProducerSendTest.scala b/core/src/test/scala/integration/kafka/api/BaseProducerSendTest.scala index 71d0dd8273992..1d50bd8910549 100644 --- a/core/src/test/scala/integration/kafka/api/BaseProducerSendTest.scala +++ b/core/src/test/scala/integration/kafka/api/BaseProducerSendTest.scala @@ -101,6 +101,13 @@ abstract class BaseProducerSendTest extends KafkaServerTestHarness { super.tearDown() } + /** + * Additional producer properties applied to every producer created via [[createProducer]]. + * Subclasses override this to run the whole suite under a different producer configuration + * (e.g. a different buffer.memory.allocation.strategy). + */ + protected def producerOverrides: Properties = new Properties() + protected def createProducer(lingerMs: Int = 0, deliveryTimeoutMs: Int = 2 * 60 * 1000, batchSize: Int = 16384, @@ -117,7 +124,8 @@ abstract class BaseProducerSendTest extends KafkaServerTestHarness { deliveryTimeoutMs = deliveryTimeoutMs, maxBlockMs = maxBlockMs, batchSize = batchSize, - bufferSize = bufferSize) + bufferSize = bufferSize, + additionalProperties = Some(producerOverrides)) registerProducer(producer) } diff --git a/core/src/test/scala/integration/kafka/api/IncrementalAllocationProducerSendTest.scala b/core/src/test/scala/integration/kafka/api/IncrementalAllocationProducerSendTest.scala new file mode 100644 index 0000000000000..74bcd892813c7 --- /dev/null +++ b/core/src/test/scala/integration/kafka/api/IncrementalAllocationProducerSendTest.scala @@ -0,0 +1,94 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.api + +import java.util.concurrent.ThreadLocalRandom +import java.util.{Properties, List => JList} +import kafka.utils.{TestInfoUtils, TestUtils} +import org.apache.kafka.clients.producer.{ProducerConfig, ProducerRecord} +import org.apache.kafka.common.TopicPartition +import org.apache.kafka.common.network.ListenerName +import org.apache.kafka.common.security.auth.SecurityProtocol +import org.junit.jupiter.api.Assertions.{assertArrayEquals, assertEquals} +import org.junit.jupiter.params.ParameterizedTest +import org.junit.jupiter.params.provider.MethodSource + +/** + * Runs the whole [[BaseProducerSendTest]] suite with the producer configured with the incremental + * buffer.memory allocation strategy, plus a few scenarios specific to chunked allocation. + */ +class IncrementalAllocationProducerSendTest extends BaseProducerSendTest { + + override protected def producerOverrides: Properties = { + val props = new Properties() + props.put(ProducerConfig.BUFFER_MEMORY_ALLOCATION_STRATEGY_CONFIG, + ProducerConfig.BUFFER_MEMORY_ALLOCATION_STRATEGY_INCREMENTAL) + props + } + + // The incremental strategy does not support compression yet; the base test would fail at + // producer construction. Overriding without the test annotations removes it from this + // subclass's run. TODO: remove this override when compression support lands. + override def testSendCompressedMessageWithCreateTime(groupProtocol: String): Unit = {} + + // batch.size=0 is below the internal chunk size, so the incremental strategy falls back to the full + // allocation path (a batch is smaller than one chunk). Verifies that fallback yields a working producer. + @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedGroupProtocolNames) + @MethodSource(Array("getTestGroupProtocolParametersConsumerGroupProtocolOnly")) + def testBatchSizeZero(groupProtocol: String): Unit = { + sendAndVerify(createProducer( + lingerMs = Int.MaxValue, + deliveryTimeoutMs = Int.MaxValue, + batchSize = 0)) + } + + @ParameterizedTest(name = TestInfoUtils.TestWithParameterizedGroupProtocolNames) + @MethodSource(Array("getTestGroupProtocolParametersConsumerGroupProtocolOnly")) + def testSendLargeRecordsSpanningMultipleChunks(groupProtocol: String): Unit = { + TestUtils.createTopicWithAdmin(admin, topic, brokers, controllerServers, 1, 2) + val partition = 0 + val tp = new TopicPartition(topic, partition) + val valueSize = 600_000 // far larger than one chunk, so each record spans many chunks + // TODO: also exercise the compressed codecs here once the incremental strategy supports compression. + + val producer = createProducer(batchSize = 1024 * 1024, bufferSize = 8L * 1024 * 1024) + val value = randomBytes(valueSize) + val metadata = producer.send(new ProducerRecord(topic, partition, "key".getBytes, value)).get + assertEquals(valueSize, metadata.serializedValueSize) + + // Verify the exact bytes round-trip. + val consumer = TestUtils.createConsumer( + bootstrapServers(listenerName = ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT)), + groupProtocolFromTestParameters()) + try { + consumer.assign(JList.of(tp)) + consumer.seekToBeginning(JList.of(tp)) + val consumed = TestUtils.consumeRecords(consumer, 1) + assertEquals(metadata.offset, consumed.head.offset) + assertArrayEquals(value, consumed.head.value) + } finally { + consumer.close() + } + } + + private def randomBytes(size: Int): Array[Byte] = { + val bytes = new Array[Byte](size) + ThreadLocalRandom.current().nextBytes(bytes) + bytes + } +} diff --git a/core/src/test/scala/unit/kafka/utils/TestUtils.scala b/core/src/test/scala/unit/kafka/utils/TestUtils.scala index bc41c2ed2aeeb..088fe6384ba6d 100755 --- a/core/src/test/scala/unit/kafka/utils/TestUtils.scala +++ b/core/src/test/scala/unit/kafka/utils/TestUtils.scala @@ -491,7 +491,8 @@ object TestUtils extends Logging { saslProperties: Option[Properties] = None, keySerializer: Serializer[K] = new ByteArraySerializer, valueSerializer: Serializer[V] = new ByteArraySerializer, - enableIdempotence: Boolean = false): KafkaProducer[K, V] = { + enableIdempotence: Boolean = false, + additionalProperties: Option[Properties] = None): KafkaProducer[K, V] = { val producerProps = new Properties producerProps.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, brokerList) producerProps.put(ProducerConfig.ACKS_CONFIG, acks.toString) @@ -504,6 +505,7 @@ object TestUtils extends Logging { producerProps.put(ProducerConfig.BATCH_SIZE_CONFIG, batchSize.toString) producerProps.put(ProducerConfig.COMPRESSION_TYPE_CONFIG, compressionType) producerProps.put(ProducerConfig.ENABLE_IDEMPOTENCE_CONFIG, enableIdempotence.toString) + additionalProperties.foreach(producerProps ++= _) producerProps ++= JaasTestUtils.producerSecurityConfigs(securityProtocol, OptionConverters.toJava(trustStoreFile), OptionConverters.toJava(saslProperties)) new KafkaProducer[K, V](producerProps, keySerializer, valueSerializer) } From b71d5b716d3427f2cdaea14839ffa16ff83a6ae6 Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Thu, 18 Jun 2026 15:54:53 -0400 Subject: [PATCH 05/61] refactor validation & upd docs --- .../kafka/clients/producer/KafkaProducer.java | 2 +- .../producer/internals/ChunkedBufferPool.java | 26 ++++++++++--------- .../internals/ChunkedProducerBatch.java | 5 ++-- .../internals/ChunkedRecordAccumulator.java | 4 +-- .../ChunkedRecordAccumulatorTest.java | 22 ++++++++-------- 5 files changed, 31 insertions(+), 28 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/KafkaProducer.java b/clients/src/main/java/org/apache/kafka/clients/producer/KafkaProducer.java index 0990a68439171..51835478ea5df 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/KafkaProducer.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/KafkaProducer.java @@ -471,7 +471,7 @@ public KafkaProducer(Properties properties, Serializer keySerializer, Seriali // Use the chunked path only when a batch is at least one full chunk // (batch.size >= CHUNK_SIZE). Below that, a batch can't fill even one chunk, so chunking // would reserve a whole chunk for a batch capped under it (over-reservation); the full - // strategy (reserving up to batch.size per batch) is both leaner and identical to trunk. + // strategy (reserving max(batch.size, record size) per batch) is both leaner and identical to trunk. if (incremental && batchSize >= ChunkedRecordAccumulator.CHUNK_SIZE) { this.accumulator = new ChunkedRecordAccumulator(logContext, batchSize, diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java index 9fa367d80e889..1c40451c28dbe 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java @@ -61,22 +61,11 @@ public ChunkedBufferPool(long memory, int chunkSize, Metrics metrics, Time time, public List allocateChunks(int totalSize, long maxTimeToBlockMs) throws InterruptedException { if (totalSize <= 0) throw new IllegalArgumentException("totalSize must be positive: " + totalSize); - if (totalSize > totalMemory()) - throw new IllegalArgumentException("Attempt to allocate " + totalSize - + " bytes across chunks, but there is a hard limit of " - + totalMemory() + " on memory allocations."); + throwIfChunksNeededExceedsPool(totalSize); int chunkSize = poolableSize(); int numChunks = (int) (((long) totalSize + chunkSize - 1L) / chunkSize); long memoryRequired = (long) numChunks * chunkSize; - // Rounding totalSize up to whole chunks can require more pool memory than buffer.memory - // holds even when totalSize alone fits (the single-buffer "full" path reserves exactly the - // record size, so it never hits this). Reject it here to fail fast, instead of blocking for - // max.block.ms and then throwing BufferExhausted for a request that can never be satisfied. - if (memoryRequired > totalMemory()) - throw new IllegalArgumentException("Attempt to allocate " + numChunks + " chunks of " - + chunkSize + " bytes (" + memoryRequired + " bytes total), but there is a hard limit of " - + totalMemory() + " on memory allocations."); // Chunks pulled from the free list; the remaining bytes are reserved against // nonPooledAvailableMemory and materialized as raw allocations after the lock is released. @@ -196,4 +185,17 @@ public List allocateChunks(int totalSize, long maxTimeToBlockMs) thr } } } + + /** + * Throw if rounding {@code totalSize} up to whole chunks would exceed the pool. + */ + private void throwIfChunksNeededExceedsPool(int totalSize) { + int chunkSize = poolableSize(); + int numChunks = (int) (((long) totalSize + chunkSize - 1L) / chunkSize); + long memoryRequired = (long) numChunks * chunkSize; + if (memoryRequired > totalMemory()) + throw new IllegalArgumentException("Attempt to allocate " + totalSize + " bytes (" + + numChunks + " chunks of " + chunkSize + " = " + memoryRequired + " bytes), but the " + + "hard limit on memory allocations is " + totalMemory() + "."); + } } diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedProducerBatch.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedProducerBatch.java index e80ed7861826d..78a187cf431d7 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedProducerBatch.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedProducerBatch.java @@ -45,8 +45,9 @@ public ChunkedProducerBatch(TopicPartition tp, MemoryRecordsBuilder recordsBuild /** * Bytes of physical buffer this batch needs before {@code tryAppend} could accept the given - * record. Returns 0 when the record fits (either logically full, or the stream's combined - * chunk capacity has room). Positive when {@code hasRoomFor} allows the record but the + * record. Returns 0 when no extension is needed: the batch is empty (first record), it is + * logically full, or the stream's combined chunk capacity already has room. Positive when + * {@code hasRoomFor} allows the record but the * chunks lack physical capacity — the accumulator allocates exactly the missing bytes * (rounded up to whole chunks) and attaches them via {@link #addBuffers} before retrying. */ diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java index a587754d80c44..a16f85dbf9ef7 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java @@ -42,8 +42,8 @@ /** * A {@link RecordAccumulator} variant that uses chunked buffer allocation: each new batch - * pre-allocates {@code ceil(recordSize / chunkSize)} chunks from a {@link ChunkedBufferPool}, - * and grows mid-batch by attaching additional chunks via + * pre-allocates the chunks needed for the first record's estimated size (batch header plus the + * record), and grows mid-batch by attaching additional chunks via * {@link ChunkedProducerBatch#addBuffers(List)} when later records would overflow. *

* This is the "incremental" buffer.memory allocation strategy: memory consumption scales with diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java index 191224cd437d8..35bc979ae86a4 100644 --- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java +++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java @@ -246,10 +246,10 @@ public List allocateChunks(int totalSize, long maxTimeToBlockMs) thr /** * Regression guard: an inflight chunked batch returns all K chunks to the pool on the * expiration / broker-disconnect path, not just one {@code initialCapacity} chunkSize. - * The parent {@code RecordAccumulator.deallocate} for inflight batches credits only - * {@code initialCapacity} back to the pool and throws; {@link ChunkedRecordAccumulator} - * overrides {@code deallocate} to return all chunks before the throw. Removing that - * override would re-introduce a (K−1)×chunkSize leak and fail this test. + * For an inflight batch the parent {@code RecordAccumulator.deallocate} calls + * {@code batch.deallocateInflightBuffer(pool)} and then throws; {@link ChunkedProducerBatch} + * overrides that hook to return all chunks (instead of donating one fresh buffer). Removing + * that override would re-introduce a (K−1)×chunkSize leak and fail this test. */ @Test public void testInflightExpirationReturnsAllChunksToPool() throws Exception { @@ -297,11 +297,11 @@ private Deque batchesFor(RecordAccumulator accum, TopicPartition /** * Sender's drain (and any other caller of {@code batch.close()}) cascades through * {@code MemoryRecordsBuilder.closeForRecordAppends()} → - * {@code appendStream.close()} → {@code bufferStream.close()}. The chunked stream's - * buffer must remain readable after that cascade so the builder can flatten chunks into - * the heap buffer it sends over the network. If the chunked stream's {@code close()} - * prematurely returns chunks to the pool, the flattened buffer is empty and the header - * write fails (or produces an empty record set). + * {@code appendStream.close()} → {@code bufferStream.close()}. The chunk data must remain + * readable through that cascade so the builder can flatten the chunks into the heap buffer it + * sends over the network — chunks are returned to the pool only at batch completion + * (deallocate), never at close. If they were released at close, the flattened buffer would be + * empty and the header write would fail (or produce an empty record set). */ @Test public void testBatchCloseDoesNotDeallocateChunksPrematurely() throws Exception { @@ -316,7 +316,7 @@ public void testBatchCloseDoesNotDeallocateChunksPrematurely() throws Exception ProducerBatch batch = dq.peekFirst(); assertNotNull(batch); - // Matches the call sites in RecordAccumulator.drain (line 941) and Sender. + // Matches the call sites in RecordAccumulator.drain and Sender. batch.close(); MemoryRecords records = batch.records(); assertTrue(records.sizeInBytes() > 0, @@ -364,7 +364,7 @@ public void testExtensionTracksCumulativeBatchSize() throws Exception { batch.close(); MemoryRecords records = batch.records(); int actualSize = records.sizeInBytes(); - // Under NONE compression, estimatedBytesWritten is exact: physical bytes ≈ header + sum + // Under NONE compression, estimatedBytesWritten is exact: physical bytes = header + sum // of per-record bytes. The chunks attached must cover that, so actualSize must be > // chunkSize for a multi-record batch with non-trivial content. assertTrue(actualSize > chunkSize, From bdc14e90f471b4c11fd881c98de4f165c009afc5 Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Mon, 22 Jun 2026 09:54:12 -0400 Subject: [PATCH 06/61] update config --- .../apache/kafka/clients/producer/KafkaProducer.java | 3 +-- .../kafka/clients/producer/ProducerConfig.java | 12 ++++++------ 2 files changed, 7 insertions(+), 8 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/KafkaProducer.java b/clients/src/main/java/org/apache/kafka/clients/producer/KafkaProducer.java index 51835478ea5df..3b13014768281 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/KafkaProducer.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/KafkaProducer.java @@ -470,8 +470,7 @@ public KafkaProducer(Properties properties, Serializer keySerializer, Seriali } // Use the chunked path only when a batch is at least one full chunk // (batch.size >= CHUNK_SIZE). Below that, a batch can't fill even one chunk, so chunking - // would reserve a whole chunk for a batch capped under it (over-reservation); the full - // strategy (reserving max(batch.size, record size) per batch) is both leaner and identical to trunk. + // would over-reserve. if (incremental && batchSize >= ChunkedRecordAccumulator.CHUNK_SIZE) { this.accumulator = new ChunkedRecordAccumulator(logContext, batchSize, diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/ProducerConfig.java b/clients/src/main/java/org/apache/kafka/clients/producer/ProducerConfig.java index 286f3e37dc823..696e9c103118f 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/ProducerConfig.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/ProducerConfig.java @@ -234,12 +234,12 @@ public class ProducerConfig extends AbstractConfig { public static final String BUFFER_MEMORY_ALLOCATION_STRATEGY_INCREMENTAL = "incremental"; private static final String BUFFER_MEMORY_ALLOCATION_STRATEGY_DOC = "Controls how the producer allocates memory from " + BUFFER_MEMORY_CONFIG + " for record batches. The following values are supported: " + "

    " - + "
  • " + BUFFER_MEMORY_ALLOCATION_STRATEGY_FULL + ": reserve " + BATCH_SIZE_CONFIG + " bytes up front for each in-progress batch, " - + "regardless of the actual size of the records appended to it.
    • " - + "
  • " + BUFFER_MEMORY_ALLOCATION_STRATEGY_INCREMENTAL + ": allocate memory on demand as records are appended, growing each batch " - + "up to " + BATCH_SIZE_CONFIG + " bytes. This keeps memory usage proportional to the data actually buffered rather than to the number of " - + "active partitions, which allows larger " + BATCH_SIZE_CONFIG + " values (e.g. for high-latency clusters) without reserving " - + "" + BATCH_SIZE_CONFIG + " bytes for every active partition.
    • " + + "
  • " + BUFFER_MEMORY_ALLOCATION_STRATEGY_FULL + ": reserves a full " + BATCH_SIZE_CONFIG + " up front when a batch is created, " + + "regardless of how much data it ends up holding. Pool memory therefore scales with the number of active partitions.
    • " + + "
  • " + BUFFER_MEMORY_ALLOCATION_STRATEGY_INCREMENTAL + ": allocates memory on demand as records are appended, growing a batch " + + "up to " + BATCH_SIZE_CONFIG + ". Pool memory therefore scales with the data actually buffered rather than the number of active " + + "partitions, allowing larger " + BATCH_SIZE_CONFIG + " values (e.g. for high-latency clusters) without reserving " + + "" + BATCH_SIZE_CONFIG + " for every active partition.
    • " + "
"; /** retry.backoff.ms */ From 4e914999d0af74b25e3714d8cc118194b550ea49 Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Mon, 22 Jun 2026 14:55:55 -0400 Subject: [PATCH 07/61] fix for partition change & misc --- .../producer/internals/BufferPool.java | 1 - .../producer/internals/ChunkedBufferPool.java | 14 ++-- .../ChunkedByteBufferOutputStream.java | 21 ++++- .../internals/ChunkedRecordAccumulator.java | 12 ++- .../ChunkedByteBufferOutputStreamTest.java | 22 +++++ .../ChunkedRecordAccumulatorTest.java | 82 +++++++++++++++++++ 6 files changed, 139 insertions(+), 13 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/BufferPool.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/BufferPool.java index f376c0fe2f4a3..c515cf0a8d79f 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/BufferPool.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/BufferPool.java @@ -59,7 +59,6 @@ public class BufferPool { private final Metrics metrics; protected final Time time; private final Sensor waitTime; - /** True once {@link #close()} has been invoked. */ protected boolean closed; /** diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java index 1c40451c28dbe..c81fef81cffbe 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java @@ -30,10 +30,7 @@ /** * A {@link BufferPool} dedicated to chunk-sized buffer reuse (chunk size = {@link #poolableSize()}). *

- * Adds {@link #allocateChunks(int, long)} to acquire multiple chunks atomically: one lock - * acquisition and a single {@link Condition} on {@link #waiters} (FIFO-fair against single-chunk - * acquirers), with any failure (timeout, close, OOM) refunding the whole reservation before the - * exception propagates. No partial holds are visible during the wait. + * Adds {@link #allocateChunks(int, long)} to acquire multiple chunks atomically. */ public class ChunkedBufferPool extends BufferPool { @@ -44,9 +41,10 @@ public ChunkedBufferPool(long memory, int chunkSize, Metrics metrics, Time time, /** * Allocate {@code ceil(totalSize / chunkSize)} chunk-sized buffers atomically, mirroring * {@link BufferPool#allocate}: satisfied immediately if memory is available, else blocks up to - * {@code maxTimeToBlockMs} for the whole request (FIFO on {@link #waiters}). The reservation is - * tracked as bytes against {@link #nonPooledAvailableMemory} plus chunks polled from - * {@link #free}; on any failure path every reserved byte is returned and the next waiter signaled. + * {@code maxTimeToBlockMs} for the whole request (FIFO on {@link #waiters}). + * The reservation is tracked as bytes against {@link #nonPooledAvailableMemory} plus chunks polled + * from {@link #free}. Any failure refunds the whole reservation and signals + * the next waiter before the exception propagates, so no partial holds are visible during the wait. * * @param totalSize minimum total bytes of capacity required across the returned chunks * @param maxTimeToBlockMs maximum time in milliseconds to block waiting for memory @@ -67,7 +65,7 @@ public List allocateChunks(int totalSize, long maxTimeToBlockMs) thr int numChunks = (int) (((long) totalSize + chunkSize - 1L) / chunkSize); long memoryRequired = (long) numChunks * chunkSize; - // Chunks pulled from the free list; the remaining bytes are reserved against + // Chunks taken from the free list. The remaining bytes are reserved against // nonPooledAvailableMemory and materialized as raw allocations after the lock is released. List pooled = new ArrayList<>(numChunks); diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java index 3b9a9ed42ff5e..89110438381f9 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java @@ -59,9 +59,8 @@ public class ChunkedByteBufferOutputStream extends ByteBufferOutputStream { * @param chunkSize the size of each chunk in bytes * @param pool the buffer pool used for deallocation */ - @SuppressWarnings("this-escape") public ChunkedByteBufferOutputStream(List initialChunks, int chunkSize, BufferPool pool) { - super(initialChunks.get(0)); + super(validatedFirstChunk(initialChunks, chunkSize)); this.chunkSize = chunkSize; this.pool = pool; this.chunks = new ArrayList<>(initialChunks); @@ -70,6 +69,21 @@ public ChunkedByteBufferOutputStream(List initialChunks, int chunkSi this.dirty = true; } + /** + * Validates the chunk contract: {@code initialChunks} non-empty, each chunk's capacity equal to + * {@code chunkSize}. Returns the first chunk. + */ + private static ByteBuffer validatedFirstChunk(List initialChunks, int chunkSize) { + if (initialChunks == null || initialChunks.isEmpty()) + throw new IllegalArgumentException("initialChunks must be non-empty"); + for (ByteBuffer chunk : initialChunks) { + if (chunk.capacity() != chunkSize) + throw new IllegalArgumentException("each chunk must have capacity " + chunkSize + + ", but found a chunk of capacity " + chunk.capacity()); + } + return initialChunks.get(0); + } + @Override public void write(int b) { ensureChunkCapacity(1); @@ -223,7 +237,8 @@ public int limit() { @Override public int initialCapacity() { - return chunks.isEmpty() ? chunkSize : chunks.get(0).capacity(); + // Every chunk is chunkSize (enforced by the constructor), so this is constant. + return chunkSize; } @Override diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java index a16f85dbf9ef7..2e74f3b096261 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java @@ -199,8 +199,18 @@ public RecordAppendResult append(String topic, } synchronized (dq) { - if (partitionChanged(topic, topicInfo, partitionInfo, dq, nowMs, cluster)) + if (partitionChanged(topic, topicInfo, partitionInfo, dq, nowMs, cluster)) { + // The partition switched while we allocated extension chunks off-lock. They + // were sized against the previous partition's tail batch, so they must not be + // attached to a different partition's tail — refund them and let the next + // iteration re-probe the new partition from scratch. + if (extensionChunks != null) { + for (ByteBuffer chunk : extensionChunks) + chunkedFree.deallocate(chunk); + extensionChunks = null; + } continue; + } if (extensionChunks != null) { ProducerBatch last = dq.peekLast(); diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java index 569c3bbcfcbdc..c96824b09c137 100644 --- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java +++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java @@ -49,6 +49,28 @@ private List chunks(ChunkedBufferPool pool, int chunkSize, int count return pool.allocateChunks(chunkSize * count, 100); } + /** + * The constructor fails fast with {@link IllegalArgumentException} when its chunk contract is + * violated — a null or empty list, or a chunk whose capacity differs from {@code chunkSize} — + * rather than the bare NPE / IndexOutOfBoundsException that {@code initialChunks.get(0)} would + * otherwise throw from the {@code super(...)} call. + */ + @Test + public void testConstructorRejectsInvalidChunks() throws Exception { + int chunkSize = 16; + ChunkedBufferPool p = pool(64, chunkSize); + + assertThrows(IllegalArgumentException.class, + () -> new ChunkedByteBufferOutputStream(null, chunkSize, p)); + assertThrows(IllegalArgumentException.class, + () -> new ChunkedByteBufferOutputStream(Collections.emptyList(), chunkSize, p)); + + // A chunk whose capacity doesn't match chunkSize violates the contract. + List wrongSize = Collections.singletonList(ByteBuffer.allocate(chunkSize + 1)); + assertThrows(IllegalArgumentException.class, + () -> new ChunkedByteBufferOutputStream(wrongSize, chunkSize, p)); + } + /** Writes that fit in the initial chunk: stream produces the bytes back via buffer(). */ @Test public void testSingleChunkWriteRoundtrip() throws Exception { diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java index 35bc979ae86a4..97bdaa09c29c0 100644 --- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java +++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java @@ -44,6 +44,7 @@ import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; import java.util.stream.Collectors; import java.util.stream.Stream; @@ -243,6 +244,87 @@ public List allocateChunks(int totalSize, long maxTimeToBlockMs) thr } } + /** + * Regression: when the sticky partition switches while extension chunks are held off-lock, + * those chunks — sized against the previous partition's tail batch — must be refunded + * to the pool on the retry, not carried into the next iteration and attached to a different + * partition's tail (which would over-extend the wrong batch). + *

+ * Deterministic reproduction without threads: the extension path is the only one that calls + * {@code allocateChunks} non-blocking ({@code maxTimeToBlockMs == 0}), so the pool flags when an + * extension allocation has happened; a {@code partitionChanged} override then fires exactly one + * spurious switch on the very next check — which is the second-sync-block check, with the + * extension chunks held. The fix refunds those chunks before the retry; without it they are + * carried and attached instead, so no refund is observed during the append. + */ + @Test + public void testPartitionSwitchRefundsHeldExtensionChunks() throws Exception { + int chunkSize = 128; + long totalMemory = 32L * chunkSize; + + AtomicBoolean extensionAllocated = new AtomicBoolean(false); + AtomicInteger chunkDeallocations = new AtomicInteger(0); + + ChunkedBufferPool pool = new ChunkedBufferPool(totalMemory, chunkSize, metrics, time, "producer-metrics") { + @Override + public List allocateChunks(int totalSize, long maxTimeToBlockMs) throws InterruptedException { + List chunks = super.allocateChunks(totalSize, maxTimeToBlockMs); + // The mid-batch extension path is the only non-blocking caller. + if (maxTimeToBlockMs == 0L) + extensionAllocated.set(true); + return chunks; + } + + @Override + public void deallocate(ByteBuffer buffer) { + chunkDeallocations.incrementAndGet(); + super.deallocate(buffer); + } + }; + + AtomicBoolean switchFired = new AtomicBoolean(false); + 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 boolean partitionChanged(String topic, TopicInfo topicInfo, + BuiltInPartitioner.StickyPartitionInfo partitionInfo, + Deque deque, long nowMs, Cluster cluster) { + // Inject one spurious switch, but only once an extension allocation has happened — + // i.e. on the second-sync-block check, while extension chunks are held. + if (extensionAllocated.get() && switchFired.compareAndSet(false, true)) + return true; + return super.partitionChanged(topic, topicInfo, partitionInfo, deque, nowMs, cluster); + } + }; + try { + // Warmup: tiny first record establishes a chunked tail (first-record path, blocking + // allocate — does not flag extensionAllocated). + accum.append(topic, partition1, 0L, key, new byte[1], Record.EMPTY_HEADERS, null, + maxBlockTimeMs, time.milliseconds(), cluster); + int deallocBeforeExtend = chunkDeallocations.get(); + + // Second record overflows the chunk → extension allocated off-lock → injected switch + // fires in the second sync block with those chunks held. + accum.append(topic, partition1, 0L, key, new byte[200], Record.EMPTY_HEADERS, null, + maxBlockTimeMs, time.milliseconds(), cluster); + + assertTrue(switchFired.get(), "the injected partition switch should have fired"); + assertTrue(chunkDeallocations.get() > deallocBeforeExtend, + "extension chunks held at the partition switch must be refunded to the pool; " + + "without the fix they are carried and attached instead"); + + // The record still appends correctly after the switch-retry. + Deque dq = batchesFor(accum, tp1); + assertEquals(1, dq.size()); + assertEquals(2, dq.peekFirst().recordCount, + "second record should still land in the batch after the switch-retry"); + } finally { + accum.close(); + } + } + /** * Regression guard: an inflight chunked batch returns all K chunks to the pool on the * expiration / broker-disconnect path, not just one {@code initialCapacity} chunkSize. From 9ccd47dc673aeaaa44d6c2b9ae6bef28ba811d77 Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Tue, 23 Jun 2026 12:59:43 -0400 Subject: [PATCH 08/61] test updates & docs --- .../producer/internals/ChunkedBufferPool.java | 6 +- .../ChunkedByteBufferOutputStream.java | 51 +++---- .../internals/ChunkedProducerBatch.java | 14 +- .../internals/ChunkedRecordAccumulator.java | 70 ++++----- .../producer/internals/RecordAccumulator.java | 12 +- .../record/internal/AbstractRecords.java | 7 +- .../record/internal/MemoryRecordsBuilder.java | 16 +-- .../internals/ChunkedBufferPoolTest.java | 51 ++----- .../ChunkedByteBufferOutputStreamTest.java | 135 ++++++++---------- .../ChunkedRecordAccumulatorTest.java | 86 +++-------- 10 files changed, 167 insertions(+), 281 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java index c81fef81cffbe..f2c8ba35f1568 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java @@ -125,8 +125,6 @@ public List allocateChunks(int totalSize, long maxTimeToBlockMs) thr remainingTimeToBlockNs -= timeNs; // Take pool chunks first, then reserve non-pool bytes for the remainder. - // Pool chunks don't bump `accumulated` because their refund path is - // `free.addFirst` in the outer catch, not nonPool. while (pooled.size() < numChunks && (long) (pooled.size() + 1) * chunkSize + accumulated <= memoryRequired && !free.isEmpty()) { @@ -140,7 +138,7 @@ public List allocateChunks(int totalSize, long maxTimeToBlockMs) thr accumulated += got; } } - // Loop terminated normally — clear the rollback marker (mirrors BufferPool.allocate). + // Clear the rollback tracker. accumulated = 0; } finally { // On failure (timeout / close / interrupt), refund the non-pool bytes taken. @@ -163,7 +161,7 @@ public List allocateChunks(int totalSize, long maxTimeToBlockMs) thr } } - // Allocate raw chunks for the reserved non-pooled portion outside the lock. On OOM, + // Allocate raw chunks for the reserved non-pooled portion outside the lock. On error, // refund all reserved bytes (memoryRequired) and let the next waiter try. int chunksStillNeeded = numChunks - pooled.size(); List result = new ArrayList<>(numChunks); diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java index 89110438381f9..97b89cd1b4416 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java @@ -23,21 +23,20 @@ import java.util.List; /** - * A {@link ByteBufferOutputStream} backed by a linked list of fixed-size chunks instead of a - * single re-allocated buffer. + * A {@link ByteBufferOutputStream} backed by a linked list of fixed-size chunks instead of a single + * re-allocated buffer. Chunks are supplied by the caller (initial chunks via the constructor, + * additional chunks via {@link #addBuffers(List)}). *

- * This stream does not grow on its own. Chunks are supplied by the caller — initial - * chunks via the constructor, additional chunks via {@link #addBuffers(List)}. When a write's - * size exceeds the stream's {@link #remaining()} (sum of free bytes across all attached - * chunks), {@link IllegalStateException} is thrown. The caller is responsible for attaching - * additional chunks before any such write. - *

- * Automatic mid-write growth (allocating from inside {@code write()}) is a follow-up tied to - * compression support, where compressor buffering can make a write exceed its reservation; until - * then the write path stays allocation-free and deterministic. - *

- * When {@link #buffer()} is called (typically at batch close), all chunks are flattened into a - * single contiguous ByteBuffer — exactly one copy operation. + * Current/temporary behavior: + *

    + *
  • The stream does not grow on its own: a write whose size exceeds the remaining free bytes + * across all attached chunks throws {@link IllegalStateException}, so the caller must attach + * enough chunks before any such write. + * TODO: KAFKA-20579 (automatic mid-write growth for compression support).
    • + *
  • {@link #buffer()} returns the written bytes as a single contiguous {@link ByteBuffer}, + * flattening all chunks into a new buffer with an extra copy. + * TODO: KAFKA-20580 (remove the extra copy on send, scatter-gather send).
    • + *
*/ public class ChunkedByteBufferOutputStream extends ByteBufferOutputStream { @@ -51,10 +50,10 @@ public class ChunkedByteBufferOutputStream extends ByteBufferOutputStream { /** * Constructs a chunked output stream backed by the given pre-allocated chunks. Ownership of - * {@code initialChunks} transfers to this stream — they will be returned to the pool via - * {@link #deallocate()}. + * {@code initialChunks} transfers to this stream (they will be returned to the pool via + * {@link #deallocate()}). * - * @param initialChunks pre-allocated chunks; must be non-empty; each chunk's capacity must + * @param initialChunks pre-allocated chunks. Must be non-empty and each chunk's capacity must * equal {@code chunkSize} * @param chunkSize the size of each chunk in bytes * @param pool the buffer pool used for deallocation @@ -123,17 +122,12 @@ private void ensureChunkCapacity(int needed) { } /** - * Advances {@code currentChunk} to the next pre-supplied chunk; throws if none is left. The - * caller must attach additional chunks via {@link #addBuffers(List)} before any write that - * exceeds {@link #remaining()}. (Mid-record growth — non-blocking pool then heap — is a - * follow-up that lands with compression support.) + * Advances {@code currentChunk} to the next pre-supplied chunk. */ private void advanceToNextChunk() { if (currentChunkIndex + 1 >= chunks.size()) { - throw new IllegalStateException( - "No more chunks available; the write exceeded the stream's remaining capacity. " - + "Caller must obtain additional chunks (e.g. from ChunkedBufferPool) and attach them " - + "via addBuffers before any write whose size exceeds remaining()"); + // TODO: KAFKA-20579. With compression support, grow here instead of throwing. + throw new IllegalStateException("write exceeded the stream's remaining chunk capacity"); } currentChunkIndex++; currentChunk = chunks.get(currentChunkIndex); @@ -154,9 +148,7 @@ public ByteBuffer buffer() { } // TODO: KAFKA-20687. This flatten runs at batch close, when the chunk set is final. // Today all chunks (used and unused) are returned to the pool only when the batch - // completes (via deallocate(pool)). Consider releasing the fully-unused chunks - // early, at close, rather than holding them until completion — removing them from - // `chunks` here so the completion-time deallocate(pool) does not double-return them. + // completes (via deallocate(pool)). Consider releasing the fully-unused chunks early, here. int totalSize = 0; for (ByteBuffer chunk : chunks) { totalSize += chunk.position(); @@ -237,14 +229,13 @@ public int limit() { @Override public int initialCapacity() { - // Every chunk is chunkSize (enforced by the constructor), so this is constant. return chunkSize; } @Override public void ensureRemaining(int remainingBytesRequired) { // A single call can guarantee at most `chunkSize` of space (the stream advances one chunk - // at a time); callers needing more attach chunks via addBuffers first. write(byte[]) loops + // at a time). Callers needing more attach chunks via addBuffers first. write(byte[]) loops // across chunks, so contiguous capacity isn't required. ensureChunkCapacity(Math.min(remainingBytesRequired, chunkSize)); } diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedProducerBatch.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedProducerBatch.java index 78a187cf431d7..05c290737fe24 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedProducerBatch.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedProducerBatch.java @@ -25,7 +25,7 @@ import java.util.List; /** - * A {@link ProducerBatch} for the incremental (chunked) buffer.memory allocation strategy, backed + * A {@link ProducerBatch} for the incremental buffer.memory allocation strategy, backed * by a {@link MemoryRecordsBuilder} whose stream is a {@link ChunkedByteBufferOutputStream}. * It adds mid-batch chunk extension support ({@link #extensionBytesNeeded} / * {@link #addBuffers}) and overrides the pool deallocation hooks so all chunks are returned to @@ -44,12 +44,12 @@ public ChunkedProducerBatch(TopicPartition tp, MemoryRecordsBuilder recordsBuild } /** - * Bytes of physical buffer this batch needs before {@code tryAppend} could accept the given - * record. Returns 0 when no extension is needed: the batch is empty (first record), it is - * logically full, or the stream's combined chunk capacity already has room. Positive when - * {@code hasRoomFor} allows the record but the - * chunks lack physical capacity — the accumulator allocates exactly the missing bytes - * (rounded up to whole chunks) and attaches them via {@link #addBuffers} before retrying. + * Bytes of chunk capacity this batch needs before {@code tryAppend} could accept the given + * record. Returns 0 when no extension is needed: the batch is empty (first record), it is at + * its batch-size limit, or the attached chunk capacity already has room. Positive when the + * record is within the batch-size limit but the attached chunks lack capacity — the accumulator + * then allocates exactly the missing bytes (rounded up to whole chunks) and attaches them via + * {@link #addBuffers} before retrying. */ int extensionBytesNeeded(long timestamp, byte[] key, byte[] value, Header[] headers) { if (recordCount == 0) diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java index 2e74f3b096261..297dac0dca91f 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java @@ -41,22 +41,12 @@ import java.util.List; /** - * A {@link RecordAccumulator} variant that uses chunked buffer allocation: each new batch - * pre-allocates the chunks needed for the first record's estimated size (batch header plus the - * record), and grows mid-batch by attaching additional chunks via - * {@link ChunkedProducerBatch#addBuffers(List)} when later records would overflow. + * A {@link RecordAccumulator} variant that backs each batch with fixed-size chunks drawn from a + * {@link ChunkedBufferPool}, attaching more chunks on demand as records are appended instead of + * reserving {@code batch.size} per batch up front. Buffered memory therefore scales with the data + * actually written rather than with {@code active_partition_count × batch.size}. *

- * This is the "incremental" buffer.memory allocation strategy: memory consumption scales with - * actual data buffered, not with {@code active_partition_count × batch.size}. - *

- * When a record arrives mid-batch (an open batch already exists) and the batch's chunks are - * full, the flow is: under the deque lock, probe that open batch — the tail of the partition's - * deque — via the {@link #tryAppend} override, which returns - * {@link RecordAppendResult#needsExtension(int)} when the batch has logical room but lacks - * physical chunk capacity, then attempt a non-blocking pool acquire for the extra chunks. If the pool is exhausted, close the current batch (so the sender can drain it) - * and route the record through the first-record path, which blocks on the pool up to - * {@code max.block.ms}. As a result, unlike the "full" strategy, {@code send()} may block (or - * fail with {@link BufferExhaustedException}) on records other than the first of a batch. + * See {@link #append} and {@link #tryAppend} for how batches are created and grown. *

* TODO: support compressed data (with mid-record growth); the constructor rejects compression for now. */ @@ -128,7 +118,7 @@ public RecordAppendResult append(String topic, appendsInProgress.incrementAndGet(); ChunkedByteBufferOutputStream bufferStream = null; List extensionChunks = null; - int extensionBytes = 0; + int extensionBytes; if (headers == null) headers = Record.EMPTY_HEADERS; try { while (true) { @@ -148,11 +138,11 @@ public RecordAppendResult append(String topic, if (partitionChanged(topic, topicInfo, partitionInfo, dq, nowMs, cluster)) continue; - // The tryAppend override probes the physical capacity of the tail batch - // (dq.peekLast(), the partition's open batch): a needsBufferExtension result - // means it has logical room but its chunks lack space for this record — fall - // through to allocate the gap outside the deque lock. A null result means the - // tail batch is full or absent — fall through to the first-record (new batch) path. + // The tryAppend override checks the open batch (dq.peekLast()) for chunk + // capacity: a needsBufferExtension result means it is within its batch-size limit + // but its chunks lack capacity for this record — fall through to allocate the gap + // outside the deque lock. A null result means there is no open batch (it is full + // or absent) — fall through to the first-record (new batch) path. RecordAppendResult appendResult = tryAppend(timestamp, key, value, headers, callbacks, dq, nowMs); if (appendResult != null && !appendResult.needsBufferExtension) { boolean enableSwitch = allBatchesFull(dq); @@ -162,7 +152,7 @@ public RecordAppendResult append(String topic, extensionBytes = appendResult == null ? 0 : appendResult.extensionBytesNeeded; } - if (extensionBytes > 0 && extensionChunks == null) { + if (extensionBytes > 0) { // Mid-batch extension: non-blocking only. The thread already holds the open // batch's chunks, so blocking here could deadlock with the Sender (which frees // pool memory by completing batches). On exhaustion, close the batch (making it @@ -201,9 +191,9 @@ public RecordAppendResult append(String topic, synchronized (dq) { if (partitionChanged(topic, topicInfo, partitionInfo, dq, nowMs, cluster)) { // The partition switched while we allocated extension chunks off-lock. They - // were sized against the previous partition's tail batch, so they must not be - // attached to a different partition's tail — refund them and let the next - // iteration re-probe the new partition from scratch. + // were sized against the previous partition's open batch, so they must not be + // attached to a different partition's open batch — refund them and let the next + // iteration re-check the new partition from scratch. if (extensionChunks != null) { for (ByteBuffer chunk : extensionChunks) chunkedFree.deallocate(chunk); @@ -214,10 +204,10 @@ public RecordAppendResult append(String topic, if (extensionChunks != null) { ProducerBatch last = dq.peekLast(); - // The off-lock allocateChunks window allows the probed tail batch to be + // The off-lock allocateChunks window allows the open batch we checked to be // drained and replaced — possibly by a split batch (a plain // ProducerBatch), which can't take extension chunks. Only attach to a - // writable chunked tail; otherwise refund the chunks and re-evaluate. + // writable chunked batch; otherwise refund the chunks and re-evaluate. if (last instanceof ChunkedProducerBatch && last.isWritable()) { ((ChunkedProducerBatch) last).addBuffers(extensionChunks); extensionChunks = null; @@ -228,27 +218,28 @@ public RecordAppendResult append(String topic, return retryResult; } // needsBufferExtension: a concurrent appender consumed capacity our - // extension was sized against — loop to re-probe. null: batch became + // extension was sized against — loop to re-check. null: batch became // full — loop into new-batch creation. Terminates because writeLimit is - // fixed: once full, the probe stops requesting extension. Regression: - // testConcurrentExtensionRaceDoesNotOverflowChunkedStream. + // fixed: once full, the check stops requesting extension. continue; } - // Tail is gone, closed, or non-chunked (e.g., a split batch). Return chunks to pool. + // The open batch is gone, closed, or non-chunked (e.g., a split batch). Return chunks to pool. for (ByteBuffer chunk : extensionChunks) chunkedFree.deallocate(chunk); extensionChunks = null; continue; } - // bufferStream is non-null here (first-record path). + // First-record path: extensionChunks == null here implies extensionBytes == 0, + // so bufferStream was allocated (this iteration or carried from a prior one). + assert bufferStream != null; int firstRecordSize = AbstractRecords.estimateSizeInBytesUpperBound( RecordBatch.CURRENT_MAGIC_VALUE, compression.type(), key, value, headers); final ChunkedByteBufferOutputStream batchStream = bufferStream; RecordAppendResult appendResult = appendNewBatch(topic, effectivePartition, dq, timestamp, key, value, headers, callbacks, () -> chunkedRecordsBuilder(batchStream, firstRecordSize), nowMs); if (appendResult.needsBufferExtension) { - // A concurrent appender created a tail batch we should extend rather + // A concurrent appender created an open batch we should extend rather // than start a new one (detected by appendNewBatch's in-lock tryAppend). // Our bufferStream was sized for a fresh batch — release it and loop so // the extension path allocates exactly the gap-sized chunks. @@ -277,8 +268,8 @@ public RecordAppendResult append(String topic, /** * Try to append to a ProducerBatch, with mid-batch chunk extension support. *

- * If the tail batch has logical room (writeLimit-wise) but its chunked stream lacks - * physical capacity, returns {@link RecordAppendResult#needsExtension(int)} without + * If the open batch is within its batch-size limit but its chunked stream lacks chunk + * capacity, returns {@link RecordAppendResult#needsExtension(int)} without * attempting the append; the caller allocates chunks outside the deque lock, attaches * them, and retries. Otherwise defers to the parent. */ @@ -289,7 +280,7 @@ protected RecordAppendResult tryAppend(long timestamp, byte[] key, byte[] value, throw new KafkaException("Producer closed while send in progress"); ProducerBatch last = deque.peekLast(); // Split batches in an incremental deque are plain ProducerBatch (heap-backed, grow-on-demand) - // and never need chunk extension, so the probe only applies to chunked batches. + // and never need chunk extension, so the check only applies to chunked batches. if (last instanceof ChunkedProducerBatch) { int extensionBytes = ((ChunkedProducerBatch) last).extensionBytesNeeded(timestamp, key, value, headers); if (extensionBytes > 0) @@ -304,10 +295,9 @@ protected ProducerBatch createProducerBatch(TopicPartition tp, MemoryRecordsBuil } /** - * Build a {@link MemoryRecordsBuilder} backed by the chunked stream. {@code writeLimit} = - * {@code max(batchSize, firstRecordSize)} (the same logical cap as the full path) bounds when - * the batch is full; physical capacity grows on demand via - * {@link ChunkedByteBufferOutputStream#addBuffers(List)}. + * Build a {@link MemoryRecordsBuilder} backed by the chunked stream. Its {@code writeLimit} + * bounds when the batch is full (the same batch-size limit as the full path), while chunk + * capacity grows on demand via {@link ChunkedByteBufferOutputStream#addBuffers(List)}. */ private MemoryRecordsBuilder chunkedRecordsBuilder(ChunkedByteBufferOutputStream bufferStream, int firstRecordSize) { diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/RecordAccumulator.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/RecordAccumulator.java index 7bb4243e24f9d..4125b29acaf5c 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/RecordAccumulator.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/RecordAccumulator.java @@ -379,9 +379,9 @@ public RecordAppendResult append(String topic, * @param headers the Headers for the record * @param callbacks The callbacks to execute * @param recordsBuilderSupplier Supplies the {@link MemoryRecordsBuilder} for the new - * batch. Invoked lazily, only when this method is committed to creating a new - * batch (i.e. after the in-lock concurrent-append race check has lost). The - * chunked subclass passes a supplier that produces a stream-backed builder. + * batch. Invoked lazily, only when a new batch is actually created. The chunked + * subclass passes a supplier that produces a builder backed by a + * {@link ChunkedByteBufferOutputStream}. * @param nowMs The current time, in milliseconds */ protected RecordAppendResult appendNewBatch(String topic, @@ -399,7 +399,7 @@ protected RecordAppendResult appendNewBatch(String topic, RecordAppendResult appendResult = tryAppend(timestamp, key, value, headers, callbacks, dq, nowMs); if (appendResult != null) { // Propagate without creating a new batch: either another thread already made us a batch - // (success), or — incremental strategy — a concurrent appender created an extendable tail + // (success), or — incremental strategy — a concurrent appender created an extendable open batch // (needsBufferExtension), so the caller releases its pre-allocated buffer and retries via // the extension path. return appendResult; @@ -1272,8 +1272,8 @@ public static final class RecordAppendResult { public final boolean batchIsFull; public final boolean newBatchCreated; /** - * Signal (incremental strategy) that the tail batch has logical room (writeLimit-wise) but - * its chunks lack physical capacity. The append was NOT attempted; the caller allocates + * Signal (incremental strategy) that the open batch is within its batch-size limit but its + * chunks lack capacity. The append was NOT attempted; the caller allocates * {@link #extensionBytesNeeded} bytes of chunk capacity, attaches them via * {@link ChunkedProducerBatch#addBuffers}, and retries. When {@code true}, {@code future} is null. */ diff --git a/clients/src/main/java/org/apache/kafka/common/record/internal/AbstractRecords.java b/clients/src/main/java/org/apache/kafka/common/record/internal/AbstractRecords.java index d8eb6c4dd0be7..4e6e32d75a556 100644 --- a/clients/src/main/java/org/apache/kafka/common/record/internal/AbstractRecords.java +++ b/clients/src/main/java/org/apache/kafka/common/record/internal/AbstractRecords.java @@ -144,11 +144,8 @@ else if (compressionType != CompressionType.NONE) } /** - * Upper-bound size of one record's contribution to a batch, excluding the batch header. - * For magic v2 it's {@link #estimateSizeInBytesUpperBound} minus - * {@link #recordBatchHeaderSizeInBytes}; for older magic it equals - * {@link #estimateSizeInBytesUpperBound}. Used by the incremental strategy's cumulative sizing, - * where the header is counted once via {@link MemoryRecordsBuilder#estimatedBytesWritten}. + * Upper-bound on the bytes a single record contributes to a batch, excluding the batch header. + * Used by the incremental strategy, which counts the header once, separately. */ public static int recordSizeUpperBound(byte magic, CompressionType compressionType, byte[] key, byte[] value, Header[] headers) { diff --git a/clients/src/main/java/org/apache/kafka/common/record/internal/MemoryRecordsBuilder.java b/clients/src/main/java/org/apache/kafka/common/record/internal/MemoryRecordsBuilder.java index 2bf48518b2d6f..611654ed3e595 100644 --- a/clients/src/main/java/org/apache/kafka/common/record/internal/MemoryRecordsBuilder.java +++ b/clients/src/main/java/org/apache/kafka/common/record/internal/MemoryRecordsBuilder.java @@ -212,9 +212,8 @@ public int initialCapacity() { } /** - * The underlying output stream. Exposed so the incremental-strategy callers can attach - * pre-allocated chunks to a {@code ChunkedByteBufferOutputStream} and route deallocation - * polymorphically. + * The underlying output stream, exposed so the incremental strategy can manage its + * chunk-backed stream. */ public ByteBufferOutputStream bufferStream() { return bufferStream; @@ -828,15 +827,12 @@ private int estimatedBytesWritten() { } /** - * Static projection of the bytes a builder with the given magic, compression type, and ratio - * would have written for {@code uncompressedRecordsSizeInBytes} of records. Exact for - * uncompressed ({@code header + uncompressedBytes}); for compressed it applies the ratio and 5% - * safety multiplier to the whole total, like the per-builder {@link #estimatedBytesWritten()}. + * Returns the projected number of bytes the builder would write for + * {@code uncompressedRecordsSizeInBytes} of records under the given magic, compression type, + * and ratio: exact for uncompressed, a ratio-aware estimate for compressed. *

* Used by the incremental strategy to size chunk reservations (first record, and mid-batch via - * {@link #estimatedBytesWrittenAfter}). This is the batch's projected total output, - * distinct from the {@link #hasRoomFor} writeLimit gate, which counts the incoming record at - * full uncompressed size; the two coincide for uncompressed data and diverge under compression. + * {@link #estimatedBytesWrittenAfter}). */ public static int estimatedBytesWritten(byte magic, CompressionType compressionType, float compressionRatio, diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPoolTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPoolTest.java index e0de8ae19e4dc..45e9aceb12cf5 100644 --- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPoolTest.java +++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPoolTest.java @@ -19,6 +19,7 @@ import org.apache.kafka.clients.producer.BufferExhaustedException; import org.apache.kafka.common.metrics.Metrics; import org.apache.kafka.common.utils.MockTime; +import org.apache.kafka.test.TestUtils; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.Test; @@ -32,6 +33,7 @@ import java.util.concurrent.atomic.AtomicReference; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertThrows; @@ -41,7 +43,6 @@ public class ChunkedBufferPoolTest { private final MockTime time = new MockTime(); private final Metrics metrics = new Metrics(time); - private final String metricGroup = "producer-metrics"; @AfterEach public void teardown() { @@ -49,6 +50,7 @@ public void teardown() { } private ChunkedBufferPool pool(long totalMemory, int chunkSize) { + String metricGroup = "producer-metrics"; return new ChunkedBufferPool(totalMemory, chunkSize, metrics, time, metricGroup); } @@ -62,7 +64,7 @@ public void testAllocateOneChunk() throws Exception { assertEquals(chunkSize, chunks.get(0).capacity()); } - /** Total size that's not a multiple of chunk size rounds up; returned chunks cover it. */ + /** Total size that's not a multiple of chunk size rounds up */ @Test public void testAllocateRoundsUpToChunkBoundary() throws Exception { int chunkSize = 64; @@ -105,14 +107,12 @@ public void testDeallocationRestoresMemory() throws Exception { assertEquals(total, p.availableMemory()); } - /** Requesting more than total memory throws IllegalArgumentException. */ @Test public void testRejectsRequestExceedingTotalMemory() { ChunkedBufferPool p = pool(128, 64); assertThrows(IllegalArgumentException.class, () -> p.allocateChunks(129, 100)); } - /** Non-positive totalSize is rejected. */ @Test public void testRejectsNonPositiveRequest() { ChunkedBufferPool p = pool(128, 64); @@ -122,7 +122,7 @@ public void testRejectsNonPositiveRequest() { /** * When a multi-chunk request can't be fully satisfied within the deadline, chunks already - * acquired during the call must be returned to the pool — the caller must not leak memory. + * acquired during the call must be returned to the pool. */ @Test public void testRollbackOnPartialFailure() throws Exception { @@ -145,10 +145,7 @@ public void testRollbackOnPartialFailure() throws Exception { /** * No partial chunk holds during the wait. While a multi-chunk request blocks, the pool's - * {@code availableMemory()} must reflect bytes the waiter has not yet "earned" — i.e., the - * reservation is via the accumulator and not via chunks pulled out of the pool's accounting. - * Distinguishes the atomic implementation from the older loop-with-rollback (which held - * partials off the accounting between iterations). + * {@code availableMemory()} must reflect bytes the waiter has not yet "earned". */ @Test public void testNoPartialHoldsDuringWait() throws Exception { @@ -173,10 +170,7 @@ public void testNoPartialHoldsDuringWait() throws Exception { }, "chunked-waiter"); t.start(); // Wait until the thread has joined the waiters queue. - long deadlineMs = System.currentTimeMillis() + 2_000; - while (p.queued() == 0 && System.currentTimeMillis() < deadlineMs) - Thread.sleep(5); - assertEquals(1, p.queued(), "waiter should be parked on the pool's queue"); + TestUtils.waitForCondition(() -> p.queued() == 1, "waiter should be parked on the pool's queue"); // While the waiter is parked, the pool's available memory must still report the // 1 chunk's worth — the waiter has NOT consumed any of it (no partial hold). @@ -235,16 +229,10 @@ public void testFifoFairnessAcrossChunkedAndSingleChunkRequests() throws Excepti tMulti.start(); assertTrue(multiStarted.await(2, TimeUnit.SECONDS)); // Wait for tMulti to be parked before starting tSingle, so the FIFO order is deterministic. - long deadlineMs = System.currentTimeMillis() + 2_000; - while (p.queued() == 0 && System.currentTimeMillis() < deadlineMs) - Thread.sleep(5); - assertEquals(1, p.queued(), "multi-chunk waiter should be the only one queued"); + TestUtils.waitForCondition(() -> p.queued() == 1, "multi-chunk waiter should be the only one queued"); tSingle.start(); // Wait for tSingle to also be parked. - deadlineMs = System.currentTimeMillis() + 2_000; - while (p.queued() < 2 && System.currentTimeMillis() < deadlineMs) - Thread.sleep(5); - assertEquals(2, p.queued(), "single-chunk waiter joined after the multi-chunk one"); + TestUtils.waitForCondition(() -> p.queued() == 2, "single-chunk waiter joined after the multi-chunk one"); // Now free chunks one at a time, allowing the multi-chunk request to accumulate. p.deallocate(h1); @@ -302,10 +290,7 @@ public void testSlowPathDoesNotDoubleRefundPolledChunksOnTimeout() throws Except }, "partial-holder"); t.start(); - long deadline = System.currentTimeMillis() + 2_000; - while (p.queued() == 0 && System.currentTimeMillis() < deadline) - Thread.sleep(5); - assertEquals(1, p.queued()); + TestUtils.waitForCondition(() -> p.queued() == 1, "waiter should be parked on the pool's queue"); // Hand the waiter 2 chunks — it polls them into `pooled` then awaits again for the 3rd. p.deallocate(h1); @@ -315,8 +300,7 @@ public void testSlowPathDoesNotDoubleRefundPolledChunksOnTimeout() throws Except Thread.sleep(100); t.join(5_000); - assertTrue(err.get() instanceof BufferExhaustedException, - "expected BufferExhaustedException, got " + err.get()); + assertInstanceOf(BufferExhaustedException.class, err.get(), "expected BufferExhaustedException, got " + err.get()); // After the throw, exactly 2 chunkSizes of physical memory should be back in the pool // (h1 + h2). h3 is still held outside the pool. The bug double-refunds: it credits @@ -358,10 +342,7 @@ public void testSlowPathSuccessDoesNotCorruptAccounting() throws Exception { }, "slow-success"); t.start(); - long deadline = System.currentTimeMillis() + 2_000; - while (p.queued() == 0 && System.currentTimeMillis() < deadline) - Thread.sleep(5); - assertEquals(1, p.queued()); + TestUtils.waitForCondition(() -> p.queued() == 1, "waiter should be parked on the pool's queue"); // Deallocate 2 chunks → waiter wakes, polls both, accumulated reaches memoryRequired, // exits normally. The success path's finally must NOT refund or decrement nonPool. @@ -406,16 +387,12 @@ public void testCloseDuringAtomicWait() throws Exception { } }, "close-during-wait"); t.start(); - long deadlineMs = System.currentTimeMillis() + 2_000; - while (p.queued() == 0 && System.currentTimeMillis() < deadlineMs) - Thread.sleep(5); - assertEquals(1, p.queued()); + TestUtils.waitForCondition(() -> p.queued() == 1, "waiter should be parked on the pool's queue"); // Close the pool: signals all waiters; the waiter must throw KafkaException. p.close(); t.join(5_000); - assertTrue(err.get() instanceof KafkaException, - "expected KafkaException, got " + err.get()); + assertInstanceOf(KafkaException.class, err.get(), "expected KafkaException, got " + err.get()); p.deallocate(h1); p.deallocate(h2); } diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java index c96824b09c137..145876730e977 100644 --- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java +++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java @@ -34,7 +34,6 @@ public class ChunkedByteBufferOutputStreamTest { private final MockTime time = new MockTime(); private final Metrics metrics = new Metrics(time); - private final String metricGroup = "test"; @AfterEach public void teardown() { @@ -42,6 +41,7 @@ public void teardown() { } private ChunkedBufferPool pool(long total, int chunkSize) { + String metricGroup = "test"; return new ChunkedBufferPool(total, chunkSize, metrics, time, metricGroup); } @@ -49,14 +49,8 @@ private List chunks(ChunkedBufferPool pool, int chunkSize, int count return pool.allocateChunks(chunkSize * count, 100); } - /** - * The constructor fails fast with {@link IllegalArgumentException} when its chunk contract is - * violated — a null or empty list, or a chunk whose capacity differs from {@code chunkSize} — - * rather than the bare NPE / IndexOutOfBoundsException that {@code initialChunks.get(0)} would - * otherwise throw from the {@code super(...)} call. - */ @Test - public void testConstructorRejectsInvalidChunks() throws Exception { + public void testConstructorRejectsInvalidChunks() { int chunkSize = 16; ChunkedBufferPool p = pool(64, chunkSize); @@ -71,123 +65,110 @@ public void testConstructorRejectsInvalidChunks() throws Exception { () -> new ChunkedByteBufferOutputStream(wrongSize, chunkSize, p)); } - /** Writes that fit in the initial chunk: stream produces the bytes back via buffer(). */ @Test public void testSingleChunkWriteRoundtrip() throws Exception { int chunkSize = 16; ChunkedBufferPool p = pool(64, chunkSize); - ChunkedByteBufferOutputStream stream = new ChunkedByteBufferOutputStream(chunks(p, chunkSize, 1), chunkSize, p); + try (ChunkedByteBufferOutputStream stream = new ChunkedByteBufferOutputStream(chunks(p, chunkSize, 1), chunkSize, p)) { - byte[] payload = new byte[]{1, 2, 3, 4, 5}; - stream.write(payload, 0, payload.length); + byte[] payload = new byte[]{1, 2, 3, 4, 5}; + stream.write(payload, 0, payload.length); - ByteBuffer flat = stream.buffer(); - flat.flip(); - byte[] out = new byte[flat.remaining()]; - flat.get(out); - assertArrayEquals(payload, out); + ByteBuffer flat = stream.buffer(); + flat.flip(); + byte[] out = new byte[flat.remaining()]; + flat.get(out); + assertArrayEquals(payload, out); - stream.deallocate(); + stream.deallocate(); + } } - /** Writes that span multiple pre-supplied chunks land in subsequent chunks in order. */ @Test - public void testWriteSpansChunks() throws Exception { + public void testWriteAcrossMultipleChunks() throws Exception { int chunkSize = 8; ChunkedBufferPool p = pool(64, chunkSize); - ChunkedByteBufferOutputStream stream = new ChunkedByteBufferOutputStream(chunks(p, chunkSize, 3), chunkSize, p); + try (ChunkedByteBufferOutputStream stream = new ChunkedByteBufferOutputStream(chunks(p, chunkSize, 3), chunkSize, p)) { - byte[] payload = new byte[20]; - for (int i = 0; i < payload.length; i++) payload[i] = (byte) i; - stream.write(payload, 0, payload.length); + byte[] payload = new byte[20]; + for (int i = 0; i < payload.length; i++) payload[i] = (byte) i; + stream.write(payload, 0, payload.length); - ByteBuffer flat = stream.buffer(); - flat.flip(); - byte[] out = new byte[flat.remaining()]; - flat.get(out); - assertArrayEquals(payload, out); + ByteBuffer flat = stream.buffer(); + flat.flip(); + byte[] out = new byte[flat.remaining()]; + flat.get(out); + assertArrayEquals(payload, out); - stream.deallocate(); + stream.deallocate(); + } } - /** remaining() sums current + queued chunk capacities and shrinks as writes consume bytes. */ @Test - public void testRemainingSumsChunks() throws Exception { + public void testRemainingSumsFreeBytesAcrossChunks() throws Exception { int chunkSize = 8; ChunkedBufferPool p = pool(64, chunkSize); - ChunkedByteBufferOutputStream stream = new ChunkedByteBufferOutputStream(chunks(p, chunkSize, 2), chunkSize, p); + try (ChunkedByteBufferOutputStream stream = new ChunkedByteBufferOutputStream(chunks(p, chunkSize, 2), chunkSize, p)) { - assertEquals(2 * chunkSize, stream.remaining()); - stream.write(new byte[3], 0, 3); - assertEquals(2 * chunkSize - 3, stream.remaining()); - stream.write(new byte[chunkSize], 0, chunkSize); // crosses into chunk 2 - assertEquals(chunkSize - 3, stream.remaining()); + assertEquals(2 * chunkSize, stream.remaining()); + stream.write(new byte[3], 0, 3); + assertEquals(2 * chunkSize - 3, stream.remaining()); + stream.write(new byte[chunkSize], 0, chunkSize); // crosses into chunk 2 + assertEquals(chunkSize - 3, stream.remaining()); - stream.deallocate(); + stream.deallocate(); + } } - /** addBuffers extends capacity so a write that would have overflowed now fits. */ @Test public void testAddBuffersExtendsStream() throws Exception { int chunkSize = 8; ChunkedBufferPool p = pool(64, chunkSize); - ChunkedByteBufferOutputStream stream = new ChunkedByteBufferOutputStream(chunks(p, chunkSize, 1), chunkSize, p); + try (ChunkedByteBufferOutputStream stream = new ChunkedByteBufferOutputStream(chunks(p, chunkSize, 1), chunkSize, p)) { - // Fill the initial chunk. - stream.write(new byte[chunkSize], 0, chunkSize); - assertEquals(0, stream.remaining()); + // Fill the initial chunk. + stream.write(new byte[chunkSize], 0, chunkSize); + assertEquals(0, stream.remaining()); - // Extend and write more — must land in the new chunk. - stream.addBuffers(Collections.singletonList(p.allocate(chunkSize, 100))); - assertEquals(chunkSize, stream.remaining()); + // Extend and write more — must land in the new chunk. + stream.addBuffers(Collections.singletonList(p.allocate(chunkSize, 100))); + assertEquals(chunkSize, stream.remaining()); - byte[] more = new byte[]{9, 9, 9}; - stream.write(more, 0, more.length); - assertEquals(chunkSize - 3, stream.remaining()); + byte[] more = new byte[]{9, 9, 9}; + stream.write(more, 0, more.length); + assertEquals(chunkSize - 3, stream.remaining()); - stream.deallocate(); + stream.deallocate(); + } } - /** Writing past all pre-supplied (and added) chunks throws — uncompressed-only contract. */ - @Test - public void testOverflowThrows() throws Exception { - int chunkSize = 4; - ChunkedBufferPool p = pool(64, chunkSize); - ChunkedByteBufferOutputStream stream = new ChunkedByteBufferOutputStream(chunks(p, chunkSize, 1), chunkSize, p); - - stream.write(new byte[chunkSize], 0, chunkSize); - assertThrows(IllegalStateException.class, () -> stream.write((byte) 1)); - - stream.deallocate(); - } - - /** position(int) at construction walks across pre-supplied chunks. */ @Test public void testPositionWalksAcrossChunks() throws Exception { int chunkSize = 4; ChunkedBufferPool p = pool(32, chunkSize); - ChunkedByteBufferOutputStream stream = new ChunkedByteBufferOutputStream(chunks(p, chunkSize, 3), chunkSize, p); + try (ChunkedByteBufferOutputStream stream = new ChunkedByteBufferOutputStream(chunks(p, chunkSize, 3), chunkSize, p)) { - stream.position(6); // straddles chunk 0 (4 bytes) and chunk 1 (2 bytes) - assertEquals(6, stream.position()); - assertEquals(6, stream.remaining()); + stream.position(6); // straddles chunk 0 (4 bytes) and chunk 1 (2 bytes) + assertEquals(6, stream.position()); + assertEquals(6, stream.remaining()); - stream.deallocate(); + stream.deallocate(); + } } - /** deallocate returns all chunks to the pool. */ @Test public void testDeallocateReturnsAllChunks() throws Exception { int chunkSize = 8; long total = 64; ChunkedBufferPool p = pool(total, chunkSize); List initial = chunks(p, chunkSize, 3); - ChunkedByteBufferOutputStream stream = new ChunkedByteBufferOutputStream(initial, chunkSize, p); - // Also attach one extra chunk so deallocate must handle the added-buffer case too. - stream.addBuffers(Collections.singletonList(p.allocate(chunkSize, 100))); - assertEquals(total - 4L * chunkSize, p.availableMemory()); + try (ChunkedByteBufferOutputStream stream = new ChunkedByteBufferOutputStream(initial, chunkSize, p)) { + // Also attach one extra chunk so deallocate must handle the added-buffer case too. + stream.addBuffers(Collections.singletonList(p.allocate(chunkSize, 100))); + assertEquals(total - 4L * chunkSize, p.availableMemory()); - stream.deallocate(); + stream.deallocate(); + } assertEquals(total, p.availableMemory()); } -} +} \ No newline at end of file diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java index 97bdaa09c29c0..d658e201083fe 100644 --- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java +++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java @@ -35,7 +35,6 @@ import java.nio.ByteBuffer; import java.util.ArrayList; -import java.util.Arrays; import java.util.Collections; import java.util.Deque; import java.util.List; @@ -65,7 +64,7 @@ public class ChunkedRecordAccumulatorTest { private final PartitionMetadata partMetadata1 = new PartitionMetadata(Errors.NONE, tp1, Optional.of(node1.id()), Optional.empty(), null, null, null); - private final List partMetadatas = new ArrayList<>(Arrays.asList(partMetadata1)); + private final List partMetadatas = new ArrayList<>(List.of(partMetadata1)); private final Map nodes = Stream.of(node1).collect(Collectors.toMap(Node::id, java.util.function.Function.identity())); private final MetadataSnapshot metadataCache = new MetadataSnapshot(null, nodes, partMetadatas, @@ -91,7 +90,6 @@ private ChunkedRecordAccumulator newAccumulator(int batchSize, int chunkSize, lo /* transactionManager */ null, pool); } - /** First record creates a batch backed by chunked allocation. */ @Test public void testFirstRecordCreatesChunkedBatch() throws Exception { int chunkSize = 256; @@ -109,12 +107,8 @@ public void testFirstRecordCreatesChunkedBatch() throws Exception { accum.close(); } - /** - * A small follow-up record fits in the existing batch's pre-allocated chunks — no extension - * needed. - */ @Test - public void testFollowupRecordFitsWithoutExtension() throws Exception { + public void testSmallFollowupRecordFitsWithoutExtension() throws Exception { int chunkSize = 256; ChunkedRecordAccumulator accum = newAccumulator(1024, chunkSize, 16L * chunkSize, Compression.NONE); @@ -125,6 +119,7 @@ public void testFollowupRecordFitsWithoutExtension() throws Exception { Deque dq = batchesFor(accum, tp1); assertEquals(1, dq.size(), "Both records should land in the same batch"); + assertNotNull(dq.peekFirst()); assertEquals(2, dq.peekFirst().recordCount); accum.close(); } @@ -144,6 +139,7 @@ public void testMidBatchExtensionGrowsExistingBatch() throws Exception { Deque dq = batchesFor(accum, tp1); assertEquals(1, dq.size()); + assertNotNull(dq.peekFirst()); int initialRecordCount = dq.peekFirst().recordCount; assertEquals(1, initialRecordCount); @@ -153,26 +149,17 @@ public void testMidBatchExtensionGrowsExistingBatch() throws Exception { // Same batch, now with the second record. assertEquals(1, dq.size(), "Should still be the same batch — extension should not roll"); + assertNotNull(dq.peekFirst()); assertEquals(2, dq.peekFirst().recordCount, "Second record should land in the extended batch"); accum.close(); } /** - * Race between two concurrent appenders on the same chunked tail. Both threads probe - * {@code extensionBytesNeeded} under the deque lock against the same physical remaining - * capacity, drop the lock, and allocate gap-sized chunks off-lock. When the first appender - * re-takes the lock and consumes most of the original remaining, the second appender's - * allocation — sized against the pre-race remaining — is short by the bytes the first - * appender consumed. Without a re-probe in the second-lock block, the second appender - * proceeds to write past the stream's physical capacity and {@code advanceToNextChunk} - * throws {@link IllegalStateException}. - *

- * Coordination: a {@link ChunkedBufferPool} subclass blocks the FIRST {@code allocateChunks} - * call after the test arms a latch (i.e. thread A's). Thread B's allocate is the second call - * and proceeds normally — B re-takes the lock, attaches its chunks, appends its record, and - * returns. Only then is thread A released, so A attaches its chunks and tries to append - * against a tail whose remaining has already shrunk. + * Two concurrent appenders race to extend the same open batch, each sizing its extension + * against the same remaining capacity off-lock; once one attaches its chunks, the other's is + * short. Verifies the second-lock re-check makes the short appender re-check rather than write + * past the stream's capacity, which would fail. */ @Test public void testConcurrentExtensionRaceDoesNotOverflowChunkedStream() throws Exception { @@ -202,7 +189,7 @@ public List allocateChunks(int totalSize, long maxTimeToBlockMs) thr /* deliveryTimeoutMs */ 3200, metrics, "producer-metrics", time, /* transactionManager */ null, pool); try { - // Warmup: tiny first record establishes a chunked tail with a known remainingInStream. + // Warmup: tiny first record establishes an open batch with a known remainingInStream. accum.append(topic, partition1, 0L, key, new byte[1], Record.EMPTY_HEADERS, null, maxBlockTimeMs, time.milliseconds(), cluster); @@ -224,8 +211,8 @@ public List allocateChunks(int totalSize, long maxTimeToBlockMs) thr assertTrue(aHoldsChunks.await(10, TimeUnit.SECONDS), "A did not reach allocateChunks"); // B runs on this thread. Its allocateChunks is the second call — armed=false now, - // so it proceeds without blocking. B probes the same R as A (A hasn't attached yet), - // attaches its own chunks, appends its record. Now the tail's remaining has shrunk + // so it proceeds without blocking. B reads the same R as A (A hasn't attached yet), + // attaches its own chunks, appends its record. Now the open batch's remaining has shrunk // by B's actual record size, but A's still-off-lock allocation didn't know that. accum.append(topic, partition1, 0L, key, new byte[recordValueSize], Record.EMPTY_HEADERS, null, maxBlockTimeMs, time.milliseconds(), cluster); @@ -245,17 +232,9 @@ public List allocateChunks(int totalSize, long maxTimeToBlockMs) thr } /** - * Regression: when the sticky partition switches while extension chunks are held off-lock, - * those chunks — sized against the previous partition's tail batch — must be refunded - * to the pool on the retry, not carried into the next iteration and attached to a different - * partition's tail (which would over-extend the wrong batch). - *

- * Deterministic reproduction without threads: the extension path is the only one that calls - * {@code allocateChunks} non-blocking ({@code maxTimeToBlockMs == 0}), so the pool flags when an - * extension allocation has happened; a {@code partitionChanged} override then fires exactly one - * spurious switch on the very next check — which is the second-sync-block check, with the - * extension chunks held. The fix refunds those chunks before the retry; without it they are - * carried and attached instead, so no refund is observed during the append. + * When the sticky partition switches while extension chunks are held off-lock, those + * chunks (sized against the previous partition's open batch) must be refunded to the pool on the + * retry (not carried over and attached to a different partition's open batch) */ @Test public void testPartitionSwitchRefundsHeldExtensionChunks() throws Exception { @@ -299,7 +278,7 @@ protected boolean partitionChanged(String topic, TopicInfo topicInfo, } }; try { - // Warmup: tiny first record establishes a chunked tail (first-record path, blocking + // Warmup: tiny first record establishes an open batch (first-record path, blocking // allocate — does not flag extensionAllocated). accum.append(topic, partition1, 0L, key, new byte[1], Record.EMPTY_HEADERS, null, maxBlockTimeMs, time.milliseconds(), cluster); @@ -318,6 +297,7 @@ protected boolean partitionChanged(String topic, TopicInfo topicInfo, // The record still appends correctly after the switch-retry. Deque dq = batchesFor(accum, tp1); assertEquals(1, dq.size()); + assertNotNull(dq.peekFirst()); assertEquals(2, dq.peekFirst().recordCount, "second record should still land in the batch after the switch-retry"); } finally { @@ -325,14 +305,6 @@ protected boolean partitionChanged(String topic, TopicInfo topicInfo, } } - /** - * Regression guard: an inflight chunked batch returns all K chunks to the pool on the - * expiration / broker-disconnect path, not just one {@code initialCapacity} chunkSize. - * For an inflight batch the parent {@code RecordAccumulator.deallocate} calls - * {@code batch.deallocateInflightBuffer(pool)} and then throws; {@link ChunkedProducerBatch} - * overrides that hook to return all chunks (instead of donating one fresh buffer). Removing - * that override would re-introduce a (K−1)×chunkSize leak and fail this test. - */ @Test public void testInflightExpirationReturnsAllChunksToPool() throws Exception { int chunkSize = 128; @@ -377,13 +349,7 @@ private Deque batchesFor(RecordAccumulator accum, TopicPartition } /** - * Sender's drain (and any other caller of {@code batch.close()}) cascades through - * {@code MemoryRecordsBuilder.closeForRecordAppends()} → - * {@code appendStream.close()} → {@code bufferStream.close()}. The chunk data must remain - * readable through that cascade so the builder can flatten the chunks into the heap buffer it - * sends over the network — chunks are returned to the pool only at batch completion - * (deallocate), never at close. If they were released at close, the flattened buffer would be - * empty and the header write would fail (or produce an empty record set). + * Test that chunks are returned to the pool only at batch completion (deallocate), never at close. */ @Test public void testBatchCloseDoesNotDeallocateChunksPrematurely() throws Exception { @@ -417,11 +383,8 @@ public void testBatchCloseDoesNotDeallocateChunksPrematurely() throws Exception } /** - * Chunks attached grow with the batch's cumulative projected output, not per-record. With a - * chunkSize smaller than the batch's eventual content, a sequence of small records should - * extend the chunks as the running total (header + uncompressed bytes for NONE) crosses - * chunk boundaries — not allocate the first record's worth and then never extend until - * physical capacity is exhausted (which the per-record formula would do). + * As small records accumulate in a batch, the attached chunks grow with the batch's cumulative + * projected output (not per-record). */ @Test public void testExtensionTracksCumulativeBatchSize() throws Exception { @@ -455,15 +418,8 @@ public void testExtensionTracksCumulativeBatchSize() throws Exception { accum.close(); } - /** - * The cumulative sizing formula counts the batch header once. The previous per-record - * formula effectively included the batch header in each record's upper bound, which - * over-allocated as records accumulated. With the cumulative formula, chunk count for N - * small records of total uncompressed size {@code U} is {@code ceil((header + U) / chunkSize)}, - * not {@code N × ceil((header + recordSize) / chunkSize)}. - */ @Test - public void testCumulativeAccountsForHeaderOnce() throws Exception { + public void testCumulativeAccountsForBatchHeaderOnce() throws Exception { int chunkSize = 256; int batchSize = 8192; long totalMemory = 64L * chunkSize; From 860bde4364550fae304a834b2c5c1dea12320af2 Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Tue, 23 Jun 2026 13:38:24 -0400 Subject: [PATCH 09/61] fixes --- .../apache/kafka/clients/producer/KafkaProducer.java | 12 +++++++----- .../internals/ChunkedByteBufferOutputStream.java | 2 ++ 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/KafkaProducer.java b/clients/src/main/java/org/apache/kafka/clients/producer/KafkaProducer.java index 3b13014768281..85bd74a4fc61d 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/KafkaProducer.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/KafkaProducer.java @@ -462,16 +462,18 @@ public KafkaProducer(Properties properties, Serializer keySerializer, Seriali String allocationStrategy = config.getString(ProducerConfig.BUFFER_MEMORY_ALLOCATION_STRATEGY_CONFIG) .toLowerCase(Locale.ROOT); boolean incremental = ProducerConfig.BUFFER_MEMORY_ALLOCATION_STRATEGY_INCREMENTAL.equals(allocationStrategy); - if (incremental && compression.type() != CompressionType.NONE) { + // Use the chunked path only when a batch is at least one full chunk + // (batch.size >= CHUNK_SIZE). Below that, a batch can't fill even one chunk, so chunking + // would over-reserve and the producer falls back to the full strategy instead. + boolean useIncremental = incremental && batchSize >= ChunkedRecordAccumulator.CHUNK_SIZE; + // The chunked path does not support compression yet (TODO: KAFKA-20579) + if (useIncremental && compression.type() != CompressionType.NONE) { throw new ConfigException("The " + ProducerConfig.BUFFER_MEMORY_ALLOCATION_STRATEGY_INCREMENTAL + " " + ProducerConfig.BUFFER_MEMORY_ALLOCATION_STRATEGY_CONFIG + " does not support compression yet. " + ProducerConfig.COMPRESSION_TYPE_CONFIG + " must be set to none."); } - // Use the chunked path only when a batch is at least one full chunk - // (batch.size >= CHUNK_SIZE). Below that, a batch can't fill even one chunk, so chunking - // would over-reserve. - if (incremental && batchSize >= ChunkedRecordAccumulator.CHUNK_SIZE) { + if (useIncremental) { this.accumulator = new ChunkedRecordAccumulator(logContext, batchSize, compression, diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java index 97b89cd1b4416..101a5b15cabdc 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java @@ -216,6 +216,8 @@ public void position(int position) { */ @Override public int remaining() { + if (currentChunk == null) // after deallocate no chunks attached, no free capacity + return 0; int total = currentChunk.remaining(); for (int i = currentChunkIndex + 1; i < chunks.size(); i++) total += chunks.get(i).remaining(); From 179e6750601c41c5155105786a2f6dad476c0d37 Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Tue, 23 Jun 2026 16:27:11 -0400 Subject: [PATCH 10/61] checkstyle --- .../clients/producer/internals/ChunkedBufferPoolTest.java | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPoolTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPoolTest.java index 45e9aceb12cf5..5b0f6f9b4fbb4 100644 --- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPoolTest.java +++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPoolTest.java @@ -17,6 +17,7 @@ package org.apache.kafka.clients.producer.internals; import org.apache.kafka.clients.producer.BufferExhaustedException; +import org.apache.kafka.common.KafkaException; import org.apache.kafka.common.metrics.Metrics; import org.apache.kafka.common.utils.MockTime; import org.apache.kafka.test.TestUtils; @@ -24,8 +25,6 @@ import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.Test; -import org.apache.kafka.common.KafkaException; - import java.nio.ByteBuffer; import java.util.List; import java.util.concurrent.CountDownLatch; From 9b45abc201bf270bd093d417bbd0589ca0768e84 Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Tue, 23 Jun 2026 17:38:44 -0400 Subject: [PATCH 11/61] fix test --- .../org/apache/kafka/tools/other/ReplicationQuotasTestRig.java | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tools/src/test/java/org/apache/kafka/tools/other/ReplicationQuotasTestRig.java b/tools/src/test/java/org/apache/kafka/tools/other/ReplicationQuotasTestRig.java index 8e646534263a7..7830889a5313d 100644 --- a/tools/src/test/java/org/apache/kafka/tools/other/ReplicationQuotasTestRig.java +++ b/tools/src/test/java/org/apache/kafka/tools/other/ReplicationQuotasTestRig.java @@ -414,7 +414,8 @@ KafkaProducer createProducer() { Option.empty(), new ByteArraySerializer(), new ByteArraySerializer(), - false + false, + Option.empty() ); } } From 59af9477d6f8436758cf3b4dd89a6df76a29e902 Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Thu, 2 Jul 2026 10:18:25 -0400 Subject: [PATCH 12/61] config internal --- .../java/org/apache/kafka/clients/producer/ProducerConfig.java | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/ProducerConfig.java b/clients/src/main/java/org/apache/kafka/clients/producer/ProducerConfig.java index 696e9c103118f..f574e0730c01f 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/ProducerConfig.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/ProducerConfig.java @@ -413,7 +413,8 @@ public class ProducerConfig extends AbstractConfig { Importance.MEDIUM, CommonClientConfigs.CLIENT_DNS_LOOKUP_DOC) .define(BUFFER_MEMORY_CONFIG, Type.LONG, 32 * 1024 * 1024L, atLeast(0L), Importance.HIGH, BUFFER_MEMORY_DOC) - .define(BUFFER_MEMORY_ALLOCATION_STRATEGY_CONFIG, + // Internal until the incremental strategy is fully implemented + .defineInternal(BUFFER_MEMORY_ALLOCATION_STRATEGY_CONFIG, Type.STRING, BUFFER_MEMORY_ALLOCATION_STRATEGY_FULL, ConfigDef.CaseInsensitiveValidString From 9e791c8bf0a8e47c574bb8b69ef0faa160e149c2 Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Thu, 2 Jul 2026 14:48:00 -0400 Subject: [PATCH 13/61] improvements for tryAppend --- .../internals/ChunkedProducerBatch.java | 28 +++++-- .../internals/ChunkedRecordAccumulator.java | 61 ++++++++------- .../producer/internals/RecordAccumulator.java | 42 +++++++--- .../ChunkedRecordAccumulatorTest.java | 78 +++++++++++++++++++ 4 files changed, 165 insertions(+), 44 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedProducerBatch.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedProducerBatch.java index 05c290737fe24..8a1d0032b90e3 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedProducerBatch.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedProducerBatch.java @@ -16,6 +16,7 @@ */ package org.apache.kafka.clients.producer.internals; +import org.apache.kafka.clients.producer.Callback; import org.apache.kafka.common.TopicPartition; import org.apache.kafka.common.header.Header; import org.apache.kafka.common.record.internal.MemoryRecordsBuilder; @@ -45,15 +46,13 @@ public ChunkedProducerBatch(TopicPartition tp, MemoryRecordsBuilder recordsBuild /** * Bytes of chunk capacity this batch needs before {@code tryAppend} could accept the given - * record. Returns 0 when no extension is needed: the batch is empty (first record), it is at - * its batch-size limit, or the attached chunk capacity already has room. Positive when the - * record is within the batch-size limit but the attached chunks lack capacity — the accumulator - * then allocates exactly the missing bytes (rounded up to whole chunks) and attaches them via - * {@link #addBuffers} before retrying. + * record. Returns 0 when no extension is needed: the batch is at its batch-size limit, or the + * attached chunk capacity already has room (always the case for an empty batch, whose stream + * is pre-sized for the first record). Positive when the record is within the batch-size limit + * but the attached chunks lack capacity — the accumulator then allocates exactly the missing + * bytes (rounded up to whole chunks) and attaches them via {@link #addBuffers} before retrying. */ int extensionBytesNeeded(long timestamp, byte[] key, byte[] value, Header[] headers) { - if (recordCount == 0) - return 0; if (!recordsBuilder.hasRoomFor(timestamp, key, value, headers)) return 0; // Size against the batch's projected total output after this record (header counted once, @@ -65,6 +64,21 @@ int extensionBytesNeeded(long timestamp, byte[] key, byte[] value, Header[] head return Math.max(0, target - totalAttachedCapacity); } + /** + * Appends the record after asserting there's capacity for it. + * At this point it's expected that the allocated chunks have + * capacity for the record because the accumulator never routes an + * append here without ensuring chunk capacity first: the first-record path pre-sizes the + * stream for the record and the mid-batch path attaches extension chunks before retrying. + */ + @Override + public FutureRecordMetadata tryAppend(long timestamp, byte[] key, byte[] value, Header[] headers, Callback callback, long now) { + assert extensionBytesNeeded(timestamp, key, value, headers) == 0 : + "Unexpected append to a chunked batch whose chunks lack capacity for the record; " + + "the accumulator should have attached extension chunks first"; + return super.tryAppend(timestamp, key, value, headers, callback, now); + } + /** * Attach pre-allocated chunks to this batch's stream so the next {@code tryAppend} can * spill into them. Ownership of the chunks transfers to the stream. diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java index 297dac0dca91f..75cdd8ee98c74 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java @@ -118,7 +118,6 @@ public RecordAppendResult append(String topic, appendsInProgress.incrementAndGet(); ChunkedByteBufferOutputStream bufferStream = null; List extensionChunks = null; - int extensionBytes; if (headers == null) headers = Record.EMPTY_HEADERS; try { while (true) { @@ -134,31 +133,35 @@ public RecordAppendResult append(String topic, setPartition(callbacks, effectivePartition); Deque dq = topicInfo.batches.computeIfAbsent(effectivePartition, k -> new ArrayDeque<>()); + RecordAppendResult appendResult; synchronized (dq) { if (partitionChanged(topic, topicInfo, partitionInfo, dq, nowMs, cluster)) continue; - // The tryAppend override checks the open batch (dq.peekLast()) for chunk - // capacity: a needsBufferExtension result means it is within its batch-size limit - // but its chunks lack capacity for this record — fall through to allocate the gap - // outside the deque lock. A null result means there is no open batch (it is full - // or absent) — fall through to the first-record (new batch) path. - RecordAppendResult appendResult = tryAppend(timestamp, key, value, headers, callbacks, dq, nowMs); - if (appendResult != null && !appendResult.needsBufferExtension) { + // The tryAppend checks the open batch (dq.peekLast()) for chunk capacity: + // a needsBufferExtension result means it is within its batch-size limit + // but its chunks lack capacity for this record, so it will allocate the gap + // outside the deque lock. A needsNewBatch result means there is no open batch + // (full or absent), so it will fall through to the first-record (new batch) path. + appendResult = tryAppend(timestamp, key, value, headers, callbacks, dq, nowMs); + if (appendResult.appended()) { boolean enableSwitch = allBatchesFull(dq); topicInfo.builtInPartitioner.updatePartitionInfo(partitionInfo, appendResult.appendedBytes, cluster, enableSwitch); return appendResult; } - extensionBytes = appendResult == null ? 0 : appendResult.extensionBytesNeeded; } - if (extensionBytes > 0) { - // Mid-batch extension: non-blocking only. The thread already holds the open - // batch's chunks, so blocking here could deadlock with the Sender (which frees - // pool memory by completing batches). On exhaustion, close the batch (making it - // drainable) and fall through to the blocking first-record path next iteration. + if (appendResult.needsBufferExtension) { + // Mid-batch extension: the open batch can still take this record so grow it in + // place. The acquire is non-blocking to fail fast when the pool is exhausted: + // close the batch and let the record block once on the new-batch path. + // A blocking call would lead to the same outcome, but would block once here + // and still need a second blocking call to start a new batch anyways + // (a first blocking call here would make all open batches drainable, including this one, + // so most probably our batch would be gone/drained by the time memory is returned to the pool, + // and we would need a new batch for our record anyways). try { - extensionChunks = chunkedFree.allocateChunks(extensionBytes, 0L); + extensionChunks = chunkedFree.allocateChunks(appendResult.extensionBytesNeeded, 0L); } catch (BufferExhaustedException e) { log.trace("Pool exhausted while extending batch for topic {} partition {}; closing existing batch", topic, effectivePartition); @@ -168,13 +171,16 @@ public RecordAppendResult append(String topic, last.closeForRecordAppends(); } } + // Continue to the next iteration that should block to start a new batch (needsNewBatch), + // given that this one has been closed for appends. continue; } nowMs = time.milliseconds(); - } else if (extensionBytes == 0 && bufferStream == null) { - // First-record path: block on the pool for enough chunks to fit this record, - // sized with the same cumulative estimator used mid-batch (header + record bytes - // for NONE, ratio-adjusted when compressed) so the two stay consistent. + } else if (appendResult.needsNewBatch && bufferStream == null) { + // The open batch is done (e.g., full, closed) so start a new one. + // Block on the pool for enough chunks to fit this record, sized with the + // same cumulative estimator used mid-batch (header + record bytes for NONE, + // ratio-adjusted when compressed) so the two stay consistent. int recordUncompressed = AbstractRecords.recordSizeUpperBound( RecordBatch.CURRENT_MAGIC_VALUE, compression.type(), key, value, headers); int size = MemoryRecordsBuilder.estimatedBytesWritten( @@ -212,15 +218,15 @@ public RecordAppendResult append(String topic, ((ChunkedProducerBatch) last).addBuffers(extensionChunks); extensionChunks = null; RecordAppendResult retryResult = tryAppend(timestamp, key, value, headers, callbacks, dq, nowMs); - if (retryResult != null && !retryResult.needsBufferExtension) { + if (retryResult.appended()) { boolean enableSwitch = allBatchesFull(dq); topicInfo.builtInPartitioner.updatePartitionInfo(partitionInfo, retryResult.appendedBytes, cluster, enableSwitch); return retryResult; } - // needsBufferExtension: a concurrent appender consumed capacity our - // extension was sized against — loop to re-check. null: batch became - // full — loop into new-batch creation. Terminates because writeLimit is - // fixed: once full, the check stops requesting extension. + // Still not appended: concurrent appenders filled the batch, + // so the extension we attached is no longer enough. + // Loop so the next iteration routes the record + // right: needsBufferExtension with a fresh gap, or needsNewBatch continue; } // The open batch is gone, closed, or non-chunked (e.g., a split batch). Return chunks to pool. @@ -230,13 +236,13 @@ public RecordAppendResult append(String topic, continue; } - // First-record path: extensionChunks == null here implies extensionBytes == 0, + // needsNewBatch path: extensionChunks == null here implies needsNewBatch, // so bufferStream was allocated (this iteration or carried from a prior one). assert bufferStream != null; int firstRecordSize = AbstractRecords.estimateSizeInBytesUpperBound( RecordBatch.CURRENT_MAGIC_VALUE, compression.type(), key, value, headers); final ChunkedByteBufferOutputStream batchStream = bufferStream; - RecordAppendResult appendResult = appendNewBatch(topic, effectivePartition, dq, timestamp, key, value, headers, callbacks, + appendResult = appendNewBatch(topic, effectivePartition, dq, timestamp, key, value, headers, callbacks, () -> chunkedRecordsBuilder(batchStream, firstRecordSize), nowMs); if (appendResult.needsBufferExtension) { // A concurrent appender created an open batch we should extend rather @@ -271,7 +277,8 @@ public RecordAppendResult append(String topic, * If the open batch is within its batch-size limit but its chunked stream lacks chunk * capacity, returns {@link RecordAppendResult#needsExtension(int)} without * attempting the append; the caller allocates chunks outside the deque lock, attaches - * them, and retries. Otherwise defers to the parent. + * them, and retries. Otherwise defers to the parent, which appends or returns + * {@link RecordAppendResult#NEEDS_NEW_BATCH}. */ @Override protected RecordAppendResult tryAppend(long timestamp, byte[] key, byte[] value, Header[] headers, diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/RecordAccumulator.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/RecordAccumulator.java index c43952e96c5a3..1551d4000c97e 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/RecordAccumulator.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/RecordAccumulator.java @@ -324,7 +324,7 @@ public RecordAppendResult append(String topic, continue; RecordAppendResult appendResult = tryAppend(timestamp, key, value, headers, callbacks, dq, nowMs); - if (appendResult != null) { + if (appendResult.appended()) { // If queue has incomplete batches we disable switch (see comments in updatePartitionInfo). boolean enableSwitch = allBatchesFull(dq); topicInfo.builtInPartitioner.updatePartitionInfo(partitionInfo, appendResult.appendedBytes, cluster, enableSwitch); @@ -398,7 +398,7 @@ protected RecordAppendResult appendNewBatch(String topic, assert partition != RecordMetadata.UNKNOWN_PARTITION; RecordAppendResult appendResult = tryAppend(timestamp, key, value, headers, callbacks, dq, nowMs); - if (appendResult != null) { + if (!appendResult.needsNewBatch) { // Propagate without creating a new batch: either another thread already made us a batch // (success), or — incremental strategy — a concurrent appender created an extendable open batch // (needsBufferExtension), so the caller releases its pre-allocated buffer and retries via @@ -437,10 +437,10 @@ protected boolean allBatchesFull(Deque deque) { /** * Try to append to a ProducerBatch. * - * If it is full, we return null and a new batch is created. We also close the batch for record appends to free up - * resources like compression buffers. The batch will be fully closed (ie. the record batch headers will be written - * and memory records built) in one of the following cases (whichever comes first): right before send, - * if it is expired, or when the producer is closed. + * If it is full (or absent), we return {@link RecordAppendResult#NEEDS_NEW_BATCH} and a new batch is created. + * We also close the batch for record appends to free up resources like compression buffers. The batch will be + * fully closed (ie. the record batch headers will be written and memory records built) in one of the following + * cases (whichever comes first): right before send, if it is expired, or when the producer is closed. */ protected RecordAppendResult tryAppend(long timestamp, byte[] key, byte[] value, Header[] headers, Callback callback, Deque deque, long nowMs) { @@ -457,7 +457,7 @@ protected RecordAppendResult tryAppend(long timestamp, byte[] key, byte[] value, return new RecordAppendResult(future, deque.size() > 1 || last.isFull(), false, appendedBytes); } } - return null; + return RecordAppendResult.NEEDS_NEW_BATCH; } private boolean isMuted(TopicPartition tp) { @@ -1266,7 +1266,10 @@ public PartitionerConfig() { } /* - * Metadata about a record just appended to the record accumulator + * Result of an attempt to append a record to the accumulator. Exactly one of three states: + * the record was appended ({@link RecordAppendResult#appended()}, {@code future} is set), the + * open batch needs more chunk capacity first ({@code needsBufferExtension}), or a new batch + * must be created for the record ({@code needsNewBatch}). */ public static final class RecordAppendResult { public final FutureRecordMetadata future; @@ -1280,13 +1283,22 @@ public static final class RecordAppendResult { */ public final boolean needsBufferExtension; public final int extensionBytesNeeded; + /** + * Signal that the record was not appended because there is no open batch that can take it + * (the batch is full, closed, or absent). The caller must create a new batch for the + * record. When {@code true}, {@code future} is null. + */ + public final boolean needsNewBatch; public final int appendedBytes; + /** The shared signal-only result for {@link #needsNewBatch}; carries no per-append state. */ + static final RecordAppendResult NEEDS_NEW_BATCH = new RecordAppendResult(null, false, false, false, 0, true, 0); + public RecordAppendResult(FutureRecordMetadata future, boolean batchIsFull, boolean newBatchCreated, int appendedBytes) { - this(future, batchIsFull, newBatchCreated, false, 0, appendedBytes); + this(future, batchIsFull, newBatchCreated, false, 0, false, appendedBytes); } private RecordAppendResult(FutureRecordMetadata future, @@ -1294,18 +1306,28 @@ private RecordAppendResult(FutureRecordMetadata future, boolean newBatchCreated, boolean needsBufferExtension, int extensionBytesNeeded, + boolean needsNewBatch, int appendedBytes) { this.future = future; this.batchIsFull = batchIsFull; this.newBatchCreated = newBatchCreated; this.needsBufferExtension = needsBufferExtension; this.extensionBytesNeeded = extensionBytesNeeded; + this.needsNewBatch = needsNewBatch; this.appendedBytes = appendedBytes; } /** A signal-only result indicating the caller must allocate more chunk capacity. */ static RecordAppendResult needsExtension(int extensionBytesNeeded) { - return new RecordAppendResult(null, false, false, true, extensionBytesNeeded, 0); + return new RecordAppendResult(null, false, false, true, extensionBytesNeeded, false, 0); + } + + /** + * @return {@code true} if the record was appended to the open batch ({@link #future} is then non-null), + * {@code false} if it wasn't ({@link #needsBufferExtension} and {@link #needsNewBatch} results). + */ + public boolean appended() { + return !needsBufferExtension && !needsNewBatch; } } diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java index d658e201083fe..455aa4222b47f 100644 --- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java +++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java @@ -17,6 +17,7 @@ package org.apache.kafka.clients.producer.internals; import org.apache.kafka.clients.MetadataSnapshot; +import org.apache.kafka.clients.producer.BufferExhaustedException; import org.apache.kafka.common.Cluster; import org.apache.kafka.common.Node; import org.apache.kafka.common.TopicPartition; @@ -24,6 +25,7 @@ import org.apache.kafka.common.metrics.Metrics; import org.apache.kafka.common.protocol.Errors; import org.apache.kafka.common.record.internal.MemoryRecords; +import org.apache.kafka.common.record.internal.MemoryRecordsBuilder; import org.apache.kafka.common.record.internal.Record; import org.apache.kafka.common.record.internal.RecordBatch; import org.apache.kafka.common.requests.MetadataResponse.PartitionMetadata; @@ -348,6 +350,82 @@ private Deque batchesFor(RecordAccumulator accum, TopicPartition return accum.getDeque(tp); } + /** + * When the pool is exhausted during a mid-batch extension, the append must not busy-loop + * retrying the non-blocking acquire: after a single failed extension acquire + * (maxTimeToBlockMs = 0), the open batch is closed and the very next pool call is the + * blocking new-batch acquire (maxTimeToBlockMs > 0), where the record lands in a new batch. + */ + @Test + public void testExhaustedExtensionFallsBackToBlockingNewBatchPath() throws Exception { + int chunkSize = 256; + List allocTimeouts = new ArrayList<>(); + AtomicInteger closeForAppendsCalls = new AtomicInteger(); + List closeCallsAtAlloc = new ArrayList<>(); + + ChunkedBufferPool pool = new ChunkedBufferPool(16L * chunkSize, chunkSize, metrics, time, "producer-metrics") { + @Override + public List allocateChunks(int totalSize, long maxTimeToBlockMs) throws InterruptedException { + allocTimeouts.add(maxTimeToBlockMs); + closeCallsAtAlloc.add(closeForAppendsCalls.get()); + // Simulate an exhausted pool for the non-blocking extension acquire only. + if (maxTimeToBlockMs == 0L) + 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) { + // Count closeForRecordAppends calls on the batches this accumulator creates. + return new ChunkedProducerBatch(tp, recordsBuilder, nowMs) { + @Override + public void closeForRecordAppends() { + closeForAppendsCalls.incrementAndGet(); + super.closeForRecordAppends(); + } + }; + } + }; + try { + // First record establishes an open batch (blocking first-record acquire). + accum.append(topic, partition1, 0L, key, new byte[100], Record.EMPTY_HEADERS, null, + maxBlockTimeMs, time.milliseconds(), cluster); + + // Second record overflows the batch's chunk so it needs extension; the injected + // exhaustion fails the non-blocking acquire, the batch is closed, and the record + // must go straight to the blocking new-batch path. + RecordAccumulator.RecordAppendResult result = accum.append(topic, partition1, 0L, key, + new byte[100], Record.EMPTY_HEADERS, null, maxBlockTimeMs, time.milliseconds(), cluster); + + // Validate the expected call sequence: blocking (first batch), non-blocking (failed extension), blocking (new batch). + assertEquals(List.of(maxBlockTimeMs, 0L, maxBlockTimeMs), allocTimeouts, + "expected a single failed extension acquire followed directly by the blocking new-batch acquire"); + assertEquals(0, closeCallsAtAlloc.get(0), "no close before the first-record acquire"); + assertEquals(0, closeCallsAtAlloc.get(1), "no close before the extension acquire"); + assertTrue(closeCallsAtAlloc.get(2) >= 1, + "the failed extension must close the batch before the blocking new-batch acquire"); + assertTrue(result.newBatchCreated, "record must land in a new batch"); + + Deque dq = batchesFor(accum, tp1); + assertEquals(2, dq.size(), "closed batch + new batch expected"); + assertNotNull(dq.peekFirst()); + assertNotNull(dq.peekLast()); + // The original batch is far below its writeLimit, so isFull() can only be true via + // the closed append stream — i.e., it was closed for appends on the failed extension. + assertTrue(dq.peekFirst().isFull(), "original batch must be closed for appends"); + assertNotNull(dq.peekFirst()); + assertEquals(1, dq.peekFirst().recordCount); + assertNotNull(dq.peekLast()); + assertEquals(1, dq.peekLast().recordCount); + } finally { + accum.close(); + } + } + /** * Test that chunks are returned to the pool only at batch completion (deallocate), never at close. */ From 35c5940190d8c7ea36eb24a0593613d0c0c41d88 Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Thu, 2 Jul 2026 16:10:29 -0400 Subject: [PATCH 14/61] release unused & test --- .../ChunkedByteBufferOutputStream.java | 30 ++++++++++----- .../producer/internals/ProducerBatch.java | 2 +- .../ChunkedByteBufferOutputStreamTest.java | 38 +++++++++++++++++++ 3 files changed, 60 insertions(+), 10 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java index 101a5b15cabdc..fc706d140a738 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java @@ -146,9 +146,6 @@ public ByteBuffer buffer() { if (flattenedBuffer != null && !dirty) { return flattenedBuffer; } - // TODO: KAFKA-20687. This flatten runs at batch close, when the chunk set is final. - // Today all chunks (used and unused) are returned to the pool only when the batch - // completes (via deallocate(pool)). Consider releasing the fully-unused chunks early, here. int totalSize = 0; for (ByteBuffer chunk : chunks) { totalSize += chunk.position(); @@ -162,15 +159,30 @@ public ByteBuffer buffer() { chunk.position(chunkPos); } dirty = false; - // The bytes are now copied into flattenedBuffer, but we intentionally do not release the - // chunks until batch completion, so the in-flight data stays reserved against the - // buffer.memory budget (the pool's available memory reflects it), consistent with the - // "full" strategy. Releasing here would return that memory to the pool while the bytes are - // still in flight in the heap copy, letting the pool admit more than buffer.memory intends. - // This flattening is an initial approach and will be removed with KAFKA-20580. + releaseUnusedChunks(); return flattenedBuffer; } + /** + * Release the fully-unused chunks (nothing written) back to the pool, now rather than at batch + * completion. Called at batch close, when the chunk set is final. + *

+ * The data-bearing chunks are intentionally kept until batch completion: they are what keeps + * the batch's in-flight data reserved against the buffer.memory budget, so releasing them at + * close would let the pool admit more data than buffer.memory intends while the batch is still + * in flight. + */ + private void releaseUnusedChunks() { + List unused = chunks.subList(currentChunkIndex + 1, chunks.size()); + if (pool != null) { + for (ByteBuffer chunk : unused) + pool.deallocate(chunk); + } + // Remove the released chunks from `chunks`, so they are + // not deallocated again on batch completion. + unused.clear(); + } + /** * Total bytes written across all chunks. */ diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ProducerBatch.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ProducerBatch.java index 6adac7a24424d..be8b54dafe2b0 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ProducerBatch.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ProducerBatch.java @@ -142,7 +142,7 @@ boolean hasLeaderChangedForTheOngoingRetry() { /** * Return this batch's buffer memory to the pool. The default single-buffer batch returns * the pooled buffer at its initial capacity; {@link ChunkedProducerBatch} overrides this - * to return all of its chunks. + * to return its remaining chunks (fully-unused chunks are released earlier, at close). */ protected void deallocateBuffer(BufferPool pool) { pool.deallocate(buffer(), initialCapacity()); diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java index 145876730e977..efa7fb6923322 100644 --- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java +++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java @@ -171,4 +171,42 @@ public void testDeallocateReturnsAllChunks() throws Exception { } assertEquals(total, p.availableMemory()); } + + /** + * Closing the batch (which invokes {@code buffer()}) returns fully-unused chunks to the pool + * right away; data-bearing chunks stay reserved until + * {@link ChunkedByteBufferOutputStream#deallocate()}, which must not return the + * already-released chunks a second time. + */ + @Test + public void testUnusedChunksReleasedAtClose() throws Exception { + int chunkSize = 8; + long total = 64; + ChunkedBufferPool p = pool(total, chunkSize); + try (ChunkedByteBufferOutputStream stream = new ChunkedByteBufferOutputStream(chunks(p, chunkSize, 3), chunkSize, p)) { + // Write into the first chunk only; chunks 2 and 3 stay untouched. + byte[] payload = new byte[]{1, 2, 3}; + stream.write(payload, 0, payload.length); + assertEquals(total - 3L * chunkSize, p.availableMemory()); + + ByteBuffer built = stream.buffer(); + assertEquals(total - chunkSize, p.availableMemory(), + "the two unused chunks should return to the pool at close"); + + // The written content is intact. + built.flip(); + byte[] out = new byte[built.remaining()]; + built.get(out); + assertArrayEquals(payload, out); + + // A repeated buffer call must not release anything further. + stream.buffer(); + assertEquals(total - chunkSize, p.availableMemory()); + + // Completion-time deallocate must return only the remaining data-bearing chunk. + stream.deallocate(); + assertEquals(total, p.availableMemory(), + "pool must be exactly restored; released chunks must not be returned twice on completion"); + } + } } \ No newline at end of file From 63de56b58752d18d6aab82158cadb490a54b683d Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Thu, 2 Jul 2026 17:43:21 -0400 Subject: [PATCH 15/61] fix to loop in ensureChunkCapacity --- .../producer/internals/ChunkedByteBufferOutputStream.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java index fc706d140a738..4f47d6723fc63 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java @@ -116,7 +116,7 @@ public void write(ByteBuffer sourceBuffer) { } private void ensureChunkCapacity(int needed) { - if (currentChunk.remaining() < needed) { + while (currentChunk.remaining() < needed) { advanceToNextChunk(); } } From a2ee30f3ebaf369ca6a7bfc6c82ca78933eb01bd Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Thu, 2 Jul 2026 17:56:47 -0400 Subject: [PATCH 16/61] fix iteration up to currentChunk --- .../internals/ChunkedByteBufferOutputStream.java | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java index 4f47d6723fc63..cf5a923f12adc 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java @@ -146,12 +146,15 @@ public ByteBuffer buffer() { if (flattenedBuffer != null && !dirty) { return flattenedBuffer; } + // Written bytes only live in chunks up to currentChunk; later chunks are untouched. + int lastDataChunk = Math.min(currentChunkIndex, chunks.size() - 1); int totalSize = 0; - for (ByteBuffer chunk : chunks) { - totalSize += chunk.position(); + for (int i = 0; i <= lastDataChunk; i++) { + totalSize += chunks.get(i).position(); } flattenedBuffer = ByteBuffer.allocate(totalSize); - for (ByteBuffer chunk : chunks) { + for (int i = 0; i <= lastDataChunk; i++) { + ByteBuffer chunk = chunks.get(i); int chunkPos = chunk.position(); chunk.flip(); flattenedBuffer.put(chunk); @@ -188,9 +191,11 @@ private void releaseUnusedChunks() { */ @Override public int position() { + // Written bytes only live in chunks up to currentChunk; later chunks are untouched. + int lastDataChunk = Math.min(currentChunkIndex, chunks.size() - 1); int total = 0; - for (ByteBuffer chunk : chunks) { - total += chunk.position(); + for (int i = 0; i <= lastDataChunk; i++) { + total += chunks.get(i).position(); } return total; } From 3c896c69c62cfcf18707f2910fe1d8109e20b257 Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Thu, 2 Jul 2026 18:25:46 -0400 Subject: [PATCH 17/61] remove unneeded freeup call --- .../kafka/clients/producer/internals/ChunkedBufferPool.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java index f2c8ba35f1568..e660beb2775cc 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java @@ -82,8 +82,8 @@ public List allocateChunks(int totalSize, long maxTimeToBlockMs) thr pooled.add(free.pollFirst()); long remainingBytes = memoryRequired - (long) pooled.size() * chunkSize; if (remainingBytes > 0) { - // remainingBytes <= memoryRequired <= totalMemory (validated above), so the int cast is safe. - freeUp((int) remainingBytes); + // remainingBytes > 0 means the free list was fully drained into `pooled`, so the + // remainder comes entirely from non-pooled memory (sufficient per the check above). this.nonPooledAvailableMemory -= remainingBytes; } } else { From 20e52d392ae66a044128a81988b107bc3ee91147 Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Fri, 3 Jul 2026 11:17:38 -0400 Subject: [PATCH 18/61] tests for concurrent reservation --- .../internals/ChunkedProducerBatch.java | 4 + .../ChunkedRecordAccumulatorTest.java | 139 +++++++++++------- 2 files changed, 91 insertions(+), 52 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedProducerBatch.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedProducerBatch.java index 8a1d0032b90e3..4691efbf9e7fa 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedProducerBatch.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedProducerBatch.java @@ -82,6 +82,10 @@ assert extensionBytesNeeded(timestamp, key, value, headers) == 0 : /** * Attach pre-allocated chunks to this batch's stream so the next {@code tryAppend} can * spill into them. Ownership of the chunks transfers to the stream. + *

+ * Concurrent appenders may over-attach (each sizes its own gap against the same remaining + * capacity), attaching more than the batch-size limit can use; chunks left fully unused are + * returned to the pool when the batch closes. */ void addBuffers(List chunks) { stream().addBuffers(chunks); diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java index 455aa4222b47f..58363f4012f33 100644 --- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java +++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java @@ -42,8 +42,6 @@ import java.util.List; import java.util.Map; import java.util.Optional; -import java.util.concurrent.CountDownLatch; -import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; @@ -51,9 +49,7 @@ import java.util.stream.Stream; import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotNull; -import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -158,77 +154,116 @@ public void testMidBatchExtensionGrowsExistingBatch() throws Exception { } /** - * Two concurrent appenders race to extend the same open batch, each sizing its extension - * against the same remaining capacity off-lock; once one attaches its chunks, the other's is - * short. Verifies the second-lock re-check makes the short appender re-check rather than write - * past the stream's capacity, which would fail. + * Pool that adds one append right after a chunk allocation returns, mocking a concurrent + * appender racing the same batch. */ - @Test - public void testConcurrentExtensionRaceDoesNotOverflowChunkedStream() throws Exception { - final int chunkSize = 256; - final int recordValueSize = 350; // sized so gap ≈ chunkSize → minimal rounding cushion - - final AtomicBoolean armed = new AtomicBoolean(false); - final CountDownLatch aHoldsChunks = new CountDownLatch(1); - final CountDownLatch bDoneAttaching = new CountDownLatch(1); - - ChunkedBufferPool pool = new ChunkedBufferPool(64L * chunkSize, chunkSize, metrics, time, "producer-metrics") { + private ChunkedBufferPool poolMockingConcurrentChunkAllocation(int chunkSize, long totalMemory, + AtomicReference injectAppendOnce, + byte[] injectedValue) { + return new ChunkedBufferPool(totalMemory, chunkSize, metrics, time, "producer-metrics") { @Override public List allocateChunks(int totalSize, long maxTimeToBlockMs) throws InterruptedException { - boolean blockThisCall = armed.compareAndSet(true, false); List chunks = super.allocateChunks(totalSize, maxTimeToBlockMs); - if (blockThisCall) { - aHoldsChunks.countDown(); - if (!bDoneAttaching.await(10, TimeUnit.SECONDS)) { - throw new InterruptedException("test timeout waiting for B"); - } + ChunkedRecordAccumulator toInject = injectAppendOnce.getAndSet(null); + if (toInject != null) { + toInject.append(topic, partition1, 0L, key, injectedValue, Record.EMPTY_HEADERS, null, + maxBlockTimeMs, time.milliseconds(), cluster); } return chunks; } }; + } + + /** + * Two concurrent appenders race to extend the same open batch, each sizing its extension + * against the same remaining capacity off-lock. Whichever attaches and appends first consumes + * that capacity, leaving the other appender's extension too small for its record. Verifies + * that the post-attach {@code tryAppend} — which attempts the write based on the batch's + * actual capacity — makes that appender extend again and land its record in the same batch, + * rather than write past the stream's chunk capacity (which would surface here as an + * {@link IllegalStateException} thrown by the append). + */ + @Test + public void testConcurrentExtensionRaceLoserExtendsAgain() throws Exception { + final int chunkSize = 256; + final byte[] value = new byte[350]; // needs 2 chunks + final AtomicReference injectAppendOnce = new AtomicReference<>(); + + ChunkedBufferPool pool = poolMockingConcurrentChunkAllocation(chunkSize, 64L * chunkSize, injectAppendOnce, value); ChunkedRecordAccumulator accum = new ChunkedRecordAccumulator(logContext, 8192, Compression.NONE, /* lingerMs */ 0, /* retryBackoffMs */ 0L, /* retryBackoffMaxMs */ 0L, /* deliveryTimeoutMs */ 3200, metrics, "producer-metrics", time, /* transactionManager */ null, pool); try { - // Warmup: tiny first record establishes an open batch with a known remainingInStream. + // Tiny first record opens the batch, both racing records extend. accum.append(topic, partition1, 0L, key, new byte[1], Record.EMPTY_HEADERS, null, maxBlockTimeMs, time.milliseconds(), cluster); - // Arm the race: the next allocateChunks call (A's) will block until B has appended. - armed.set(true); + // This append sizes its gap, then the injected append wins the race while the gap + // chunks are held off-lock, shrinking the remaining capacity the gap was sized + // against. Even with the over-allocation, this append shouldn't write past the stream's + // chunk capacity. + injectAppendOnce.set(accum); + accum.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null, + maxBlockTimeMs, time.milliseconds(), cluster); - AtomicReference aError = new AtomicReference<>(); - Thread tA = new Thread(() -> { - try { - accum.append(topic, partition1, 0L, key, new byte[recordValueSize], - Record.EMPTY_HEADERS, null, maxBlockTimeMs, time.milliseconds(), cluster); - } catch (Throwable t) { - aError.set(t); - } - }, "test-appender-A"); - tA.start(); + Deque dq = batchesFor(accum, tp1); + assertEquals(1, dq.size(), "Expecting a single batch"); + assertNotNull(dq.peekFirst()); + assertEquals(3, dq.peekFirst().recordCount, "Expecting 3 records in the batch: the initial one plus the 2 racing ones."); + } finally { + accum.close(); + } + } - // Wait for A to be parked inside allocateChunks holding its (pre-attach) chunks. - assertTrue(aHoldsChunks.await(10, TimeUnit.SECONDS), "A did not reach allocateChunks"); + /** + * Concurrent extensions can over-reserve: each appender sizes its own gap against the same + * remaining capacity, so a batch can end up with more chunk capacity than its batch-size limit + * admits. The limit must still be enforced on append: a record that no longer fits after the + * race winner's append rolls to a new batch even though the over-reserved chunks could + * physically hold it, and the unused surplus returns to the pool when the batch closes. + */ + @Test + public void testConcurrentExtensionRaceLoserStartsNewBatch() throws Exception { + final int chunkSize = 256; + final int batchSize = 1024; + // Sized so each record fits the batch-size limit individually, but not both together. + final byte[] value = new byte[500]; + final AtomicReference injectAppendOnce = new AtomicReference<>(); - // B runs on this thread. Its allocateChunks is the second call — armed=false now, - // so it proceeds without blocking. B reads the same R as A (A hasn't attached yet), - // attaches its own chunks, appends its record. Now the open batch's remaining has shrunk - // by B's actual record size, but A's still-off-lock allocation didn't know that. - accum.append(topic, partition1, 0L, key, new byte[recordValueSize], Record.EMPTY_HEADERS, null, + ChunkedBufferPool pool = poolMockingConcurrentChunkAllocation(chunkSize, 64L * chunkSize, injectAppendOnce, value); + ChunkedRecordAccumulator accum = new ChunkedRecordAccumulator(logContext, batchSize, Compression.NONE, + /* lingerMs */ 0, /* retryBackoffMs */ 0L, /* retryBackoffMaxMs */ 0L, + /* deliveryTimeoutMs */ 3200, metrics, "producer-metrics", time, + /* transactionManager */ null, pool); + try { + // Open the batch with a tiny record; both racing records will extend it. + accum.append(topic, partition1, 0L, key, new byte[1], Record.EMPTY_HEADERS, null, maxBlockTimeMs, time.milliseconds(), cluster); - // Release A. Without the fix in the second-lock block, A's tryAppendToExisting - // will overflow the chunked stream and throw IllegalStateException. - bDoneAttaching.countDown(); - tA.join(10_000); - assertFalse(tA.isAlive(), "Thread A did not complete in time"); + // This append sizes its gap, then the injected append wins the race while the gap + // chunks are held off-lock. The retry finds the batch over its batch-size limit, so + // the record must roll to a new batch despite the attached (over-reserved) capacity. + injectAppendOnce.set(accum); + accum.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null, + maxBlockTimeMs, time.milliseconds(), cluster); - assertNull(aError.get(), - "Thread A failed: " + (aError.get() == null ? "" : aError.get().toString())); + Deque dq = batchesFor(accum, tp1); + assertEquals(2, dq.size(), "The losing record must roll to a new batch, not exceed batch.size"); + ProducerBatch first = dq.peekFirst(); + ProducerBatch second = dq.peekLast(); + assertNotNull(first); + assertNotNull(second); + assertEquals(2, first.recordCount, "First batch should have the initial record plus the race winner"); + assertEquals(1, second.recordCount, "Second batch should have the loser record"); + + // The loser's extension chunks were attached but never used (the over-reservation); + // closing the first batch returns that unused surplus to the pool. + long beforeClose = pool.availableMemory(); + first.close(); + assertTrue(pool.availableMemory() > beforeClose, + "The over-reserved unused chunks should return to the pool at close"); } finally { - // Drain any state regardless of outcome. accum.close(); } } From 8d2f5e2d5fd7cdcb1e9d79150ea025ef1f572172 Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Fri, 3 Jul 2026 11:24:05 -0400 Subject: [PATCH 19/61] fix for int max value cap --- .../kafka/clients/producer/internals/ChunkedBufferPool.java | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java index e660beb2775cc..0bd540b09d963 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java @@ -132,7 +132,9 @@ public List allocateChunks(int totalSize, long maxTimeToBlockMs) thr } long stillNeeded = memoryRequired - (long) pooled.size() * chunkSize - accumulated; if (stillNeeded > 0) { - freeUp((int) stillNeeded); + // stillNeeded can exceed Integer.MAX_VALUE (memoryRequired rounds totalSize + // up to whole chunks) + freeUp((int) Math.min(stillNeeded, Integer.MAX_VALUE)); long got = Math.min(stillNeeded, this.nonPooledAvailableMemory); this.nonPooledAvailableMemory -= got; accumulated += got; From 3e1fcd83c15991e8c7f7bffce980a2ef64dfa9d6 Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Fri, 3 Jul 2026 11:39:25 -0400 Subject: [PATCH 20/61] improve reservation when no mem available --- .../producer/internals/ChunkedBufferPool.java | 31 ++++++++++--------- 1 file changed, 17 insertions(+), 14 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java index 0bd540b09d963..1c36bc87b6d29 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java @@ -92,8 +92,9 @@ public List allocateChunks(int totalSize, long maxTimeToBlockMs) thr // A single Condition is added to the waiter's list to ensure FIFO fairness at the request level. // // `accumulated` tracks bytes drawn from nonPooledAvailableMemory only (pool chunks - // already taken live in `pooled`). Matches BufferPool.allocate's semantics: on - // failure, `accumulated` is exactly the amount to refund; on success it is reset to 0. + // already taken live in `pooled`), and is always a whole-chunk multiple. Matches + // BufferPool.allocate's semantics: on failure, `accumulated` is exactly the amount + // to refund; on success it is reset to 0. long accumulated = 0; Condition moreMemory = lock.newCondition(); try { @@ -124,20 +125,22 @@ public List allocateChunks(int totalSize, long maxTimeToBlockMs) thr remainingTimeToBlockNs -= timeNs; - // Take pool chunks first, then reserve non-pool bytes for the remainder. - while (pooled.size() < numChunks - && (long) (pooled.size() + 1) * chunkSize + accumulated <= memoryRequired - && !free.isEmpty()) { + // Reuse pooled chunks first, preferring them over raw reservations: if a + // taken chunk covers a slot already reserved as raw bytes in an earlier + // iteration, hand that raw reservation back to the pool. + while (pooled.size() < numChunks && !free.isEmpty()) { pooled.add(free.pollFirst()); + if (accumulated >= chunkSize) { // accumulated is always chunk-aligned + accumulated -= chunkSize; + this.nonPooledAvailableMemory += chunkSize; + } } - long stillNeeded = memoryRequired - (long) pooled.size() * chunkSize - accumulated; - if (stillNeeded > 0) { - // stillNeeded can exceed Integer.MAX_VALUE (memoryRequired rounds totalSize - // up to whole chunks) - freeUp((int) Math.min(stillNeeded, Integer.MAX_VALUE)); - long got = Math.min(stillNeeded, this.nonPooledAvailableMemory); - this.nonPooledAvailableMemory -= got; - accumulated += got; + // Reserve non-pooled memory for the still-uncovered chunks, in whole chunks + // (the buffers themselves are allocated after the lock is released). + while (pooled.size() + (int) (accumulated / chunkSize) < numChunks + && this.nonPooledAvailableMemory >= chunkSize) { + this.nonPooledAvailableMemory -= chunkSize; + accumulated += chunkSize; } } // Clear the rollback tracker. From 36b168d3a7c42fcc1e7e9606b9656b6c44109819 Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Fri, 3 Jul 2026 12:03:03 -0400 Subject: [PATCH 21/61] consolidate refund of pooled and non-pooled mem --- .../producer/internals/ChunkedBufferPool.java | 32 +++++++++---------- 1 file changed, 15 insertions(+), 17 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java index 1c36bc87b6d29..08940bb515c99 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java @@ -43,8 +43,8 @@ public ChunkedBufferPool(long memory, int chunkSize, Metrics metrics, Time time, * {@link BufferPool#allocate}: satisfied immediately if memory is available, else blocks up to * {@code maxTimeToBlockMs} for the whole request (FIFO on {@link #waiters}). * The reservation is tracked as bytes against {@link #nonPooledAvailableMemory} plus chunks polled - * from {@link #free}. Any failure refunds the whole reservation and signals - * the next waiter before the exception propagates, so no partial holds are visible during the wait. + * from {@link #free}. Any failure refunds the whole reservation and signals the next waiter + * before the exception propagates, so a failed request leaves nothing reserved. * * @param totalSize minimum total bytes of capacity required across the returned chunks * @param maxTimeToBlockMs maximum time in milliseconds to block waiting for memory @@ -92,10 +92,11 @@ public List allocateChunks(int totalSize, long maxTimeToBlockMs) thr // A single Condition is added to the waiter's list to ensure FIFO fairness at the request level. // // `accumulated` tracks bytes drawn from nonPooledAvailableMemory only (pool chunks - // already taken live in `pooled`), and is always a whole-chunk multiple. Matches - // BufferPool.allocate's semantics: on failure, `accumulated` is exactly the amount - // to refund; on success it is reset to 0. + // already taken live in `pooled`), and is always a whole-chunk multiple. If the wait + // does not complete (timeout / close / interrupt), the finally refunds the whole + // reservation: `accumulated` back to non-pooled memory, `pooled` back to the free chunks list. long accumulated = 0; + boolean allocationCompleted = false; Condition moreMemory = lock.newCondition(); try { long remainingTimeToBlockNs = TimeUnit.MILLISECONDS.toNanos(maxTimeToBlockMs); @@ -125,7 +126,7 @@ public List allocateChunks(int totalSize, long maxTimeToBlockMs) thr remainingTimeToBlockNs -= timeNs; - // Reuse pooled chunks first, preferring them over raw reservations: if a + // Reuse free-list chunks first, preferring them over raw reservations: if a // taken chunk covers a slot already reserved as raw bytes in an earlier // iteration, hand that raw reservation back to the pool. while (pooled.size() < numChunks && !free.isEmpty()) { @@ -143,21 +144,18 @@ public List allocateChunks(int totalSize, long maxTimeToBlockMs) thr accumulated += chunkSize; } } - // Clear the rollback tracker. - accumulated = 0; + allocationCompleted = true; } finally { - // On failure (timeout / close / interrupt), refund the non-pool bytes taken. - // Pool chunks already in `pooled` are returned to `free` separately by the - // outer catch. - this.nonPooledAvailableMemory += accumulated; + if (!allocationCompleted) { + // Refund all that was reserved (pooled chunks and non-pooled bytes) + this.nonPooledAvailableMemory += accumulated; + for (ByteBuffer chunk : pooled) + free.addFirst(chunk); + pooled.clear(); + } waiters.remove(moreMemory); } } - } catch (RuntimeException | InterruptedException e) { - // Any chunks taken from the pool must be returned. - for (ByteBuffer chunk : pooled) - free.addFirst(chunk); - throw e; } finally { try { signalNextWaiterIfMemoryAvailable(); From c4dd9283993b7eabe55b4f2d6d49973cdeb11eac Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Wed, 8 Jul 2026 15:14:51 -0400 Subject: [PATCH 22/61] add missing metric + test --- .../clients/producer/internals/BufferPool.java | 6 +++++- .../producer/internals/ChunkedBufferPool.java | 1 + .../internals/ChunkedBufferPoolTest.java | 16 ++++++++++++++++ 3 files changed, 22 insertions(+), 1 deletion(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/BufferPool.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/BufferPool.java index c515cf0a8d79f..bf975b47692ba 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/BufferPool.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/BufferPool.java @@ -160,7 +160,7 @@ public ByteBuffer allocate(int size, long maxTimeToBlockMs) throws InterruptedEx throw new KafkaException("Producer closed while allocating memory"); if (waitingTimeElapsed) { - this.metrics.sensor("buffer-exhausted-records").record(); + recordBufferExhausted(); throw new BufferExhaustedException("Failed to allocate " + size + " bytes within the configured max blocking time " + maxTimeToBlockMs + " ms. Total memory: " + totalMemory() + " bytes. Available memory: " + availableMemory() + " bytes. Poolable size: " + poolableSize() + " bytes"); @@ -213,6 +213,10 @@ protected void recordWaitTime(long timeNs) { this.waitTime.record(timeNs, time.milliseconds()); } + protected void recordBufferExhausted() { + this.metrics.sensor("buffer-exhausted-records").record(); + } + /** * Wake the longest-waiting thread if any memory (pooled or non-pooled) is available. * Must be called with {@link #lock} held. No-op if no waiters or no memory is free. diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java index 08940bb515c99..436c7c325332c 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java @@ -117,6 +117,7 @@ public List allocateChunks(int totalSize, long maxTimeToBlockMs) thr throw new KafkaException("Producer closed while allocating memory"); if (waitingTimeElapsed) { + recordBufferExhausted(); throw new BufferExhaustedException("Failed to allocate " + memoryRequired + " bytes (" + numChunks + " chunks of " + chunkSize + ") within the configured max blocking time " + maxTimeToBlockMs diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPoolTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPoolTest.java index 5b0f6f9b4fbb4..7ba4f5919b83c 100644 --- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPoolTest.java +++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPoolTest.java @@ -18,6 +18,7 @@ import org.apache.kafka.clients.producer.BufferExhaustedException; import org.apache.kafka.common.KafkaException; +import org.apache.kafka.common.metrics.KafkaMetric; import org.apache.kafka.common.metrics.Metrics; import org.apache.kafka.common.utils.MockTime; import org.apache.kafka.test.TestUtils; @@ -142,6 +143,21 @@ public void testRollbackOnPartialFailure() throws Exception { assertEquals(total, p.availableMemory()); } + /** A drop due to pool exhaustion increments the shared buffer-exhausted metric. */ + @Test + public void testBufferExhaustedMetricRecordedOnTimeout() throws Exception { + int chunkSize = 64; + ChunkedBufferPool p = pool(chunkSize, chunkSize); // room for exactly 1 chunk + p.allocate(chunkSize, 100); // drain the pool + + KafkaMetric exhausted = metrics.metric(metrics.metricName("buffer-exhausted-total", "producer-metrics")); + assertEquals(0.0, (double) exhausted.metricValue()); + + // Zero deadline → immediate timeout → BufferExhaustedException, metric recorded. + assertThrows(BufferExhaustedException.class, () -> p.allocateChunks(chunkSize, 0)); + assertEquals(1.0, (double) exhausted.metricValue()); + } + /** * No partial chunk holds during the wait. While a multi-chunk request blocks, the pool's * {@code availableMemory()} must reflect bytes the waiter has not yet "earned". From a246f9d3e8793432ca7fe290cac3b4cf77ca3645 Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Wed, 8 Jul 2026 15:15:25 -0400 Subject: [PATCH 23/61] O(1) capacity calculation --- .../producer/internals/ChunkedByteBufferOutputStream.java | 8 ++++++++ .../clients/producer/internals/ChunkedProducerBatch.java | 5 +---- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java index cf5a923f12adc..68be16c82290f 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java @@ -228,6 +228,14 @@ public void position(int position) { dirty = true; } + /** + * Total capacity across all attached chunks (written + free). + * Every chunk has the same size, so this equals {@code position() + remaining()} without walking the list. + */ + int attachedCapacity() { + return chunks.size() * chunkSize; + } + /** * Total bytes available across the current chunk and every queued (not-yet-active) chunk. */ diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedProducerBatch.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedProducerBatch.java index 4691efbf9e7fa..62d8425240789 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedProducerBatch.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedProducerBatch.java @@ -20,7 +20,6 @@ import org.apache.kafka.common.TopicPartition; import org.apache.kafka.common.header.Header; import org.apache.kafka.common.record.internal.MemoryRecordsBuilder; -import org.apache.kafka.common.utils.ByteBufferOutputStream; import java.nio.ByteBuffer; import java.util.List; @@ -59,9 +58,7 @@ int extensionBytesNeeded(long timestamp, byte[] key, byte[] value, Header[] head // ratio-adjusted when compressed), not per-record. Per-record sizing would over-count the // header and miss the compressor's flush-accumulation behavior. int target = recordsBuilder.estimatedBytesWrittenAfter(key, value, headers); - ByteBufferOutputStream stream = recordsBuilder.bufferStream(); - int totalAttachedCapacity = stream.position() + stream.remaining(); - return Math.max(0, target - totalAttachedCapacity); + return Math.max(0, target - stream().attachedCapacity()); } /** From c645114956b43a9fa08b365e1d05c41934234e28 Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Thu, 9 Jul 2026 12:12:21 -0400 Subject: [PATCH 24/61] fix for overcounting and test --- .../producer/internals/BufferPool.java | 5 +++ .../producer/internals/ChunkedBufferPool.java | 5 ++- .../internals/ChunkedRecordAccumulator.java | 10 +++++- .../internals/ChunkedBufferPoolTest.java | 16 ---------- .../ChunkedRecordAccumulatorTest.java | 32 +++++++++++++++++++ 5 files changed, 50 insertions(+), 18 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/BufferPool.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/BufferPool.java index bf975b47692ba..9bd36c9410ff0 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/BufferPool.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/BufferPool.java @@ -213,6 +213,11 @@ protected void recordWaitTime(long timeNs) { this.waitTime.record(timeNs, time.milliseconds()); } + /** + * Record that a record send was dropped because the buffer pool was exhausted. Shared by the + * full strategy (allocate) and the incremental strategy (ChunkedRecordAccumulator), + * so both update the same buffer-exhausted metrics. + */ protected void recordBufferExhausted() { this.metrics.sensor("buffer-exhausted-records").record(); } diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java index 436c7c325332c..749a27c9d46be 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java @@ -117,7 +117,10 @@ public List allocateChunks(int totalSize, long maxTimeToBlockMs) thr throw new KafkaException("Producer closed while allocating memory"); if (waitingTimeElapsed) { - recordBufferExhausted(); + // Intentionally not recording the buffer-exhausted metric here. + // This may be called on the extension path, which recovers without + // dropping the record, so the caller is the one recording + // the drop if needed. throw new BufferExhaustedException("Failed to allocate " + memoryRequired + " bytes (" + numChunks + " chunks of " + chunkSize + ") within the configured max blocking time " + maxTimeToBlockMs diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java index 75cdd8ee98c74..c5ecaf6f17e41 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java @@ -189,7 +189,15 @@ public RecordAppendResult append(String topic, recordUncompressed); log.trace("Allocating {} byte chunked buffer ({} byte chunks) for topic {} partition {} with remaining timeout {}ms", size, chunkedFree.poolableSize(), topic, effectivePartition, maxTimeToBlock); - List initialChunks = chunkedFree.allocateChunks(size, maxTimeToBlock); + List initialChunks; + try { + initialChunks = chunkedFree.allocateChunks(size, maxTimeToBlock); + } catch (BufferExhaustedException e) { + // The blocking new-batch acquire was not able to get memory within + // max.block.ms. Record it in the buffer-exhausted metrics. + chunkedFree.recordBufferExhausted(); + throw e; + } nowMs = time.milliseconds(); bufferStream = new ChunkedByteBufferOutputStream(initialChunks, chunkedFree.poolableSize(), chunkedFree); } diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPoolTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPoolTest.java index 7ba4f5919b83c..5b0f6f9b4fbb4 100644 --- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPoolTest.java +++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPoolTest.java @@ -18,7 +18,6 @@ import org.apache.kafka.clients.producer.BufferExhaustedException; import org.apache.kafka.common.KafkaException; -import org.apache.kafka.common.metrics.KafkaMetric; import org.apache.kafka.common.metrics.Metrics; import org.apache.kafka.common.utils.MockTime; import org.apache.kafka.test.TestUtils; @@ -143,21 +142,6 @@ public void testRollbackOnPartialFailure() throws Exception { assertEquals(total, p.availableMemory()); } - /** A drop due to pool exhaustion increments the shared buffer-exhausted metric. */ - @Test - public void testBufferExhaustedMetricRecordedOnTimeout() throws Exception { - int chunkSize = 64; - ChunkedBufferPool p = pool(chunkSize, chunkSize); // room for exactly 1 chunk - p.allocate(chunkSize, 100); // drain the pool - - KafkaMetric exhausted = metrics.metric(metrics.metricName("buffer-exhausted-total", "producer-metrics")); - assertEquals(0.0, (double) exhausted.metricValue()); - - // Zero deadline → immediate timeout → BufferExhaustedException, metric recorded. - assertThrows(BufferExhaustedException.class, () -> p.allocateChunks(chunkSize, 0)); - assertEquals(1.0, (double) exhausted.metricValue()); - } - /** * No partial chunk holds during the wait. While a multi-chunk request blocks, the pool's * {@code availableMemory()} must reflect bytes the waiter has not yet "earned". diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java index 58363f4012f33..6fa8a7c29fe4a 100644 --- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java +++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java @@ -22,6 +22,7 @@ import org.apache.kafka.common.Node; import org.apache.kafka.common.TopicPartition; import org.apache.kafka.common.compress.Compression; +import org.apache.kafka.common.metrics.KafkaMetric; import org.apache.kafka.common.metrics.Metrics; import org.apache.kafka.common.protocol.Errors; import org.apache.kafka.common.record.internal.MemoryRecords; @@ -461,6 +462,37 @@ public void closeForRecordAppends() { } } + /** + * A single dropped record is counted exactly once even when it first fails the extension attempt + * (recovered) and then fails the new-batch acquire. + * Uses a real (non-overridden) pool so the actual allocateChunks path runs on both acquires. + */ + @Test + public void testBufferExhaustedNotDoubleCountedAcrossExtensionAndNewBatch() throws Exception { + int chunkSize = 256; + // Pool holds exactly one chunk: the first record consumes it, leaving the pool empty so both + // the second record's extension attempt and its new-batch acquire fail. + ChunkedRecordAccumulator accum = newAccumulator(8192, chunkSize, chunkSize, Compression.NONE); + try { + KafkaMetric exhausted = metrics.metric(metrics.metricName("buffer-exhausted-total", "producer-metrics")); + + // First record fills the one available chunk (pool now empty), opening the batch. + accum.append(topic, partition1, 0L, key, new byte[150], Record.EMPTY_HEADERS, null, + maxBlockTimeMs, time.milliseconds(), cluster); + assertEquals(0.0, (double) exhausted.metricValue()); + + // Second record overflows the batch's chunk, so extension attempt fails (pool empty, + // recovered by closing the batch), then the new-batch acquire fails too (maxTimeToBlock + // = 0). The two failed acquires must count as a single dropped record. + assertThrows(BufferExhaustedException.class, () -> accum.append(topic, partition1, 0L, key, + new byte[150], Record.EMPTY_HEADERS, null, 0L, time.milliseconds(), cluster)); + assertEquals(1.0, (double) exhausted.metricValue(), + "the dropped record must be counted once, not once per failed acquire"); + } finally { + accum.close(); + } + } + /** * Test that chunks are returned to the pool only at batch completion (deallocate), never at close. */ From 803816a0363df49cbb41f3c3d543fc0cc6e17be5 Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Mon, 13 Jul 2026 09:47:55 -0400 Subject: [PATCH 25/61] release unused on stream close & tests --- .../ChunkedByteBufferOutputStream.java | 20 +++--- .../ChunkedByteBufferOutputStreamTest.java | 62 ++++++++++--------- .../ChunkedRecordAccumulatorTest.java | 47 +++++++++----- 3 files changed, 76 insertions(+), 53 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java index 68be16c82290f..74a0388b253ed 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java @@ -162,20 +162,24 @@ public ByteBuffer buffer() { chunk.position(chunkPos); } dirty = false; - releaseUnusedChunks(); return flattenedBuffer; } /** - * Release the fully-unused chunks (nothing written) back to the pool, now rather than at batch - * completion. Called at batch close, when the chunk set is final. - *

- * The data-bearing chunks are intentionally kept until batch completion: they are what keeps - * the batch's in-flight data reserved against the buffer.memory budget, so releasing them at - * close would let the pool admit more data than buffer.memory intends while the batch is still - * in flight. + * Releases the fully-unused chunks, given that the stream is closed for appends. + */ + @Override + public void close() { + releaseUnusedChunks(); + } + + /** + * Return the fully-unused chunks to the pool. The data-bearing chunks are + * kept until batch completion ({@link #deallocate()}), as they hold the in-flight data. */ private void releaseUnusedChunks() { + if (currentChunk == null) // already deallocated; nothing attached + return; List unused = chunks.subList(currentChunkIndex + 1, chunks.size()); if (pool != null) { for (ByteBuffer chunk : unused) diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java index efa7fb6923322..f00e74b687b92 100644 --- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java +++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java @@ -173,40 +173,42 @@ public void testDeallocateReturnsAllChunks() throws Exception { } /** - * Closing the batch (which invokes {@code buffer()}) returns fully-unused chunks to the pool - * right away; data-bearing chunks stay reserved until - * {@link ChunkedByteBufferOutputStream#deallocate()}, which must not return the - * already-released chunks a second time. + * The fully-unused chunks are returned to the pool on {@link ChunkedByteBufferOutputStream#close()} + * (the stream is closed for appends), not as a side effect of reading {@link + * ChunkedByteBufferOutputStream#buffer()}. Data-bearing chunks stay reserved until + * {@link ChunkedByteBufferOutputStream#deallocate()}, which must not return the already-released + * chunks a second time. */ @Test - public void testUnusedChunksReleasedAtClose() throws Exception { + public void testUnusedChunksReleasedOnCloseNotOnBuffer() throws Exception { int chunkSize = 8; long total = 64; ChunkedBufferPool p = pool(total, chunkSize); - try (ChunkedByteBufferOutputStream stream = new ChunkedByteBufferOutputStream(chunks(p, chunkSize, 3), chunkSize, p)) { - // Write into the first chunk only; chunks 2 and 3 stay untouched. - byte[] payload = new byte[]{1, 2, 3}; - stream.write(payload, 0, payload.length); - assertEquals(total - 3L * chunkSize, p.availableMemory()); - - ByteBuffer built = stream.buffer(); - assertEquals(total - chunkSize, p.availableMemory(), - "the two unused chunks should return to the pool at close"); - - // The written content is intact. - built.flip(); - byte[] out = new byte[built.remaining()]; - built.get(out); - assertArrayEquals(payload, out); - - // A repeated buffer call must not release anything further. - stream.buffer(); - assertEquals(total - chunkSize, p.availableMemory()); - - // Completion-time deallocate must return only the remaining data-bearing chunk. - stream.deallocate(); - assertEquals(total, p.availableMemory(), - "pool must be exactly restored; released chunks must not be returned twice on completion"); - } + ChunkedByteBufferOutputStream stream = new ChunkedByteBufferOutputStream(chunks(p, chunkSize, 3), chunkSize, p); + // Write into the first chunk only; chunks 2 and 3 stay unused. + byte[] payload = new byte[]{1, 2, 3}; + stream.write(payload, 0, payload.length); + assertEquals(total - 3L * chunkSize, p.availableMemory()); + + // reading buffer() on a still-open stream must not release chunks. + ByteBuffer built = stream.buffer(); + assertEquals(total - 3L * chunkSize, p.availableMemory(), + "buffer() must not release chunks on a still-open stream"); + built.flip(); + byte[] out = new byte[built.remaining()]; + built.get(out); + assertArrayEquals(payload, out); + + // close() (appends done) releases the two unused chunks; a second close() is a no-op. + stream.close(); + assertEquals(total - chunkSize, p.availableMemory(), + "the two unused chunks should return to the pool on close"); + stream.close(); + assertEquals(total - chunkSize, p.availableMemory()); + + // Completion-time deallocate returns only the remaining data-bearing chunk (no double free). + stream.deallocate(); + assertEquals(total, p.availableMemory(), + "pool must be exactly restored; released chunks must not be returned twice on completion"); } } \ No newline at end of file diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java index 6fa8a7c29fe4a..3d0570f0ecad4 100644 --- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java +++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java @@ -222,7 +222,8 @@ public void testConcurrentExtensionRaceLoserExtendsAgain() throws Exception { * remaining capacity, so a batch can end up with more chunk capacity than its batch-size limit * admits. The limit must still be enforced on append: a record that no longer fits after the * race winner's append rolls to a new batch even though the over-reserved chunks could - * physically hold it, and the unused surplus returns to the pool when the batch closes. + * physically hold it, and the unused surplus returns to the pool as soon as the batch is closed + * for appends during that roll (before completion). */ @Test public void testConcurrentExtensionRaceLoserStartsNewBatch() throws Exception { @@ -231,8 +232,25 @@ public void testConcurrentExtensionRaceLoserStartsNewBatch() throws Exception { // Sized so each record fits the batch-size limit individually, but not both together. final byte[] value = new byte[500]; final AtomicReference injectAppendOnce = new AtomicReference<>(); + final AtomicInteger chunkDeallocations = new AtomicInteger(); - ChunkedBufferPool pool = poolMockingConcurrentChunkAllocation(chunkSize, 64L * chunkSize, injectAppendOnce, value); + ChunkedBufferPool pool = new ChunkedBufferPool(64L * chunkSize, chunkSize, metrics, time, "producer-metrics") { + @Override + public List allocateChunks(int totalSize, long maxTimeToBlockMs) throws InterruptedException { + List chunks = super.allocateChunks(totalSize, maxTimeToBlockMs); + ChunkedRecordAccumulator toInject = injectAppendOnce.getAndSet(null); + if (toInject != null) + toInject.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null, + maxBlockTimeMs, time.milliseconds(), cluster); + return chunks; + } + + @Override + public void deallocate(ByteBuffer buffer) { + chunkDeallocations.incrementAndGet(); + super.deallocate(buffer); + } + }; ChunkedRecordAccumulator accum = new ChunkedRecordAccumulator(logContext, batchSize, Compression.NONE, /* lingerMs */ 0, /* retryBackoffMs */ 0L, /* retryBackoffMaxMs */ 0L, /* deliveryTimeoutMs */ 3200, metrics, "producer-metrics", time, @@ -241,6 +259,7 @@ public void testConcurrentExtensionRaceLoserStartsNewBatch() throws Exception { // Open the batch with a tiny record; both racing records will extend it. accum.append(topic, partition1, 0L, key, new byte[1], Record.EMPTY_HEADERS, null, maxBlockTimeMs, time.milliseconds(), cluster); + int deallocsAfterOpen = chunkDeallocations.get(); // This append sizes its gap, then the injected append wins the race while the gap // chunks are held off-lock. The retry finds the batch over its batch-size limit, so @@ -251,19 +270,17 @@ public void testConcurrentExtensionRaceLoserStartsNewBatch() throws Exception { Deque dq = batchesFor(accum, tp1); assertEquals(2, dq.size(), "The losing record must roll to a new batch, not exceed batch.size"); - ProducerBatch first = dq.peekFirst(); - ProducerBatch second = dq.peekLast(); - assertNotNull(first); - assertNotNull(second); - assertEquals(2, first.recordCount, "First batch should have the initial record plus the race winner"); - assertEquals(1, second.recordCount, "Second batch should have the loser record"); - - // The loser's extension chunks were attached but never used (the over-reservation); - // closing the first batch returns that unused surplus to the pool. - long beforeClose = pool.availableMemory(); - first.close(); - assertTrue(pool.availableMemory() > beforeClose, - "The over-reserved unused chunks should return to the pool at close"); + assertNotNull(dq.peekFirst()); + assertNotNull(dq.peekLast()); + assertEquals(2, dq.peekFirst().recordCount, "First batch should have the initial record plus the race winner"); + assertNotNull(dq.peekLast()); + assertEquals(1, dq.peekLast().recordCount, "Second batch should have the loser record"); + + // The loser's extension chunks were attached to the first batch but never used (the + // over-reservation). They are freed as soon as the first batch is closed for appends + // (during the roll to the new batch) — before completion, not held until deallocate. + assertTrue(chunkDeallocations.get() > deallocsAfterOpen, + "the over-reserved unused chunks are freed when the first batch is closed for appends"); } finally { accum.close(); } From 73019661fe9e6c80ea64f7263869ac309e6d56d3 Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Wed, 15 Jul 2026 08:31:15 -0400 Subject: [PATCH 26/61] indentation, docs & asserts --- .../internals/ChunkedProducerBatch.java | 11 ++++++---- .../internals/ChunkedRecordAccumulator.java | 5 ++++- .../producer/internals/RecordAccumulator.java | 22 +++++++++++-------- 3 files changed, 24 insertions(+), 14 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedProducerBatch.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedProducerBatch.java index 62d8425240789..d273afbae1edb 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedProducerBatch.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedProducerBatch.java @@ -62,17 +62,20 @@ int extensionBytesNeeded(long timestamp, byte[] key, byte[] value, Header[] head } /** - * Appends the record after asserting there's capacity for it. + * Appends the record after checking there's chunk capacity for it. * At this point it's expected that the allocated chunks have * capacity for the record because the accumulator never routes an * append here without ensuring chunk capacity first: the first-record path pre-sizes the * stream for the record and the mid-batch path attaches extension chunks before retrying. + * + * @throws IllegalStateException if the attached chunks lack capacity for the record */ @Override public FutureRecordMetadata tryAppend(long timestamp, byte[] key, byte[] value, Header[] headers, Callback callback, long now) { - assert extensionBytesNeeded(timestamp, key, value, headers) == 0 : - "Unexpected append to a chunked batch whose chunks lack capacity for the record; " + - "the accumulator should have attached extension chunks first"; + if (extensionBytesNeeded(timestamp, key, value, headers) != 0) + throw new IllegalStateException( + "Unexpected append to a chunked batch whose chunks lack capacity for the record; " + + "the accumulator should have attached extension chunks first"); return super.tryAppend(timestamp, key, value, headers, callback, now); } diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java index c5ecaf6f17e41..2c8664128e067 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java @@ -246,7 +246,8 @@ public RecordAppendResult append(String topic, // needsNewBatch path: extensionChunks == null here implies needsNewBatch, // so bufferStream was allocated (this iteration or carried from a prior one). - assert bufferStream != null; + if (bufferStream == null) + throw new IllegalStateException("needsNewBatch path reached without an allocated buffer stream"); int firstRecordSize = AbstractRecords.estimateSizeInBytesUpperBound( RecordBatch.CURRENT_MAGIC_VALUE, compression.type(), key, value, headers); final ChunkedByteBufferOutputStream batchStream = bufferStream; @@ -261,6 +262,8 @@ public RecordAppendResult append(String topic, bufferStream = null; continue; } + if (appendResult.needsNewBatch) + throw new IllegalStateException("appendNewBatch must not return a needsNewBatch result"); if (appendResult.newBatchCreated) bufferStream = null; boolean enableSwitch = allBatchesFull(dq); diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/RecordAccumulator.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/RecordAccumulator.java index 2cf7bbaac66d8..465bb0bf56b0d 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/RecordAccumulator.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/RecordAccumulator.java @@ -384,17 +384,21 @@ public RecordAppendResult append(String topic, * subclass passes a supplier that produces a builder backed by a * {@link ChunkedByteBufferOutputStream}. * @param nowMs The current time, in milliseconds + * @return the append result, which never has {@code needsNewBatch=true}: the method either + * propagates a non-{@code needsNewBatch} result from an open batch created concurrently + * (a success, or — incremental strategy — a {@code needsBufferExtension} signal), or it + * creates the new batch here and returns the successful append to it. */ protected RecordAppendResult appendNewBatch(String topic, - int partition, - Deque dq, - long timestamp, - byte[] key, - byte[] value, - Header[] headers, - AppendCallbacks callbacks, - Supplier recordsBuilderSupplier, - long nowMs) { + int partition, + Deque dq, + long timestamp, + byte[] key, + byte[] value, + Header[] headers, + AppendCallbacks callbacks, + Supplier recordsBuilderSupplier, + long nowMs) { assert partition != RecordMetadata.UNKNOWN_PARTITION; RecordAppendResult appendResult = tryAppend(timestamp, key, value, headers, callbacks, dq, nowMs); From 1d1dcd87cde503bab133b712907095d1386ed6a0 Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Wed, 15 Jul 2026 08:33:13 -0400 Subject: [PATCH 27/61] rename --- .../producer/internals/ChunkedBufferPool.java | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java index 749a27c9d46be..ab7abf489fdff 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java @@ -91,17 +91,17 @@ public List allocateChunks(int totalSize, long maxTimeToBlockMs) thr // Same as in BufferPool.allocate, but wait to acquire the memory needed for all the chunks. // A single Condition is added to the waiter's list to ensure FIFO fairness at the request level. // - // `accumulated` tracks bytes drawn from nonPooledAvailableMemory only (pool chunks + // `nonPoolAccumulated` tracks bytes drawn from nonPooledAvailableMemory only (pool chunks // already taken live in `pooled`), and is always a whole-chunk multiple. If the wait // does not complete (timeout / close / interrupt), the finally refunds the whole - // reservation: `accumulated` back to non-pooled memory, `pooled` back to the free chunks list. - long accumulated = 0; + // reservation: `nonPoolAccumulated` back to non-pooled memory, `pooled` back to the free chunks list. + long nonPoolAccumulated = 0; boolean allocationCompleted = false; Condition moreMemory = lock.newCondition(); try { long remainingTimeToBlockNs = TimeUnit.MILLISECONDS.toNanos(maxTimeToBlockMs); waiters.addLast(moreMemory); - while ((long) pooled.size() * chunkSize + accumulated < memoryRequired) { + while ((long) pooled.size() * chunkSize + nonPoolAccumulated < memoryRequired) { long startWaitNs = time.nanoseconds(); long timeNs; boolean waitingTimeElapsed; @@ -135,24 +135,24 @@ public List allocateChunks(int totalSize, long maxTimeToBlockMs) thr // iteration, hand that raw reservation back to the pool. while (pooled.size() < numChunks && !free.isEmpty()) { pooled.add(free.pollFirst()); - if (accumulated >= chunkSize) { // accumulated is always chunk-aligned - accumulated -= chunkSize; + if (nonPoolAccumulated >= chunkSize) { // nonPoolAccumulated is always chunk-aligned + nonPoolAccumulated -= chunkSize; this.nonPooledAvailableMemory += chunkSize; } } // Reserve non-pooled memory for the still-uncovered chunks, in whole chunks // (the buffers themselves are allocated after the lock is released). - while (pooled.size() + (int) (accumulated / chunkSize) < numChunks + while (pooled.size() + (int) (nonPoolAccumulated / chunkSize) < numChunks && this.nonPooledAvailableMemory >= chunkSize) { this.nonPooledAvailableMemory -= chunkSize; - accumulated += chunkSize; + nonPoolAccumulated += chunkSize; } } allocationCompleted = true; } finally { if (!allocationCompleted) { // Refund all that was reserved (pooled chunks and non-pooled bytes) - this.nonPooledAvailableMemory += accumulated; + this.nonPooledAvailableMemory += nonPoolAccumulated; for (ByteBuffer chunk : pooled) free.addFirst(chunk); pooled.clear(); From dea7c880acf9f068025672ac3c83cb82a53749f2 Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Wed, 15 Jul 2026 10:17:50 -0400 Subject: [PATCH 28/61] close for writes after buffer & test --- .../ChunkedByteBufferOutputStream.java | 50 +++++++++++++------ .../ChunkedByteBufferOutputStreamTest.java | 22 ++++++++ 2 files changed, 58 insertions(+), 14 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java index 74a0388b253ed..6e2ac7f880924 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java @@ -45,8 +45,11 @@ public class ChunkedByteBufferOutputStream extends ByteBufferOutputStream { private final BufferPool pool; private ByteBuffer currentChunk; private int currentChunkIndex; + // Set once the content has been read for send via buffer(). + private boolean finalized; + // Single-buffer view produced by flatten() and cached here so repeat buffer() calls + // return the same instance. To be removed once scatter-gather (KAFKA-20580) is implemented. private ByteBuffer flattenedBuffer; - private boolean dirty; /** * Constructs a chunked output stream backed by the given pre-allocated chunks. Ownership of @@ -65,7 +68,6 @@ public ChunkedByteBufferOutputStream(List initialChunks, int chunkSi this.chunks = new ArrayList<>(initialChunks); this.currentChunk = this.chunks.get(0); this.currentChunkIndex = 0; - this.dirty = true; } /** @@ -85,13 +87,14 @@ private static ByteBuffer validatedFirstChunk(List initialChunks, in @Override public void write(int b) { + ensureWritable(); ensureChunkCapacity(1); currentChunk.put((byte) b); - dirty = true; } @Override public void write(byte[] bytes, int off, int len) { + ensureWritable(); while (len > 0) { ensureChunkCapacity(1); int toWrite = Math.min(len, currentChunk.remaining()); @@ -99,11 +102,11 @@ public void write(byte[] bytes, int off, int len) { off += toWrite; len -= toWrite; } - dirty = true; } @Override public void write(ByteBuffer sourceBuffer) { + ensureWritable(); while (sourceBuffer.hasRemaining()) { ensureChunkCapacity(1); int toWrite = Math.min(sourceBuffer.remaining(), currentChunk.remaining()); @@ -112,7 +115,14 @@ public void write(ByteBuffer sourceBuffer) { currentChunk.put(sourceBuffer); sourceBuffer.limit(oldLimit); } - dirty = true; + } + + /** + * Guards against writes after the stream has been finalized with a call to {@link #buffer()}. + */ + private void ensureWritable() { + if (finalized) + throw new IllegalStateException("cannot write after buffer() has been called"); } private void ensureChunkCapacity(int needed) { @@ -141,28 +151,41 @@ public void addBuffers(List newChunks) { chunks.addAll(newChunks); } + /** + * Returns the written bytes as a single contiguous buffer. Calling this finalizes the stream: no + * further writes are allowed (see {@link #ensureWritable()}) and later calls return the same + * instance, which callers such as {@code MemoryRecordsBuilder#writeDefaultBatchHeader} rely on + * when they write the batch header directly into the returned buffer. + */ @Override public ByteBuffer buffer() { - if (flattenedBuffer != null && !dirty) { - return flattenedBuffer; - } - // Written bytes only live in chunks up to currentChunk; later chunks are untouched. + finalized = true; + if (flattenedBuffer == null) + flattenedBuffer = flatten(); + return flattenedBuffer; + } + + /** + * Flattens the written bytes across the data-bearing chunks into a single new buffer (an extra + * copy). This will be removed once scatter-gather send (KAFKA-20580) is implemented. + */ + private ByteBuffer flatten() { + // Written bytes only live in chunks up to currentChunk, later chunks are untouched. int lastDataChunk = Math.min(currentChunkIndex, chunks.size() - 1); int totalSize = 0; for (int i = 0; i <= lastDataChunk; i++) { totalSize += chunks.get(i).position(); } - flattenedBuffer = ByteBuffer.allocate(totalSize); + ByteBuffer flattened = ByteBuffer.allocate(totalSize); for (int i = 0; i <= lastDataChunk; i++) { ByteBuffer chunk = chunks.get(i); int chunkPos = chunk.position(); chunk.flip(); - flattenedBuffer.put(chunk); + flattened.put(chunk); chunk.limit(chunk.capacity()); chunk.position(chunkPos); } - dirty = false; - return flattenedBuffer; + return flattened; } /** @@ -229,7 +252,6 @@ public void position(int position) { } currentChunkIndex = idx; currentChunk = chunks.get(idx); - dirty = true; } /** diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java index f00e74b687b92..598be830925f7 100644 --- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java +++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java @@ -28,6 +28,7 @@ import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertThrows; public class ChunkedByteBufferOutputStreamTest { @@ -50,6 +51,7 @@ private List chunks(ChunkedBufferPool pool, int chunkSize, int count } @Test + @SuppressWarnings({"try", "resource"}) // Each constructor call throws, so no stream is ever created public void testConstructorRejectsInvalidChunks() { int chunkSize = 16; ChunkedBufferPool p = pool(64, chunkSize); @@ -65,6 +67,26 @@ public void testConstructorRejectsInvalidChunks() { () -> new ChunkedByteBufferOutputStream(wrongSize, chunkSize, p)); } + @Test + public void testWritesDisallowedAfterBuffer() throws Exception { + int chunkSize = 16; + ChunkedBufferPool p = pool(64, chunkSize); + try (ChunkedByteBufferOutputStream stream = new ChunkedByteBufferOutputStream(chunks(p, chunkSize, 2), chunkSize, p)) { + stream.write(new byte[]{1, 2, 3}, 0, 3); + + // buffer() finalizes the stream: it may be called repeatedly (returning the same + // instance), but any subsequent write must fail. + ByteBuffer first = stream.buffer(); + assertSame(first, stream.buffer(), "buffer() must return the same cached instance once built"); + + assertThrows(IllegalStateException.class, () -> stream.write(1)); + assertThrows(IllegalStateException.class, () -> stream.write(new byte[]{4}, 0, 1)); + assertThrows(IllegalStateException.class, () -> stream.write(ByteBuffer.wrap(new byte[]{5}))); + + stream.deallocate(); + } + } + @Test public void testSingleChunkWriteRoundtrip() throws Exception { int chunkSize = 16; From 1e5a006f82cfbc61c80b57454e35019ae2e9c05c Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Wed, 15 Jul 2026 10:34:32 -0400 Subject: [PATCH 29/61] clear index --- .../producer/internals/ChunkedByteBufferOutputStream.java | 1 + 1 file changed, 1 insertion(+) diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java index 6e2ac7f880924..053815c0cda1b 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java @@ -304,6 +304,7 @@ public void deallocate(BufferPool pool) { } chunks.clear(); currentChunk = null; + currentChunkIndex = 0; flattenedBuffer = null; } From 1387a8ba3f8f899e77a6512caa66e09744fd9e77 Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Wed, 15 Jul 2026 10:52:24 -0400 Subject: [PATCH 30/61] validate deallocated & test --- .../ChunkedByteBufferOutputStream.java | 20 +++++++++++-- .../ChunkedByteBufferOutputStreamTest.java | 30 +++++++++++++++++++ 2 files changed, 48 insertions(+), 2 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java index 053815c0cda1b..afbcd120e57cf 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java @@ -87,6 +87,7 @@ private static ByteBuffer validatedFirstChunk(List initialChunks, in @Override public void write(int b) { + ensureNotDeallocated(); ensureWritable(); ensureChunkCapacity(1); currentChunk.put((byte) b); @@ -94,6 +95,7 @@ public void write(int b) { @Override public void write(byte[] bytes, int off, int len) { + ensureNotDeallocated(); ensureWritable(); while (len > 0) { ensureChunkCapacity(1); @@ -106,6 +108,7 @@ public void write(byte[] bytes, int off, int len) { @Override public void write(ByteBuffer sourceBuffer) { + ensureNotDeallocated(); ensureWritable(); while (sourceBuffer.hasRemaining()) { ensureChunkCapacity(1); @@ -125,6 +128,14 @@ private void ensureWritable() { throw new IllegalStateException("cannot write after buffer() has been called"); } + /** + * Guards against any use after {@link #deallocate()} has returned the chunks. + */ + private void ensureNotDeallocated() { + if (currentChunk == null) + throw new IllegalStateException("operation not allowed after the stream has been deallocated"); + } + private void ensureChunkCapacity(int needed) { while (currentChunk.remaining() < needed) { advanceToNextChunk(); @@ -148,6 +159,7 @@ private void advanceToNextChunk() { * the stream; they will be returned to the pool via {@link #deallocate()}. */ public void addBuffers(List newChunks) { + ensureNotDeallocated(); chunks.addAll(newChunks); } @@ -159,6 +171,7 @@ public void addBuffers(List newChunks) { */ @Override public ByteBuffer buffer() { + ensureNotDeallocated(); finalized = true; if (flattenedBuffer == null) flattenedBuffer = flatten(); @@ -218,6 +231,7 @@ private void releaseUnusedChunks() { */ @Override public int position() { + ensureNotDeallocated(); // Written bytes only live in chunks up to currentChunk; later chunks are untouched. int lastDataChunk = Math.min(currentChunkIndex, chunks.size() - 1); int total = 0; @@ -233,6 +247,7 @@ public int position() { */ @Override public void position(int position) { + ensureNotDeallocated(); if (currentChunkIndex != 0 || currentChunk.position() != 0) { throw new IllegalStateException("position() can only be called before any writes"); } @@ -259,6 +274,7 @@ public void position(int position) { * Every chunk has the same size, so this equals {@code position() + remaining()} without walking the list. */ int attachedCapacity() { + ensureNotDeallocated(); return chunks.size() * chunkSize; } @@ -267,8 +283,7 @@ int attachedCapacity() { */ @Override public int remaining() { - if (currentChunk == null) // after deallocate no chunks attached, no free capacity - return 0; + ensureNotDeallocated(); int total = currentChunk.remaining(); for (int i = currentChunkIndex + 1; i < chunks.size(); i++) total += chunks.get(i).remaining(); @@ -287,6 +302,7 @@ public int initialCapacity() { @Override public void ensureRemaining(int remainingBytesRequired) { + ensureNotDeallocated(); // A single call can guarantee at most `chunkSize` of space (the stream advances one chunk // at a time). Callers needing more attach chunks via addBuffers first. write(byte[]) loops // across chunks, so contiguous capacity isn't required. diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java index 598be830925f7..49b58a08ce6d0 100644 --- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java +++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java @@ -27,6 +27,7 @@ import java.util.List; import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertThrows; @@ -67,6 +68,35 @@ public void testConstructorRejectsInvalidChunks() { () -> new ChunkedByteBufferOutputStream(wrongSize, chunkSize, p)); } + @Test + public void testOperationsRejectedAfterDeallocate() throws Exception { + int chunkSize = 8; + ChunkedBufferPool p = pool(64, chunkSize); + try (ChunkedByteBufferOutputStream stream = new ChunkedByteBufferOutputStream(chunks(p, chunkSize, 2), chunkSize, p)) { + stream.write(new byte[]{1, 2, 3}, 0, 3); + + stream.deallocate(); + + // Every query/write operation must reject use after deallocation. + assertThrows(IllegalStateException.class, stream::remaining); + assertThrows(IllegalStateException.class, stream::position); + assertThrows(IllegalStateException.class, stream::buffer); + assertThrows(IllegalStateException.class, stream::attachedCapacity); + assertThrows(IllegalStateException.class, () -> stream.position(1)); + assertThrows(IllegalStateException.class, () -> stream.ensureRemaining(1)); + assertThrows(IllegalStateException.class, () -> stream.write(1)); + assertThrows(IllegalStateException.class, () -> stream.write(new byte[]{4}, 0, 1)); + assertThrows(IllegalStateException.class, () -> stream.write(ByteBuffer.wrap(new byte[]{5}))); + assertThrows(IllegalStateException.class, () -> stream.addBuffers(Collections.singletonList(ByteBuffer.allocate(chunkSize)))); + + // Lifecycle calls stay idempotent no-ops: a second close()/deallocate() must not throw. + assertDoesNotThrow(() -> { + stream.close(); + stream.deallocate(); + }); + } + } + @Test public void testWritesDisallowedAfterBuffer() throws Exception { int chunkSize = 16; From 9996c7749d3557d8d0826f85ba2dfef3c800a328 Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Wed, 15 Jul 2026 11:12:18 -0400 Subject: [PATCH 31/61] fix accessor & simplify currentChunkIndex --- .../internals/ChunkedByteBufferOutputStream.java | 10 ++++------ .../producer/internals/ChunkedProducerBatch.java | 4 ---- .../clients/producer/internals/ProducerBatch.java | 2 +- 3 files changed, 5 insertions(+), 11 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java index afbcd120e57cf..4056e53342cc4 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java @@ -184,13 +184,12 @@ public ByteBuffer buffer() { */ private ByteBuffer flatten() { // Written bytes only live in chunks up to currentChunk, later chunks are untouched. - int lastDataChunk = Math.min(currentChunkIndex, chunks.size() - 1); int totalSize = 0; - for (int i = 0; i <= lastDataChunk; i++) { + for (int i = 0; i <= currentChunkIndex; i++) { totalSize += chunks.get(i).position(); } ByteBuffer flattened = ByteBuffer.allocate(totalSize); - for (int i = 0; i <= lastDataChunk; i++) { + for (int i = 0; i <= currentChunkIndex; i++) { ByteBuffer chunk = chunks.get(i); int chunkPos = chunk.position(); chunk.flip(); @@ -232,10 +231,9 @@ private void releaseUnusedChunks() { @Override public int position() { ensureNotDeallocated(); - // Written bytes only live in chunks up to currentChunk; later chunks are untouched. - int lastDataChunk = Math.min(currentChunkIndex, chunks.size() - 1); + // Written bytes only live in chunks up to currentChunk, later chunks are untouched. int total = 0; - for (int i = 0; i <= lastDataChunk; i++) { + for (int i = 0; i <= currentChunkIndex; i++) { total += chunks.get(i).position(); } return total; diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedProducerBatch.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedProducerBatch.java index d273afbae1edb..33a3b20268154 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedProducerBatch.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedProducerBatch.java @@ -35,12 +35,8 @@ */ public class ChunkedProducerBatch extends ProducerBatch { - // The parent's builder reference is private; keep our own to access stream capacity state. - private final MemoryRecordsBuilder recordsBuilder; - public ChunkedProducerBatch(TopicPartition tp, MemoryRecordsBuilder recordsBuilder, long createdMs) { super(tp, recordsBuilder, createdMs); - this.recordsBuilder = recordsBuilder; } /** diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ProducerBatch.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ProducerBatch.java index 241f5e25f2e84..bad4c81ccc545 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ProducerBatch.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ProducerBatch.java @@ -68,7 +68,7 @@ private enum FinalState { ABORTED, FAILED, SUCCEEDED } final ProduceRequestResult produceFuture; private final List thunks = new ArrayList<>(); - private final MemoryRecordsBuilder recordsBuilder; + protected final MemoryRecordsBuilder recordsBuilder; private final AtomicInteger attempts = new AtomicInteger(0); private final boolean isSplitBatch; private final AtomicReference finalState = new AtomicReference<>(null); From 210d0a4b17137e069f133bc603b8462d10e64bfb Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Wed, 15 Jul 2026 13:24:39 -0400 Subject: [PATCH 32/61] remove dup math on numChunks --- .../clients/producer/internals/ChunkedBufferPool.java | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java index ab7abf489fdff..a1961a76d3746 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java @@ -59,11 +59,11 @@ public ChunkedBufferPool(long memory, int chunkSize, Metrics metrics, Time time, public List allocateChunks(int totalSize, long maxTimeToBlockMs) throws InterruptedException { if (totalSize <= 0) throw new IllegalArgumentException("totalSize must be positive: " + totalSize); - throwIfChunksNeededExceedsPool(totalSize); int chunkSize = poolableSize(); int numChunks = (int) (((long) totalSize + chunkSize - 1L) / chunkSize); long memoryRequired = (long) numChunks * chunkSize; + throwIfChunksNeededExceedsPool(totalSize, numChunks, chunkSize, memoryRequired); // Chunks taken from the free list. The remaining bytes are reserved against // nonPooledAvailableMemory and materialized as raw allocations after the lock is released. @@ -190,12 +190,9 @@ public List allocateChunks(int totalSize, long maxTimeToBlockMs) thr } /** - * Throw if rounding {@code totalSize} up to whole chunks would exceed the pool. + * Throw if the request memory rounded up to whole chunks would exceed the pool. */ - private void throwIfChunksNeededExceedsPool(int totalSize) { - int chunkSize = poolableSize(); - int numChunks = (int) (((long) totalSize + chunkSize - 1L) / chunkSize); - long memoryRequired = (long) numChunks * chunkSize; + private void throwIfChunksNeededExceedsPool(int totalSize, int numChunks, int chunkSize, long memoryRequired) { if (memoryRequired > totalMemory()) throw new IllegalArgumentException("Attempt to allocate " + totalSize + " bytes (" + numChunks + " chunks of " + chunkSize + " = " + memoryRequired + " bytes), but the " From 58a0bfcfceef3fb01c303dd285d1913830873df2 Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Wed, 15 Jul 2026 13:32:24 -0400 Subject: [PATCH 33/61] consolidate deallocateExtensionChunks --- .../internals/ChunkedRecordAccumulator.java | 25 +++++++++++-------- 1 file changed, 14 insertions(+), 11 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java index 2c8664128e067..ed3250de2000e 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java @@ -208,11 +208,8 @@ public RecordAppendResult append(String topic, // were sized against the previous partition's open batch, so they must not be // attached to a different partition's open batch — refund them and let the next // iteration re-check the new partition from scratch. - if (extensionChunks != null) { - for (ByteBuffer chunk : extensionChunks) - chunkedFree.deallocate(chunk); - extensionChunks = null; - } + deallocateExtensionChunks(extensionChunks); + extensionChunks = null; continue; } @@ -238,8 +235,7 @@ public RecordAppendResult append(String topic, continue; } // The open batch is gone, closed, or non-chunked (e.g., a split batch). Return chunks to pool. - for (ByteBuffer chunk : extensionChunks) - chunkedFree.deallocate(chunk); + deallocateExtensionChunks(extensionChunks); extensionChunks = null; continue; } @@ -274,14 +270,21 @@ public RecordAppendResult append(String topic, } finally { if (bufferStream != null) bufferStream.deallocate(); - if (extensionChunks != null) { - for (ByteBuffer chunk : extensionChunks) - chunkedFree.deallocate(chunk); - } + deallocateExtensionChunks(extensionChunks); appendsInProgress.decrementAndGet(); } } + /** + * Return any held extension chunks to the pool. No-op when none are held. + */ + private void deallocateExtensionChunks(List extensionChunks) { + if (extensionChunks == null) + return; + for (ByteBuffer chunk : extensionChunks) + chunkedFree.deallocate(chunk); + } + /** * Try to append to a ProducerBatch, with mid-batch chunk extension support. *

From e8c78a57e40bf9110ef0555692956b4d995fd989 Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Wed, 15 Jul 2026 13:44:07 -0400 Subject: [PATCH 34/61] clarify over-allocation on concurrent appends --- .../producer/internals/ChunkedRecordAccumulator.java | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java index ed3250de2000e..4764db664cf8f 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java @@ -220,6 +220,14 @@ public RecordAppendResult append(String topic, // ProducerBatch), which can't take extension chunks. Only attach to a // writable chunked batch; otherwise refund the chunks and re-evaluate. if (last instanceof ChunkedProducerBatch && last.isWritable()) { + // Attach the chunks we sized off-lock without re-checking whether a + // concurrent appender already grew the batch enough to fit our record. + // This optimizes for the common (uncontended) case. Under concurrency it + // can temporarily over-allocate: two appenders each sizing their own gap + // against the same observed capacity both attach, so the batch ends up + // with more capacity than the records need (one appender's chunks could + // have sufficed for both). The surplus is fully unused and returned to + // the pool when the batch is closed for appends. ((ChunkedProducerBatch) last).addBuffers(extensionChunks); extensionChunks = null; RecordAppendResult retryResult = tryAppend(timestamp, key, value, headers, callbacks, dq, nowMs); From 1d8bdd3175413b9acc489595fb84b495d4cc3b5a Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Wed, 15 Jul 2026 14:12:43 -0400 Subject: [PATCH 35/61] refactor append result --- .../internals/ChunkedRecordAccumulator.java | 19 ++-- .../producer/internals/RecordAccumulator.java | 90 +++++++++++-------- .../clients/producer/KafkaProducerTest.java | 2 +- 3 files changed, 61 insertions(+), 50 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java index 4764db664cf8f..3f938eaffa056 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java @@ -151,7 +151,7 @@ public RecordAppendResult append(String topic, } } - if (appendResult.needsBufferExtension) { + if (appendResult.needsBufferExtension()) { // Mid-batch extension: the open batch can still take this record so grow it in // place. The acquire is non-blocking to fail fast when the pool is exhausted: // close the batch and let the record block once on the new-batch path. @@ -176,7 +176,7 @@ public RecordAppendResult append(String topic, continue; } nowMs = time.milliseconds(); - } else if (appendResult.needsNewBatch && bufferStream == null) { + } else if (appendResult.needsNewBatch() && bufferStream == null) { // The open batch is done (e.g., full, closed) so start a new one. // Block on the pool for enough chunks to fit this record, sized with the // same cumulative estimator used mid-batch (header + record bytes for NONE, @@ -221,13 +221,10 @@ public RecordAppendResult append(String topic, // writable chunked batch; otherwise refund the chunks and re-evaluate. if (last instanceof ChunkedProducerBatch && last.isWritable()) { // Attach the chunks we sized off-lock without re-checking whether a - // concurrent appender already grew the batch enough to fit our record. - // This optimizes for the common (uncontended) case. Under concurrency it - // can temporarily over-allocate: two appenders each sizing their own gap - // against the same observed capacity both attach, so the batch ends up - // with more capacity than the records need (one appender's chunks could - // have sufficed for both). The surplus is fully unused and returned to - // the pool when the batch is closed for appends. + // concurrent appender already grew the batch enough. Optimizes for the + // uncontended case; under concurrency two appenders can each attach their + // own gap (temporary over-allocation, one could have been enough for both). + // The unused chunks returns to the pool when the batch closes for appends. ((ChunkedProducerBatch) last).addBuffers(extensionChunks); extensionChunks = null; RecordAppendResult retryResult = tryAppend(timestamp, key, value, headers, callbacks, dq, nowMs); @@ -257,7 +254,7 @@ public RecordAppendResult append(String topic, final ChunkedByteBufferOutputStream batchStream = bufferStream; appendResult = appendNewBatch(topic, effectivePartition, dq, timestamp, key, value, headers, callbacks, () -> chunkedRecordsBuilder(batchStream, firstRecordSize), nowMs); - if (appendResult.needsBufferExtension) { + if (appendResult.needsBufferExtension()) { // A concurrent appender created an open batch we should extend rather // than start a new one (detected by appendNewBatch's in-lock tryAppend). // Our bufferStream was sized for a fresh batch — release it and loop so @@ -266,7 +263,7 @@ public RecordAppendResult append(String topic, bufferStream = null; continue; } - if (appendResult.needsNewBatch) + if (appendResult.needsNewBatch()) throw new IllegalStateException("appendNewBatch must not return a needsNewBatch result"); if (appendResult.newBatchCreated) bufferStream = null; diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/RecordAccumulator.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/RecordAccumulator.java index 465bb0bf56b0d..72a7b17229dc7 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/RecordAccumulator.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/RecordAccumulator.java @@ -402,7 +402,7 @@ protected RecordAppendResult appendNewBatch(String topic, assert partition != RecordMetadata.UNKNOWN_PARTITION; RecordAppendResult appendResult = tryAppend(timestamp, key, value, headers, callbacks, dq, nowMs); - if (!appendResult.needsNewBatch) { + if (!appendResult.needsNewBatch()) { // Propagate without creating a new batch: either another thread already made us a batch // (success), or — incremental strategy — a concurrent appender created an extendable open batch // (needsBufferExtension), so the caller releases its pre-allocated buffer and retries via @@ -418,7 +418,7 @@ protected RecordAppendResult appendNewBatch(String topic, dq.addLast(batch); incomplete.add(batch); - return new RecordAppendResult(future, dq.size() > 1 || batch.isFull(), true, batch.estimatedSizeInBytes()); + return RecordAppendResult.appended(future, dq.size() > 1 || batch.isFull(), true, batch.estimatedSizeInBytes()); } /** @@ -458,7 +458,7 @@ protected RecordAppendResult tryAppend(long timestamp, byte[] key, byte[] value, last.closeForRecordAppends(); } else { int appendedBytes = last.estimatedSizeInBytes() - initialBytes; - return new RecordAppendResult(future, deque.size() > 1 || last.isFull(), false, appendedBytes); + return RecordAppendResult.appended(future, deque.size() > 1 || last.isFull(), false, appendedBytes); } } return RecordAppendResult.NEEDS_NEW_BATCH; @@ -1270,68 +1270,82 @@ public PartitionerConfig() { } /* - * Result of an attempt to append a record to the accumulator. Exactly one of three states: - * the record was appended ({@link RecordAppendResult#appended()}, {@code future} is set), the - * open batch needs more chunk capacity first ({@code needsBufferExtension}), or a new batch - * must be created for the record ({@code needsNewBatch}). + * Result of an attempt to append a record to the accumulator. Carries exactly one of three + * mutually-exclusive outcomes: the record was appended ({@link RecordAppendResult#appended()}, {@code future} is set), + * the open batch needs more chunk capacity first ({@link RecordAppendResult#needsBufferExtension()}), + * or a new batch must be created for the record ({@link RecordAppendResult#needsNewBatch()}). */ public static final class RecordAppendResult { + /** + * The three mutually-exclusive outcomes of an append attempt. + */ + public enum Outcome { APPENDED, NEEDS_BUFFER_EXTENSION, NEEDS_NEW_BATCH } + + public final Outcome outcome; public final FutureRecordMetadata future; public final boolean batchIsFull; public final boolean newBatchCreated; /** - * Signal (incremental strategy) that the open batch is within its batch-size limit but its - * chunks lack capacity. The append was NOT attempted; the caller allocates - * {@link #extensionBytesNeeded} bytes of chunk capacity, attaches them via - * {@link ChunkedProducerBatch#addBuffers}, and retries. When {@code true}, {@code future} is null. + * Bytes of chunk capacity the open batch needs before the record fits (incremental + * strategy). Meaningful only for {@link Outcome#NEEDS_BUFFER_EXTENSION}: the append was NOT + * attempted ({@code future} is null); the caller allocates this many bytes, attaches them + * via {@link ChunkedProducerBatch#addBuffers}, and retries. */ - public final boolean needsBufferExtension; public final int extensionBytesNeeded; - /** - * Signal that the record was not appended because there is no open batch that can take it - * (the batch is full, closed, or absent). The caller must create a new batch for the - * record. When {@code true}, {@code future} is null. - */ - public final boolean needsNewBatch; public final int appendedBytes; - /** The shared signal-only result for {@link #needsNewBatch}; carries no per-append state. */ - static final RecordAppendResult NEEDS_NEW_BATCH = new RecordAppendResult(null, false, false, false, 0, true, 0); - - public RecordAppendResult(FutureRecordMetadata future, - boolean batchIsFull, - boolean newBatchCreated, - int appendedBytes) { - this(future, batchIsFull, newBatchCreated, false, 0, false, appendedBytes); - } + /** The shared signal-only result for {@link Outcome#NEEDS_NEW_BATCH}; carries no per-append state. */ + public static final RecordAppendResult NEEDS_NEW_BATCH = + new RecordAppendResult(Outcome.NEEDS_NEW_BATCH, null, false, false, 0, 0); - private RecordAppendResult(FutureRecordMetadata future, + private RecordAppendResult(Outcome outcome, + FutureRecordMetadata future, boolean batchIsFull, boolean newBatchCreated, - boolean needsBufferExtension, int extensionBytesNeeded, - boolean needsNewBatch, int appendedBytes) { + this.outcome = outcome; this.future = future; this.batchIsFull = batchIsFull; this.newBatchCreated = newBatchCreated; - this.needsBufferExtension = needsBufferExtension; this.extensionBytesNeeded = extensionBytesNeeded; - this.needsNewBatch = needsNewBatch; this.appendedBytes = appendedBytes; } - /** A signal-only result indicating the caller must allocate more chunk capacity. */ - static RecordAppendResult needsExtension(int extensionBytesNeeded) { - return new RecordAppendResult(null, false, false, true, extensionBytesNeeded, false, 0); + /** Result of a successful append to an open batch */ + public static RecordAppendResult appended(FutureRecordMetadata future, + boolean batchIsFull, + boolean newBatchCreated, + int appendedBytes) { + Objects.requireNonNull(future, "future must be non-null for an appended result"); + return new RecordAppendResult(Outcome.APPENDED, future, batchIsFull, newBatchCreated, 0, appendedBytes); } /** - * @return {@code true} if the record was appended to the open batch ({@link #future} is then non-null), - * {@code false} if it wasn't ({@link #needsBufferExtension} and {@link #needsNewBatch} results). + * Signal-only result (incremental strategy): the open batch is within its batch-size limit + * but its chunks lack capacity, so the caller must allocate {@code extensionBytesNeeded} + * bytes of chunk capacity and retry. The append was not attempted. + */ + public static RecordAppendResult needsExtension(int extensionBytesNeeded) { + return new RecordAppendResult(Outcome.NEEDS_BUFFER_EXTENSION, null, false, false, extensionBytesNeeded, 0); + } + + /** + * @return {@code true} if the record was appended to the open batch ({@link #future} is then + * non-null), {@code false} for the {@link #needsBufferExtension()} and {@link #needsNewBatch()} results. */ public boolean appended() { - return !needsBufferExtension && !needsNewBatch; + return outcome == Outcome.APPENDED; + } + + /** @return {@code true} if the open batch needs more chunk capacity before the record fits. */ + public boolean needsBufferExtension() { + return outcome == Outcome.NEEDS_BUFFER_EXTENSION; + } + + /** @return {@code true} if there is no open batch that can take the record, so a new one must be created. */ + public boolean needsNewBatch() { + return outcome == Outcome.NEEDS_NEW_BATCH; } } diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/KafkaProducerTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/KafkaProducerTest.java index 628caabb0cfb9..e89b767b6f3a0 100644 --- a/clients/src/test/java/org/apache/kafka/clients/producer/KafkaProducerTest.java +++ b/clients/src/test/java/org/apache/kafka/clients/producer/KafkaProducerTest.java @@ -2953,7 +2953,7 @@ private FutureRecordMetadata expectAppend( RecordAccumulator.AppendCallbacks callbacks = (RecordAccumulator.AppendCallbacks) invocation.getArguments()[6]; callbacks.setPartition(initialSelectedPartition.partition()); - return new RecordAccumulator.RecordAppendResult( + return RecordAccumulator.RecordAppendResult.appended( futureRecordMetadata, false, false, From 4f3543a733015a73a8a422381d2a11feebdacb6c Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Wed, 15 Jul 2026 14:53:51 -0400 Subject: [PATCH 36/61] docs --- .../producer/internals/ChunkedRecordAccumulator.java | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java index 3f938eaffa056..8b92031efbdf2 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java @@ -321,9 +321,11 @@ protected ProducerBatch createProducerBatch(TopicPartition tp, MemoryRecordsBuil } /** - * Build a {@link MemoryRecordsBuilder} backed by the chunked stream. Its {@code writeLimit} - * bounds when the batch is full (the same batch-size limit as the full path), while chunk - * capacity grows on demand via {@link ChunkedByteBufferOutputStream#addBuffers(List)}. + * Build a {@link MemoryRecordsBuilder} backed by the chunked stream. + * + * @param bufferStream the chunked stream backing the batch + * @param firstRecordSize the first record's uncompressed size upper bound. Used to set the + * builder's write limit used by {@code hasRoomFor}/{@code isFull} */ private MemoryRecordsBuilder chunkedRecordsBuilder(ChunkedByteBufferOutputStream bufferStream, int firstRecordSize) { From 2136f0bc9640666a1bcf06ae780720de91d7f66c Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Wed, 15 Jul 2026 16:10:28 -0400 Subject: [PATCH 37/61] fix package name --- .../producer/internals/ChunkedByteBufferOutputStream.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java index 4056e53342cc4..a9fc1e2757d65 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java @@ -16,7 +16,7 @@ */ package org.apache.kafka.clients.producer.internals; -import org.apache.kafka.common.utils.ByteBufferOutputStream; +import org.apache.kafka.common.utils.internals.ByteBufferOutputStream; import java.nio.ByteBuffer; import java.util.ArrayList; From 9eced7d1d1a1e06357d5ad8b6e0a3ceaf7c02048 Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Mon, 20 Jul 2026 11:27:11 -0400 Subject: [PATCH 38/61] consolidate BufferPool --- checkstyle/suppressions.xml | 2 +- .../kafka/clients/producer/KafkaProducer.java | 3 +- .../producer/internals/BufferPool.java | 204 ++++++++++++++++-- .../producer/internals/ChunkedBufferPool.java | 201 ----------------- .../internals/ChunkedRecordAccumulator.java | 12 +- ...ava => BufferPoolChunkAllocationTest.java} | 32 +-- .../ChunkedByteBufferOutputStreamTest.java | 26 +-- .../ChunkedRecordAccumulatorTest.java | 18 +- 8 files changed, 230 insertions(+), 268 deletions(-) delete mode 100644 clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java rename clients/src/test/java/org/apache/kafka/clients/producer/internals/{ChunkedBufferPoolTest.java => BufferPoolChunkAllocationTest.java} (94%) diff --git a/checkstyle/suppressions.xml b/checkstyle/suppressions.xml index 1d4cba639c3a2..425ed76accc29 100644 --- a/checkstyle/suppressions.xml +++ b/checkstyle/suppressions.xml @@ -71,7 +71,7 @@ files="(Utils|Topic|Lz4BlockOutputStream|JoinGroupRequest).java"/> + files="(AbstractFetch|BufferPool|ClientTelemetryReporter|ConsumerCoordinator|CommitRequestManager|FetchCollector|OffsetFetcherUtils|KafkaProducer|Sender|ConfigDef|KerberosLogin|AbstractRequest|AbstractResponse|Selector|SslFactory|SslTransportLayer|SaslClientAuthenticator|SaslClientCallbackHandler|SaslServerAuthenticator|AbstractCoordinator|TransactionManager|AbstractStickyAssignor|DefaultSslEngineFactory|Authorizer|RecordAccumulator|MemoryRecords|FetchSessionHandler|MockAdminClient).java"/> diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/KafkaProducer.java b/clients/src/main/java/org/apache/kafka/clients/producer/KafkaProducer.java index d42d86856a436..a64eedb94e18c 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/KafkaProducer.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/KafkaProducer.java @@ -29,7 +29,6 @@ import org.apache.kafka.clients.consumer.OffsetCommitCallback; import org.apache.kafka.clients.producer.internals.BufferPool; import org.apache.kafka.clients.producer.internals.BuiltInPartitioner; -import org.apache.kafka.clients.producer.internals.ChunkedBufferPool; import org.apache.kafka.clients.producer.internals.ChunkedRecordAccumulator; import org.apache.kafka.clients.producer.internals.KafkaProducerMetrics; import org.apache.kafka.clients.producer.internals.ProducerInterceptors; @@ -488,7 +487,7 @@ public KafkaProducer(Properties properties, Serializer keySerializer, Seriali PRODUCER_METRIC_GROUP_NAME, time, transactionManager, - new ChunkedBufferPool(this.totalMemorySize, ChunkedRecordAccumulator.CHUNK_SIZE, metrics, time, PRODUCER_METRIC_GROUP_NAME)); + new BufferPool(this.totalMemorySize, ChunkedRecordAccumulator.CHUNK_SIZE, metrics, time, PRODUCER_METRIC_GROUP_NAME)); } else { this.accumulator = new RecordAccumulator(logContext, batchSize, diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/BufferPool.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/BufferPool.java index 9bd36c9410ff0..c2b2a1ec840a9 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/BufferPool.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/BufferPool.java @@ -26,10 +26,13 @@ import java.nio.ByteBuffer; import java.util.ArrayDeque; +import java.util.ArrayList; import java.util.Deque; +import java.util.List; import java.util.concurrent.TimeUnit; import java.util.concurrent.locks.Condition; import java.util.concurrent.locks.ReentrantLock; +import java.util.function.Supplier; /** @@ -145,28 +148,10 @@ public ByteBuffer allocate(int size, long maxTimeToBlockMs) throws InterruptedEx // loop over and over until we have a buffer or have reserved // enough memory to allocate one while (accumulated < size) { - long startWaitNs = time.nanoseconds(); - long timeNs; - boolean waitingTimeElapsed; - try { - waitingTimeElapsed = !moreMemory.await(remainingTimeToBlockNs, TimeUnit.NANOSECONDS); - } finally { - long endWaitNs = time.nanoseconds(); - timeNs = Math.max(0L, endWaitNs - startWaitNs); - recordWaitTime(timeNs); - } - - if (this.closed) - throw new KafkaException("Producer closed while allocating memory"); - - if (waitingTimeElapsed) { - recordBufferExhausted(); - throw new BufferExhaustedException("Failed to allocate " + size + " bytes within the configured max blocking time " + remainingTimeToBlockNs -= awaitMemory(moreMemory, remainingTimeToBlockNs, true, + () -> "Failed to allocate " + size + " bytes within the configured max blocking time " + maxTimeToBlockMs + " ms. Total memory: " + totalMemory() + " bytes. Available memory: " + availableMemory() + " bytes. Poolable size: " + poolableSize() + " bytes"); - } - - remainingTimeToBlockNs -= timeNs; // check if we can satisfy this request from the free list, // otherwise allocate memory @@ -208,6 +193,185 @@ public ByteBuffer allocate(int size, long maxTimeToBlockMs) throws InterruptedEx return buffer; } + /** + * Block once on {@code moreMemory} for up to {@code remainingTimeToBlockNs}, recording the wait + * time. Shared by {@link #allocate} and {@link #allocateChunks}. + * + * @return the nanos actually waited, for the caller to deduct from its blocking budget + * @throws KafkaException if the pool was closed during the wait + * @throws BufferExhaustedException if the wait timed out; the buffer-exhausted metric is recorded + * only when {@code recordExhaustedOnTimeout} is true (the incremental extension path + * recovers without dropping the record, so it records the drop itself if needed) + */ + private long awaitMemory(Condition moreMemory, long remainingTimeToBlockNs, + boolean recordExhaustedOnTimeout, Supplier exhaustedMessage) throws InterruptedException { + long startWaitNs = time.nanoseconds(); + long timeNs; + boolean waitingTimeElapsed; + try { + waitingTimeElapsed = !moreMemory.await(remainingTimeToBlockNs, TimeUnit.NANOSECONDS); + } finally { + long endWaitNs = time.nanoseconds(); + timeNs = Math.max(0L, endWaitNs - startWaitNs); + recordWaitTime(timeNs); + } + + if (this.closed) + throw new KafkaException("Producer closed while allocating memory"); + + if (waitingTimeElapsed) { + if (recordExhaustedOnTimeout) + recordBufferExhausted(); + throw new BufferExhaustedException(exhaustedMessage.get()); + } + + return timeNs; + } + + /** + * Allocate {@code ceil(totalSize / poolableSize)} poolable-sized buffers atomically, mirroring + * {@link #allocate}: satisfied immediately if memory is available, else blocks up to + * {@code maxTimeToBlockMs} for the whole request (FIFO on {@link #waiters}). The reservation is + * tracked as bytes against {@link #nonPooledAvailableMemory} plus chunks polled from {@link #free}. + * Any failure refunds the whole reservation and signals the next waiter before the exception + * propagates, so a failed request leaves nothing reserved. + *

+ * Used by the incremental buffer.memory allocation strategy; the poolable size is the chunk size. + * + * @param totalSize minimum total bytes of capacity required across the returned chunks + * @param maxTimeToBlockMs maximum time in milliseconds to block waiting for memory + * @return list of {@code ceil(totalSize / poolableSize())} {@code ByteBuffer}s, each of capacity + * {@code poolableSize()} + * @throws InterruptedException if interrupted while waiting + * @throws IllegalArgumentException if {@code totalSize <= 0}, or if the request rounded up to + * whole chunks exceeds {@code totalMemory()} + * @throws BufferExhaustedException if the request can't be satisfied within {@code maxTimeToBlockMs} + * @throws KafkaException if the pool is closed during the wait + */ + public List allocateChunks(int totalSize, long maxTimeToBlockMs) throws InterruptedException { + if (totalSize <= 0) + throw new IllegalArgumentException("totalSize must be positive: " + totalSize); + + int chunkSize = poolableSize(); + int numChunks = (int) (((long) totalSize + chunkSize - 1L) / chunkSize); + long memoryRequired = (long) numChunks * chunkSize; + throwIfChunksNeededExceedsPool(totalSize, numChunks, chunkSize, memoryRequired); + + // Chunks taken from the free list. The remaining bytes are reserved against + // nonPooledAvailableMemory and materialized as raw allocations after the lock is released. + List pooled = new ArrayList<>(numChunks); + + lock.lock(); + if (this.closed) { + lock.unlock(); + throw new KafkaException("Producer closed while allocating memory"); + } + try { + long freeListBytes = (long) free.size() * chunkSize; + if (this.nonPooledAvailableMemory + freeListBytes >= memoryRequired) { + // Enough memory available to allocate the chunks needed + while (pooled.size() < numChunks && !free.isEmpty()) + pooled.add(free.pollFirst()); + long remainingBytes = memoryRequired - (long) pooled.size() * chunkSize; + if (remainingBytes > 0) { + // remainingBytes > 0 means the free list was fully drained into `pooled`, so the + // remainder comes entirely from non-pooled memory (sufficient per the check above). + this.nonPooledAvailableMemory -= remainingBytes; + } + } else { + // Not enough memory available, so we wait to acquire the memory needed for all the chunks. + // Same as allocate, but for the whole multi-chunk request. A single Condition is added to + // the waiter's list to ensure FIFO fairness at the request level. + // + // `nonPoolAccumulated` tracks bytes drawn from nonPooledAvailableMemory only (pool chunks + // already taken live in `pooled`), and is always a whole-chunk multiple. If the wait + // does not complete (timeout / close / interrupt), the finally refunds the whole + // reservation: `nonPoolAccumulated` back to non-pooled memory, `pooled` back to the free chunks list. + long nonPoolAccumulated = 0; + boolean allocationCompleted = false; + Condition moreMemory = lock.newCondition(); + try { + long remainingTimeToBlockNs = TimeUnit.MILLISECONDS.toNanos(maxTimeToBlockMs); + waiters.addLast(moreMemory); + while ((long) pooled.size() * chunkSize + nonPoolAccumulated < memoryRequired) { + // Not recording the buffer-exhausted metric on timeout (recordExhaustedOnTimeout=false): + // this may be the extension path, which recovers without dropping the record, so the + // caller records the drop if needed. + remainingTimeToBlockNs -= awaitMemory(moreMemory, remainingTimeToBlockNs, false, + () -> "Failed to allocate " + memoryRequired + " bytes (" + numChunks + " chunks of " + + chunkSize + ") within the configured max blocking time " + maxTimeToBlockMs + + " ms. Total memory: " + totalMemory() + " bytes. Available memory: " + + availableMemory() + " bytes."); + + // Reuse free-list chunks first, preferring them over raw reservations: if a + // taken chunk covers a slot already reserved as raw bytes in an earlier + // iteration, hand that raw reservation back to the pool. + while (pooled.size() < numChunks && !free.isEmpty()) { + pooled.add(free.pollFirst()); + if (nonPoolAccumulated >= chunkSize) { // nonPoolAccumulated is always chunk-aligned + nonPoolAccumulated -= chunkSize; + this.nonPooledAvailableMemory += chunkSize; + } + } + // Reserve non-pooled memory for the still-uncovered chunks, in whole chunks + // (the buffers themselves are allocated after the lock is released). + while (pooled.size() + (int) (nonPoolAccumulated / chunkSize) < numChunks + && this.nonPooledAvailableMemory >= chunkSize) { + this.nonPooledAvailableMemory -= chunkSize; + nonPoolAccumulated += chunkSize; + } + } + allocationCompleted = true; + } finally { + if (!allocationCompleted) { + // Refund all that was reserved (pooled chunks and non-pooled bytes) + this.nonPooledAvailableMemory += nonPoolAccumulated; + for (ByteBuffer chunk : pooled) + free.addFirst(chunk); + pooled.clear(); + } + waiters.remove(moreMemory); + } + } + } finally { + try { + signalNextWaiterIfMemoryAvailable(); + } finally { + lock.unlock(); + } + } + + // Allocate raw chunks for the reserved non-pooled portion outside the lock. On error, + // refund all reserved bytes (memoryRequired) and let the next waiter try. + int chunksStillNeeded = numChunks - pooled.size(); + List result = new ArrayList<>(numChunks); + result.addAll(pooled); + boolean error = true; + try { + for (int i = 0; i < chunksStillNeeded; i++) + result.add(allocateByteBuffer(chunkSize)); + error = false; + return result; + } finally { + if (error) { + // The pooled buffers we already drained are also lost on this path; mirror + // safeAllocateByteBuffer's behaviour (the bytes return, the ByteBuffer + // instances become garbage). + releaseReservedBytes(memoryRequired); + } + } + } + + /** + * Throw if the request memory rounded up to whole chunks would exceed the pool. + */ + private void throwIfChunksNeededExceedsPool(int totalSize, int numChunks, int chunkSize, long memoryRequired) { + if (memoryRequired > totalMemory()) + throw new IllegalArgumentException("Attempt to allocate " + totalSize + " bytes (" + + numChunks + " chunks of " + chunkSize + " = " + memoryRequired + " bytes), but the " + + "hard limit on memory allocations is " + totalMemory() + "."); + } + // Protected for testing protected void recordWaitTime(long timeNs) { this.waitTime.record(timeNs, time.milliseconds()); diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java deleted file mode 100644 index a1961a76d3746..0000000000000 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java +++ /dev/null @@ -1,201 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one or more - * contributor license agreements. See the NOTICE file distributed with - * this work for additional information regarding copyright ownership. - * The ASF licenses this file to You under the Apache License, Version 2.0 - * (the "License"); you may not use this file except in compliance with - * the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package org.apache.kafka.clients.producer.internals; - -import org.apache.kafka.clients.producer.BufferExhaustedException; -import org.apache.kafka.common.KafkaException; -import org.apache.kafka.common.metrics.Metrics; -import org.apache.kafka.common.utils.Time; - -import java.nio.ByteBuffer; -import java.util.ArrayList; -import java.util.List; -import java.util.concurrent.TimeUnit; -import java.util.concurrent.locks.Condition; - -/** - * A {@link BufferPool} dedicated to chunk-sized buffer reuse (chunk size = {@link #poolableSize()}). - *

- * Adds {@link #allocateChunks(int, long)} to acquire multiple chunks atomically. - */ -public class ChunkedBufferPool extends BufferPool { - - public ChunkedBufferPool(long memory, int chunkSize, Metrics metrics, Time time, String metricGrpName) { - super(memory, chunkSize, metrics, time, metricGrpName); - } - - /** - * Allocate {@code ceil(totalSize / chunkSize)} chunk-sized buffers atomically, mirroring - * {@link BufferPool#allocate}: satisfied immediately if memory is available, else blocks up to - * {@code maxTimeToBlockMs} for the whole request (FIFO on {@link #waiters}). - * The reservation is tracked as bytes against {@link #nonPooledAvailableMemory} plus chunks polled - * from {@link #free}. Any failure refunds the whole reservation and signals the next waiter - * before the exception propagates, so a failed request leaves nothing reserved. - * - * @param totalSize minimum total bytes of capacity required across the returned chunks - * @param maxTimeToBlockMs maximum time in milliseconds to block waiting for memory - * @return list of {@code ceil(totalSize / chunkSize)} {@code ByteBuffer}s, each of capacity - * {@code chunkSize} - * @throws InterruptedException if interrupted while waiting - * @throws IllegalArgumentException if {@code totalSize <= 0}, or if the request rounded up to - * whole chunks exceeds {@code totalMemory()} - * @throws BufferExhaustedException if the request can't be satisfied within {@code maxTimeToBlockMs} - * @throws KafkaException if the pool is closed during the wait - */ - public List allocateChunks(int totalSize, long maxTimeToBlockMs) throws InterruptedException { - if (totalSize <= 0) - throw new IllegalArgumentException("totalSize must be positive: " + totalSize); - - int chunkSize = poolableSize(); - int numChunks = (int) (((long) totalSize + chunkSize - 1L) / chunkSize); - long memoryRequired = (long) numChunks * chunkSize; - throwIfChunksNeededExceedsPool(totalSize, numChunks, chunkSize, memoryRequired); - - // Chunks taken from the free list. The remaining bytes are reserved against - // nonPooledAvailableMemory and materialized as raw allocations after the lock is released. - List pooled = new ArrayList<>(numChunks); - - lock.lock(); - if (this.closed) { - lock.unlock(); - throw new KafkaException("Producer closed while allocating memory"); - } - try { - long freeListBytes = (long) free.size() * chunkSize; - if (this.nonPooledAvailableMemory + freeListBytes >= memoryRequired) { - // Enough memory available to allocate the chunks needed - while (pooled.size() < numChunks && !free.isEmpty()) - pooled.add(free.pollFirst()); - long remainingBytes = memoryRequired - (long) pooled.size() * chunkSize; - if (remainingBytes > 0) { - // remainingBytes > 0 means the free list was fully drained into `pooled`, so the - // remainder comes entirely from non-pooled memory (sufficient per the check above). - this.nonPooledAvailableMemory -= remainingBytes; - } - } else { - // Not enough memory available to allocate the chunks needed, so we need to wait for memory. - // Same as in BufferPool.allocate, but wait to acquire the memory needed for all the chunks. - // A single Condition is added to the waiter's list to ensure FIFO fairness at the request level. - // - // `nonPoolAccumulated` tracks bytes drawn from nonPooledAvailableMemory only (pool chunks - // already taken live in `pooled`), and is always a whole-chunk multiple. If the wait - // does not complete (timeout / close / interrupt), the finally refunds the whole - // reservation: `nonPoolAccumulated` back to non-pooled memory, `pooled` back to the free chunks list. - long nonPoolAccumulated = 0; - boolean allocationCompleted = false; - Condition moreMemory = lock.newCondition(); - try { - long remainingTimeToBlockNs = TimeUnit.MILLISECONDS.toNanos(maxTimeToBlockMs); - waiters.addLast(moreMemory); - while ((long) pooled.size() * chunkSize + nonPoolAccumulated < memoryRequired) { - long startWaitNs = time.nanoseconds(); - long timeNs; - boolean waitingTimeElapsed; - try { - waitingTimeElapsed = !moreMemory.await(remainingTimeToBlockNs, TimeUnit.NANOSECONDS); - } finally { - long endWaitNs = time.nanoseconds(); - timeNs = Math.max(0L, endWaitNs - startWaitNs); - recordWaitTime(timeNs); - } - - if (this.closed) - throw new KafkaException("Producer closed while allocating memory"); - - if (waitingTimeElapsed) { - // Intentionally not recording the buffer-exhausted metric here. - // This may be called on the extension path, which recovers without - // dropping the record, so the caller is the one recording - // the drop if needed. - throw new BufferExhaustedException("Failed to allocate " + memoryRequired - + " bytes (" + numChunks + " chunks of " + chunkSize - + ") within the configured max blocking time " + maxTimeToBlockMs - + " ms. Total memory: " + totalMemory() + " bytes. Available memory: " - + availableMemory() + " bytes."); - } - - remainingTimeToBlockNs -= timeNs; - - // Reuse free-list chunks first, preferring them over raw reservations: if a - // taken chunk covers a slot already reserved as raw bytes in an earlier - // iteration, hand that raw reservation back to the pool. - while (pooled.size() < numChunks && !free.isEmpty()) { - pooled.add(free.pollFirst()); - if (nonPoolAccumulated >= chunkSize) { // nonPoolAccumulated is always chunk-aligned - nonPoolAccumulated -= chunkSize; - this.nonPooledAvailableMemory += chunkSize; - } - } - // Reserve non-pooled memory for the still-uncovered chunks, in whole chunks - // (the buffers themselves are allocated after the lock is released). - while (pooled.size() + (int) (nonPoolAccumulated / chunkSize) < numChunks - && this.nonPooledAvailableMemory >= chunkSize) { - this.nonPooledAvailableMemory -= chunkSize; - nonPoolAccumulated += chunkSize; - } - } - allocationCompleted = true; - } finally { - if (!allocationCompleted) { - // Refund all that was reserved (pooled chunks and non-pooled bytes) - this.nonPooledAvailableMemory += nonPoolAccumulated; - for (ByteBuffer chunk : pooled) - free.addFirst(chunk); - pooled.clear(); - } - waiters.remove(moreMemory); - } - } - } finally { - try { - signalNextWaiterIfMemoryAvailable(); - } finally { - lock.unlock(); - } - } - - // Allocate raw chunks for the reserved non-pooled portion outside the lock. On error, - // refund all reserved bytes (memoryRequired) and let the next waiter try. - int chunksStillNeeded = numChunks - pooled.size(); - List result = new ArrayList<>(numChunks); - result.addAll(pooled); - boolean error = true; - try { - for (int i = 0; i < chunksStillNeeded; i++) - result.add(allocateByteBuffer(chunkSize)); - error = false; - return result; - } finally { - if (error) { - // The pooled buffers we already drained are also lost on this path; mirror - // BufferPool.safeAllocateByteBuffer's behaviour (the bytes return, the ByteBuffer - // instances become garbage). - releaseReservedBytes(memoryRequired); - } - } - } - - /** - * Throw if the request memory rounded up to whole chunks would exceed the pool. - */ - private void throwIfChunksNeededExceedsPool(int totalSize, int numChunks, int chunkSize, long memoryRequired) { - if (memoryRequired > totalMemory()) - throw new IllegalArgumentException("Attempt to allocate " + totalSize + " bytes (" - + numChunks + " chunks of " + chunkSize + " = " + memoryRequired + " bytes), but the " - + "hard limit on memory allocations is " + totalMemory() + "."); - } -} diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java index 8b92031efbdf2..d7e6b8ea59c05 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java @@ -42,7 +42,7 @@ /** * A {@link RecordAccumulator} variant that backs each batch with fixed-size chunks drawn from a - * {@link ChunkedBufferPool}, attaching more chunks on demand as records are appended instead of + * {@link BufferPool}, attaching more chunks on demand as records are appended instead of * reserving {@code batch.size} per batch up front. Buffered memory therefore scales with the data * actually written rather than with {@code active_partition_count × batch.size}. *

@@ -59,7 +59,7 @@ public class ChunkedRecordAccumulator extends RecordAccumulator { */ public static final int CHUNK_SIZE = 16 * 1024; - private final ChunkedBufferPool chunkedFree; + private final BufferPool chunkedFree; public ChunkedRecordAccumulator(LogContext logContext, int batchSize, @@ -73,7 +73,7 @@ public ChunkedRecordAccumulator(LogContext logContext, String metricGrpName, Time time, TransactionManager transactionManager, - ChunkedBufferPool bufferPool) { + BufferPool bufferPool) { super(logContext, batchSize, compression, lingerMs, retryBackoffMs, retryBackoffMaxMs, deliveryTimeoutMs, partitionerConfig, metrics, metricGrpName, time, transactionManager, bufferPool); // TODO: drop this once the incremental strategy supports compressed data (with the @@ -95,7 +95,7 @@ public ChunkedRecordAccumulator(LogContext logContext, String metricGrpName, Time time, TransactionManager transactionManager, - ChunkedBufferPool bufferPool) { + BufferPool bufferPool) { this(logContext, batchSize, compression, lingerMs, retryBackoffMs, retryBackoffMaxMs, deliveryTimeoutMs, new PartitionerConfig(), metrics, metricGrpName, time, transactionManager, bufferPool); @@ -254,6 +254,8 @@ public RecordAppendResult append(String topic, final ChunkedByteBufferOutputStream batchStream = bufferStream; appendResult = appendNewBatch(topic, effectivePartition, dq, timestamp, key, value, headers, callbacks, () -> chunkedRecordsBuilder(batchStream, firstRecordSize), nowMs); + if (appendResult.needsNewBatch()) + throw new IllegalStateException("appendNewBatch must not return a needsNewBatch result"); if (appendResult.needsBufferExtension()) { // A concurrent appender created an open batch we should extend rather // than start a new one (detected by appendNewBatch's in-lock tryAppend). @@ -263,8 +265,6 @@ public RecordAppendResult append(String topic, bufferStream = null; continue; } - if (appendResult.needsNewBatch()) - throw new IllegalStateException("appendNewBatch must not return a needsNewBatch result"); if (appendResult.newBatchCreated) bufferStream = null; boolean enableSwitch = allBatchesFull(dq); diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPoolTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/BufferPoolChunkAllocationTest.java similarity index 94% rename from clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPoolTest.java rename to clients/src/test/java/org/apache/kafka/clients/producer/internals/BufferPoolChunkAllocationTest.java index 5b0f6f9b4fbb4..c71688e9301a7 100644 --- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPoolTest.java +++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/BufferPoolChunkAllocationTest.java @@ -38,7 +38,7 @@ import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; -public class ChunkedBufferPoolTest { +public class BufferPoolChunkAllocationTest { private final MockTime time = new MockTime(); private final Metrics metrics = new Metrics(time); @@ -48,16 +48,16 @@ public void teardown() { metrics.close(); } - private ChunkedBufferPool pool(long totalMemory, int chunkSize) { + private BufferPool pool(long totalMemory, int chunkSize) { String metricGroup = "producer-metrics"; - return new ChunkedBufferPool(totalMemory, chunkSize, metrics, time, metricGroup); + return new BufferPool(totalMemory, chunkSize, metrics, time, metricGroup); } /** Single-chunk request returns a list of one buffer at the chunk size. */ @Test public void testAllocateOneChunk() throws Exception { int chunkSize = 64; - ChunkedBufferPool p = pool(1024, chunkSize); + BufferPool p = pool(1024, chunkSize); List chunks = p.allocateChunks(chunkSize, 100); assertEquals(1, chunks.size()); assertEquals(chunkSize, chunks.get(0).capacity()); @@ -67,7 +67,7 @@ public void testAllocateOneChunk() throws Exception { @Test public void testAllocateRoundsUpToChunkBoundary() throws Exception { int chunkSize = 64; - ChunkedBufferPool p = pool(1024, chunkSize); + BufferPool p = pool(1024, chunkSize); // 65 bytes requested → 2 chunks (128 bytes total). List chunks = p.allocateChunks(65, 100); assertEquals(2, chunks.size()); @@ -79,7 +79,7 @@ public void testAllocateRoundsUpToChunkBoundary() throws Exception { @Test public void testAllocateMultipleChunks() throws Exception { int chunkSize = 64; - ChunkedBufferPool p = pool(1024, chunkSize); + BufferPool p = pool(1024, chunkSize); List chunks = p.allocateChunks(4 * chunkSize, 100); assertEquals(4, chunks.size()); } @@ -89,7 +89,7 @@ public void testAllocateMultipleChunks() throws Exception { public void testAvailableMemoryAfterAllocation() throws Exception { int chunkSize = 64; long total = 256; - ChunkedBufferPool p = pool(total, chunkSize); + BufferPool p = pool(total, chunkSize); p.allocateChunks(3 * chunkSize, 100); assertEquals(total - 3L * chunkSize, p.availableMemory()); } @@ -99,7 +99,7 @@ public void testAvailableMemoryAfterAllocation() throws Exception { public void testDeallocationRestoresMemory() throws Exception { int chunkSize = 64; long total = 256; - ChunkedBufferPool p = pool(total, chunkSize); + BufferPool p = pool(total, chunkSize); List chunks = p.allocateChunks(2 * chunkSize, 100); for (ByteBuffer chunk : chunks) p.deallocate(chunk); @@ -108,13 +108,13 @@ public void testDeallocationRestoresMemory() throws Exception { @Test public void testRejectsRequestExceedingTotalMemory() { - ChunkedBufferPool p = pool(128, 64); + BufferPool p = pool(128, 64); assertThrows(IllegalArgumentException.class, () -> p.allocateChunks(129, 100)); } @Test public void testRejectsNonPositiveRequest() { - ChunkedBufferPool p = pool(128, 64); + BufferPool p = pool(128, 64); assertThrows(IllegalArgumentException.class, () -> p.allocateChunks(0, 100)); assertThrows(IllegalArgumentException.class, () -> p.allocateChunks(-1, 100)); } @@ -127,7 +127,7 @@ public void testRejectsNonPositiveRequest() { public void testRollbackOnPartialFailure() throws Exception { int chunkSize = 64; long total = 2 * chunkSize; // only 2 chunks worth of memory - ChunkedBufferPool p = pool(total, chunkSize); + BufferPool p = pool(total, chunkSize); // Reserve one chunk so the pool has only 1 left. ByteBuffer held = p.allocate(chunkSize, 100); @@ -150,7 +150,7 @@ public void testRollbackOnPartialFailure() throws Exception { public void testNoPartialHoldsDuringWait() throws Exception { int chunkSize = 64; long total = 4L * chunkSize; - ChunkedBufferPool p = pool(total, chunkSize); + BufferPool p = pool(total, chunkSize); // Drain the pool down to 1 chunk's worth so a 2-chunk request can't be satisfied immediately. ByteBuffer h1 = p.allocate(chunkSize, 100); ByteBuffer h2 = p.allocate(chunkSize, 100); @@ -193,7 +193,7 @@ public void testNoPartialHoldsDuringWait() throws Exception { public void testFifoFairnessAcrossChunkedAndSingleChunkRequests() throws Exception { int chunkSize = 64; long total = 4L * chunkSize; - ChunkedBufferPool p = pool(total, chunkSize); + BufferPool p = pool(total, chunkSize); // Drain the pool entirely so any new request must wait. ByteBuffer h1 = p.allocate(chunkSize, 100); ByteBuffer h2 = p.allocate(chunkSize, 100); @@ -267,7 +267,7 @@ public void testFifoFairnessAcrossChunkedAndSingleChunkRequests() throws Excepti public void testSlowPathDoesNotDoubleRefundPolledChunksOnTimeout() throws Exception { int chunkSize = 64; long total = 3L * chunkSize; - ChunkedBufferPool p = pool(total, chunkSize); + BufferPool p = pool(total, chunkSize); // Drain so the next allocateChunks must wait. ByteBuffer h1 = p.allocate(chunkSize, 100); @@ -322,7 +322,7 @@ public void testSlowPathDoesNotDoubleRefundPolledChunksOnTimeout() throws Except public void testSlowPathSuccessDoesNotCorruptAccounting() throws Exception { int chunkSize = 64; long total = 3L * chunkSize; - ChunkedBufferPool p = pool(total, chunkSize); + BufferPool p = pool(total, chunkSize); // Drain so a 2-chunk request must wait. ByteBuffer h1 = p.allocate(chunkSize, 100); @@ -371,7 +371,7 @@ public void testSlowPathSuccessDoesNotCorruptAccounting() throws Exception { public void testCloseDuringAtomicWait() throws Exception { int chunkSize = 64; long total = 2L * chunkSize; - ChunkedBufferPool p = pool(total, chunkSize); + BufferPool p = pool(total, chunkSize); // Drain so the next request must wait. ByteBuffer h1 = p.allocate(chunkSize, 100); ByteBuffer h2 = p.allocate(chunkSize, 100); diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java index 49b58a08ce6d0..441f6001d012a 100644 --- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java +++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java @@ -42,12 +42,12 @@ public void teardown() { metrics.close(); } - private ChunkedBufferPool pool(long total, int chunkSize) { + private BufferPool pool(long total, int chunkSize) { String metricGroup = "test"; - return new ChunkedBufferPool(total, chunkSize, metrics, time, metricGroup); + return new BufferPool(total, chunkSize, metrics, time, metricGroup); } - private List chunks(ChunkedBufferPool pool, int chunkSize, int count) throws InterruptedException { + private List chunks(BufferPool pool, int chunkSize, int count) throws InterruptedException { return pool.allocateChunks(chunkSize * count, 100); } @@ -55,7 +55,7 @@ private List chunks(ChunkedBufferPool pool, int chunkSize, int count @SuppressWarnings({"try", "resource"}) // Each constructor call throws, so no stream is ever created public void testConstructorRejectsInvalidChunks() { int chunkSize = 16; - ChunkedBufferPool p = pool(64, chunkSize); + BufferPool p = pool(64, chunkSize); assertThrows(IllegalArgumentException.class, () -> new ChunkedByteBufferOutputStream(null, chunkSize, p)); @@ -71,7 +71,7 @@ public void testConstructorRejectsInvalidChunks() { @Test public void testOperationsRejectedAfterDeallocate() throws Exception { int chunkSize = 8; - ChunkedBufferPool p = pool(64, chunkSize); + BufferPool p = pool(64, chunkSize); try (ChunkedByteBufferOutputStream stream = new ChunkedByteBufferOutputStream(chunks(p, chunkSize, 2), chunkSize, p)) { stream.write(new byte[]{1, 2, 3}, 0, 3); @@ -100,7 +100,7 @@ public void testOperationsRejectedAfterDeallocate() throws Exception { @Test public void testWritesDisallowedAfterBuffer() throws Exception { int chunkSize = 16; - ChunkedBufferPool p = pool(64, chunkSize); + BufferPool p = pool(64, chunkSize); try (ChunkedByteBufferOutputStream stream = new ChunkedByteBufferOutputStream(chunks(p, chunkSize, 2), chunkSize, p)) { stream.write(new byte[]{1, 2, 3}, 0, 3); @@ -120,7 +120,7 @@ public void testWritesDisallowedAfterBuffer() throws Exception { @Test public void testSingleChunkWriteRoundtrip() throws Exception { int chunkSize = 16; - ChunkedBufferPool p = pool(64, chunkSize); + BufferPool p = pool(64, chunkSize); try (ChunkedByteBufferOutputStream stream = new ChunkedByteBufferOutputStream(chunks(p, chunkSize, 1), chunkSize, p)) { byte[] payload = new byte[]{1, 2, 3, 4, 5}; @@ -139,7 +139,7 @@ public void testSingleChunkWriteRoundtrip() throws Exception { @Test public void testWriteAcrossMultipleChunks() throws Exception { int chunkSize = 8; - ChunkedBufferPool p = pool(64, chunkSize); + BufferPool p = pool(64, chunkSize); try (ChunkedByteBufferOutputStream stream = new ChunkedByteBufferOutputStream(chunks(p, chunkSize, 3), chunkSize, p)) { byte[] payload = new byte[20]; @@ -159,7 +159,7 @@ public void testWriteAcrossMultipleChunks() throws Exception { @Test public void testRemainingSumsFreeBytesAcrossChunks() throws Exception { int chunkSize = 8; - ChunkedBufferPool p = pool(64, chunkSize); + BufferPool p = pool(64, chunkSize); try (ChunkedByteBufferOutputStream stream = new ChunkedByteBufferOutputStream(chunks(p, chunkSize, 2), chunkSize, p)) { assertEquals(2 * chunkSize, stream.remaining()); @@ -175,7 +175,7 @@ public void testRemainingSumsFreeBytesAcrossChunks() throws Exception { @Test public void testAddBuffersExtendsStream() throws Exception { int chunkSize = 8; - ChunkedBufferPool p = pool(64, chunkSize); + BufferPool p = pool(64, chunkSize); try (ChunkedByteBufferOutputStream stream = new ChunkedByteBufferOutputStream(chunks(p, chunkSize, 1), chunkSize, p)) { // Fill the initial chunk. @@ -197,7 +197,7 @@ public void testAddBuffersExtendsStream() throws Exception { @Test public void testPositionWalksAcrossChunks() throws Exception { int chunkSize = 4; - ChunkedBufferPool p = pool(32, chunkSize); + BufferPool p = pool(32, chunkSize); try (ChunkedByteBufferOutputStream stream = new ChunkedByteBufferOutputStream(chunks(p, chunkSize, 3), chunkSize, p)) { stream.position(6); // straddles chunk 0 (4 bytes) and chunk 1 (2 bytes) @@ -212,7 +212,7 @@ public void testPositionWalksAcrossChunks() throws Exception { public void testDeallocateReturnsAllChunks() throws Exception { int chunkSize = 8; long total = 64; - ChunkedBufferPool p = pool(total, chunkSize); + BufferPool p = pool(total, chunkSize); List initial = chunks(p, chunkSize, 3); try (ChunkedByteBufferOutputStream stream = new ChunkedByteBufferOutputStream(initial, chunkSize, p)) { // Also attach one extra chunk so deallocate must handle the added-buffer case too. @@ -235,7 +235,7 @@ public void testDeallocateReturnsAllChunks() throws Exception { public void testUnusedChunksReleasedOnCloseNotOnBuffer() throws Exception { int chunkSize = 8; long total = 64; - ChunkedBufferPool p = pool(total, chunkSize); + BufferPool p = pool(total, chunkSize); ChunkedByteBufferOutputStream stream = new ChunkedByteBufferOutputStream(chunks(p, chunkSize, 3), chunkSize, p); // Write into the first chunk only; chunks 2 and 3 stay unused. byte[] payload = new byte[]{1, 2, 3}; diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java index 3d0570f0ecad4..b9e4d3ed29b09 100644 --- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java +++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java @@ -82,7 +82,7 @@ public void teardown() { } private ChunkedRecordAccumulator newAccumulator(int batchSize, int chunkSize, long totalMemory, Compression compression) { - ChunkedBufferPool pool = new ChunkedBufferPool(totalMemory, chunkSize, metrics, time, "producer-metrics"); + BufferPool pool = new BufferPool(totalMemory, chunkSize, metrics, time, "producer-metrics"); return new ChunkedRecordAccumulator(logContext, batchSize, compression, /* lingerMs */ 0, /* retryBackoffMs */ 0L, /* retryBackoffMaxMs */ 0L, /* deliveryTimeoutMs */ 3200, metrics, "producer-metrics", time, @@ -158,10 +158,10 @@ public void testMidBatchExtensionGrowsExistingBatch() throws Exception { * Pool that adds one append right after a chunk allocation returns, mocking a concurrent * appender racing the same batch. */ - private ChunkedBufferPool poolMockingConcurrentChunkAllocation(int chunkSize, long totalMemory, + private BufferPool poolMockingConcurrentChunkAllocation(int chunkSize, long totalMemory, AtomicReference injectAppendOnce, byte[] injectedValue) { - return new ChunkedBufferPool(totalMemory, chunkSize, metrics, time, "producer-metrics") { + return new BufferPool(totalMemory, chunkSize, metrics, time, "producer-metrics") { @Override public List allocateChunks(int totalSize, long maxTimeToBlockMs) throws InterruptedException { List chunks = super.allocateChunks(totalSize, maxTimeToBlockMs); @@ -190,7 +190,7 @@ public void testConcurrentExtensionRaceLoserExtendsAgain() throws Exception { final byte[] value = new byte[350]; // needs 2 chunks final AtomicReference injectAppendOnce = new AtomicReference<>(); - ChunkedBufferPool pool = poolMockingConcurrentChunkAllocation(chunkSize, 64L * chunkSize, injectAppendOnce, value); + BufferPool pool = poolMockingConcurrentChunkAllocation(chunkSize, 64L * chunkSize, injectAppendOnce, value); ChunkedRecordAccumulator accum = new ChunkedRecordAccumulator(logContext, 8192, Compression.NONE, /* lingerMs */ 0, /* retryBackoffMs */ 0L, /* retryBackoffMaxMs */ 0L, /* deliveryTimeoutMs */ 3200, metrics, "producer-metrics", time, @@ -234,7 +234,7 @@ public void testConcurrentExtensionRaceLoserStartsNewBatch() throws Exception { final AtomicReference injectAppendOnce = new AtomicReference<>(); final AtomicInteger chunkDeallocations = new AtomicInteger(); - ChunkedBufferPool pool = new ChunkedBufferPool(64L * chunkSize, chunkSize, metrics, time, "producer-metrics") { + BufferPool pool = new BufferPool(64L * chunkSize, chunkSize, metrics, time, "producer-metrics") { @Override public List allocateChunks(int totalSize, long maxTimeToBlockMs) throws InterruptedException { List chunks = super.allocateChunks(totalSize, maxTimeToBlockMs); @@ -299,7 +299,7 @@ public void testPartitionSwitchRefundsHeldExtensionChunks() throws Exception { AtomicBoolean extensionAllocated = new AtomicBoolean(false); AtomicInteger chunkDeallocations = new AtomicInteger(0); - ChunkedBufferPool pool = new ChunkedBufferPool(totalMemory, chunkSize, metrics, time, "producer-metrics") { + BufferPool pool = new BufferPool(totalMemory, chunkSize, metrics, time, "producer-metrics") { @Override public List allocateChunks(int totalSize, long maxTimeToBlockMs) throws InterruptedException { List chunks = super.allocateChunks(totalSize, maxTimeToBlockMs); @@ -364,7 +364,7 @@ protected boolean partitionChanged(String topic, TopicInfo topicInfo, public void testInflightExpirationReturnsAllChunksToPool() throws Exception { int chunkSize = 128; long totalMemory = 32L * chunkSize; - ChunkedBufferPool pool = new ChunkedBufferPool(totalMemory, chunkSize, metrics, time, "producer-metrics"); + BufferPool pool = new BufferPool(totalMemory, chunkSize, metrics, time, "producer-metrics"); ChunkedRecordAccumulator accum = new ChunkedRecordAccumulator(logContext, 8192, Compression.NONE, /* lingerMs */ 0, /* retryBackoffMs */ 0L, /* retryBackoffMaxMs */ 0L, /* deliveryTimeoutMs */ 3200, metrics, "producer-metrics", time, @@ -416,7 +416,7 @@ public void testExhaustedExtensionFallsBackToBlockingNewBatchPath() throws Excep AtomicInteger closeForAppendsCalls = new AtomicInteger(); List closeCallsAtAlloc = new ArrayList<>(); - ChunkedBufferPool pool = new ChunkedBufferPool(16L * chunkSize, chunkSize, metrics, time, "producer-metrics") { + BufferPool pool = new BufferPool(16L * chunkSize, chunkSize, metrics, time, "producer-metrics") { @Override public List allocateChunks(int totalSize, long maxTimeToBlockMs) throws InterruptedException { allocTimeouts.add(maxTimeToBlockMs); @@ -585,7 +585,7 @@ public void testCumulativeAccountsForBatchHeaderOnce() throws Exception { int chunkSize = 256; int batchSize = 8192; long totalMemory = 64L * chunkSize; - ChunkedBufferPool pool = new ChunkedBufferPool(totalMemory, chunkSize, metrics, time, "producer-metrics"); + BufferPool pool = new BufferPool(totalMemory, chunkSize, metrics, time, "producer-metrics"); ChunkedRecordAccumulator accum = new ChunkedRecordAccumulator(logContext, batchSize, Compression.NONE, /* lingerMs */ 0, /* retryBackoffMs */ 0L, /* retryBackoffMaxMs */ 0L, /* deliveryTimeoutMs */ 3200, metrics, "producer-metrics", time, From 7a5afd10089aaa0b7cedb2f1fbf5e51178e02deb Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Mon, 20 Jul 2026 12:18:14 -0400 Subject: [PATCH 39/61] addressing comments --- .../ChunkedByteBufferOutputStream.java | 31 ++++++++++++------- .../producer/internals/RecordAccumulator.java | 16 +++++----- .../ChunkedByteBufferOutputStreamTest.java | 17 +++++++--- 3 files changed, 40 insertions(+), 24 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java index a9fc1e2757d65..4202b40c44448 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java @@ -45,8 +45,8 @@ public class ChunkedByteBufferOutputStream extends ByteBufferOutputStream { private final BufferPool pool; private ByteBuffer currentChunk; private int currentChunkIndex; - // Set once the content has been read for send via buffer(). - private boolean finalized; + // Set once the stream is closed for appends via close(); no further writes or addBuffers are allowed. + private boolean closed; // Single-buffer view produced by flatten() and cached here so repeat buffer() calls // return the same instance. To be removed once scatter-gather (KAFKA-20580) is implemented. private ByteBuffer flattenedBuffer; @@ -121,11 +121,12 @@ public void write(ByteBuffer sourceBuffer) { } /** - * Guards against writes after the stream has been finalized with a call to {@link #buffer()}. + * Guards against writes (and {@link #addBuffers}) after the stream has been closed for appends + * via {@link #close()}. */ private void ensureWritable() { - if (finalized) - throw new IllegalStateException("cannot write after buffer() has been called"); + if (closed) + throw new IllegalStateException("cannot write after the stream has been closed"); } /** @@ -160,19 +161,21 @@ private void advanceToNextChunk() { */ public void addBuffers(List newChunks) { ensureNotDeallocated(); + ensureWritable(); chunks.addAll(newChunks); } /** - * Returns the written bytes as a single contiguous buffer. Calling this finalizes the stream: no - * further writes are allowed (see {@link #ensureWritable()}) and later calls return the same - * instance, which callers such as {@code MemoryRecordsBuilder#writeDefaultBatchHeader} rely on - * when they write the batch header directly into the returned buffer. + * Returns the written bytes as a {@link ByteBuffer}. + *

+ * Currently the chunks are flattened into a single new buffer, built once and cached so repeat + * calls return the same instance, which callers such as + * {@code MemoryRecordsBuilder#writeDefaultBatchHeader} rely on when they write the batch header + * directly into the returned buffer. */ @Override public ByteBuffer buffer() { ensureNotDeallocated(); - finalized = true; if (flattenedBuffer == null) flattenedBuffer = flatten(); return flattenedBuffer; @@ -201,10 +204,12 @@ private ByteBuffer flatten() { } /** - * Releases the fully-unused chunks, given that the stream is closed for appends. + * Closes the stream for appends: no further writes or {@link #addBuffers} are allowed, and the + * fully-unused chunks are released to the pool. */ @Override public void close() { + closed = true; releaseUnusedChunks(); } @@ -290,11 +295,13 @@ public int remaining() { @Override public int limit() { + ensureNotDeallocated(); return Integer.MAX_VALUE; } @Override public int initialCapacity() { + ensureNotDeallocated(); return chunkSize; } @@ -318,7 +325,7 @@ public void deallocate(BufferPool pool) { } chunks.clear(); currentChunk = null; - currentChunkIndex = 0; + currentChunkIndex = -1; flattenedBuffer = null; } diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/RecordAccumulator.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/RecordAccumulator.java index 72a7b17229dc7..c8d09ad1ac312 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/RecordAccumulator.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/RecordAccumulator.java @@ -384,10 +384,11 @@ public RecordAppendResult append(String topic, * subclass passes a supplier that produces a builder backed by a * {@link ChunkedByteBufferOutputStream}. * @param nowMs The current time, in milliseconds - * @return the append result, which never has {@code needsNewBatch=true}: the method either - * propagates a non-{@code needsNewBatch} result from an open batch created concurrently - * (a success, or — incremental strategy — a {@code needsBufferExtension} signal), or it - * creates the new batch here and returns the successful append to it. + * @return the append result, which is never {@code needsNewBatch}. It is either {@code appended} + * — the record was appended, whether to a batch another thread created concurrently or to + * the batch this method creates — or, in the incremental strategy, {@code needsBufferExtension}: + * a concurrent appender created an extendable open batch, so the caller releases its + * pre-allocated buffer and retries via the extension path. */ protected RecordAppendResult appendNewBatch(String topic, int partition, @@ -1277,11 +1278,12 @@ public PartitionerConfig() { */ public static final class RecordAppendResult { /** - * The three mutually-exclusive outcomes of an append attempt. + * The three mutually-exclusive outcomes of an append attempt. Internal representation only; + * callers use {@link #appended()}, {@link #needsBufferExtension()}, and {@link #needsNewBatch()}. */ - public enum Outcome { APPENDED, NEEDS_BUFFER_EXTENSION, NEEDS_NEW_BATCH } + private enum Outcome { APPENDED, NEEDS_BUFFER_EXTENSION, NEEDS_NEW_BATCH } - public final Outcome outcome; + private final Outcome outcome; public final FutureRecordMetadata future; public final boolean batchIsFull; public final boolean newBatchCreated; diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java index 441f6001d012a..2f554a2bda135 100644 --- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java +++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java @@ -82,6 +82,8 @@ public void testOperationsRejectedAfterDeallocate() throws Exception { assertThrows(IllegalStateException.class, stream::position); assertThrows(IllegalStateException.class, stream::buffer); assertThrows(IllegalStateException.class, stream::attachedCapacity); + assertThrows(IllegalStateException.class, stream::limit); + assertThrows(IllegalStateException.class, stream::initialCapacity); assertThrows(IllegalStateException.class, () -> stream.position(1)); assertThrows(IllegalStateException.class, () -> stream.ensureRemaining(1)); assertThrows(IllegalStateException.class, () -> stream.write(1)); @@ -98,20 +100,25 @@ public void testOperationsRejectedAfterDeallocate() throws Exception { } @Test - public void testWritesDisallowedAfterBuffer() throws Exception { + public void testWritesDisallowedAfterClose() throws Exception { int chunkSize = 16; BufferPool p = pool(64, chunkSize); try (ChunkedByteBufferOutputStream stream = new ChunkedByteBufferOutputStream(chunks(p, chunkSize, 2), chunkSize, p)) { stream.write(new byte[]{1, 2, 3}, 0, 3); - // buffer() finalizes the stream: it may be called repeatedly (returning the same - // instance), but any subsequent write must fail. - ByteBuffer first = stream.buffer(); - assertSame(first, stream.buffer(), "buffer() must return the same cached instance once built"); + // close() closes the stream for appends; any subsequent write must fail. + stream.close(); assertThrows(IllegalStateException.class, () -> stream.write(1)); assertThrows(IllegalStateException.class, () -> stream.write(new byte[]{4}, 0, 1)); assertThrows(IllegalStateException.class, () -> stream.write(ByteBuffer.wrap(new byte[]{5}))); + // Attaching more chunks is a write-preparation step, so it is disallowed once closed too. + assertThrows(IllegalStateException.class, + () -> stream.addBuffers(Collections.singletonList(ByteBuffer.allocate(chunkSize)))); + + // buffer() still works after close and returns the same cached instance on repeat calls. + ByteBuffer first = stream.buffer(); + assertSame(first, stream.buffer(), "buffer() must return the same cached instance once built"); stream.deallocate(); } From 68d3fee9f30492b7296ba96bd0d6bc662f800a06 Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Mon, 20 Jul 2026 13:24:56 -0400 Subject: [PATCH 40/61] reuse size --- .../producer/internals/ChunkedRecordAccumulator.java | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java index d7e6b8ea59c05..97f46bbf9aa6f 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java @@ -117,6 +117,7 @@ public RecordAppendResult append(String topic, appendsInProgress.incrementAndGet(); ChunkedByteBufferOutputStream bufferStream = null; + int newBatchSize = 0; // set when bufferStream is allocated; reused as the write-limit basis List extensionChunks = null; if (headers == null) headers = Record.EMPTY_HEADERS; try { @@ -183,15 +184,15 @@ public RecordAppendResult append(String topic, // ratio-adjusted when compressed) so the two stay consistent. int recordUncompressed = AbstractRecords.recordSizeUpperBound( RecordBatch.CURRENT_MAGIC_VALUE, compression.type(), key, value, headers); - int size = MemoryRecordsBuilder.estimatedBytesWritten( + newBatchSize = MemoryRecordsBuilder.estimatedBytesWritten( RecordBatch.CURRENT_MAGIC_VALUE, compression.type(), CompressionRatioEstimator.estimation(topic, compression.type()), recordUncompressed); log.trace("Allocating {} byte chunked buffer ({} byte chunks) for topic {} partition {} with remaining timeout {}ms", - size, chunkedFree.poolableSize(), topic, effectivePartition, maxTimeToBlock); + newBatchSize, chunkedFree.poolableSize(), topic, effectivePartition, maxTimeToBlock); List initialChunks; try { - initialChunks = chunkedFree.allocateChunks(size, maxTimeToBlock); + initialChunks = chunkedFree.allocateChunks(newBatchSize, maxTimeToBlock); } catch (BufferExhaustedException e) { // The blocking new-batch acquire was not able to get memory within // max.block.ms. Record it in the buffer-exhausted metrics. @@ -249,8 +250,9 @@ public RecordAppendResult append(String topic, // so bufferStream was allocated (this iteration or carried from a prior one). if (bufferStream == null) throw new IllegalStateException("needsNewBatch path reached without an allocated buffer stream"); - int firstRecordSize = AbstractRecords.estimateSizeInBytesUpperBound( - RecordBatch.CURRENT_MAGIC_VALUE, compression.type(), key, value, headers); + // Reuse the new-batch size estimate as the write-limit basis (equals the record's + // uncompressed upper bound without compression). TODO: review once compression lands. + final int firstRecordSize = newBatchSize; final ChunkedByteBufferOutputStream batchStream = bufferStream; appendResult = appendNewBatch(topic, effectivePartition, dq, timestamp, key, value, headers, callbacks, () -> chunkedRecordsBuilder(batchStream, firstRecordSize), nowMs); From d9917980786044c3468be246738682cb40f7a757 Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Wed, 22 Jul 2026 13:54:52 -0400 Subject: [PATCH 41/61] dedup on RecordAccumulator --- .../internals/ChunkedRecordAccumulator.java | 21 +++------- .../producer/internals/RecordAccumulator.java | 39 ++++++++++++------- 2 files changed, 32 insertions(+), 28 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java index 97f46bbf9aa6f..563f43e74d85d 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java @@ -112,8 +112,7 @@ public RecordAppendResult append(String topic, long maxTimeToBlock, long nowMs, Cluster cluster) throws InterruptedException { - TopicInfo topicInfo = topicInfoMap.computeIfAbsent(topic, - k -> new TopicInfo(createBuiltInPartitioner(logContext, k, batchSize, partitionerRackAware, rack))); + TopicInfo topicInfo = topicInfoFor(topic); appendsInProgress.incrementAndGet(); ChunkedByteBufferOutputStream bufferStream = null; @@ -145,11 +144,8 @@ public RecordAppendResult append(String topic, // outside the deque lock. A needsNewBatch result means there is no open batch // (full or absent), so it will fall through to the first-record (new batch) path. appendResult = tryAppend(timestamp, key, value, headers, callbacks, dq, nowMs); - if (appendResult.appended()) { - boolean enableSwitch = allBatchesFull(dq); - topicInfo.builtInPartitioner.updatePartitionInfo(partitionInfo, appendResult.appendedBytes, cluster, enableSwitch); - return appendResult; - } + if (appendResult.appended()) + return updatePartitionInfoOnAppend(appendResult, topicInfo, partitionInfo, dq, cluster); } if (appendResult.needsBufferExtension()) { @@ -229,11 +225,8 @@ public RecordAppendResult append(String topic, ((ChunkedProducerBatch) last).addBuffers(extensionChunks); extensionChunks = null; RecordAppendResult retryResult = tryAppend(timestamp, key, value, headers, callbacks, dq, nowMs); - if (retryResult.appended()) { - boolean enableSwitch = allBatchesFull(dq); - topicInfo.builtInPartitioner.updatePartitionInfo(partitionInfo, retryResult.appendedBytes, cluster, enableSwitch); - return retryResult; - } + if (retryResult.appended()) + return updatePartitionInfoOnAppend(retryResult, topicInfo, partitionInfo, dq, cluster); // Still not appended: concurrent appenders filled the batch, // so the extension we attached is no longer enough. // Loop so the next iteration routes the record @@ -269,9 +262,7 @@ public RecordAppendResult append(String topic, } if (appendResult.newBatchCreated) bufferStream = null; - boolean enableSwitch = allBatchesFull(dq); - topicInfo.builtInPartitioner.updatePartitionInfo(partitionInfo, appendResult.appendedBytes, cluster, enableSwitch); - return appendResult; + return updatePartitionInfoOnAppend(appendResult, topicInfo, partitionInfo, dq, cluster); } } } finally { diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/RecordAccumulator.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/RecordAccumulator.java index c8d09ad1ac312..316943bbfbb97 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/RecordAccumulator.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/RecordAccumulator.java @@ -289,7 +289,7 @@ public RecordAppendResult append(String topic, long maxTimeToBlock, long nowMs, Cluster cluster) throws InterruptedException { - TopicInfo topicInfo = topicInfoMap.computeIfAbsent(topic, k -> new TopicInfo(createBuiltInPartitioner(logContext, k, batchSize, partitionerRackAware, rack))); + TopicInfo topicInfo = topicInfoFor(topic); // We keep track of the number of appending thread to make sure we do not miss batches in // abortIncompleteBatches(). @@ -324,12 +324,8 @@ public RecordAppendResult append(String topic, continue; RecordAppendResult appendResult = tryAppend(timestamp, key, value, headers, callbacks, dq, nowMs); - if (appendResult.appended()) { - // If queue has incomplete batches we disable switch (see comments in updatePartitionInfo). - boolean enableSwitch = allBatchesFull(dq); - topicInfo.builtInPartitioner.updatePartitionInfo(partitionInfo, appendResult.appendedBytes, cluster, enableSwitch); - return appendResult; - } + if (appendResult.appended()) + return updatePartitionInfoOnAppend(appendResult, topicInfo, partitionInfo, dq, cluster); } if (buffer == null) { @@ -356,10 +352,7 @@ public RecordAppendResult append(String topic, // Set buffer to null, so that deallocate doesn't return it back to free pool, since it's used in the batch. if (appendResult.newBatchCreated) buffer = null; - // If queue has incomplete batches we disable switch (see comments in updatePartitionInfo). - boolean enableSwitch = allBatchesFull(dq); - topicInfo.builtInPartitioner.updatePartitionInfo(partitionInfo, appendResult.appendedBytes, cluster, enableSwitch); - return appendResult; + return updatePartitionInfoOnAppend(appendResult, topicInfo, partitionInfo, dq, cluster); } } } finally { @@ -368,6 +361,27 @@ public RecordAppendResult append(String topic, } } + /** + * The {@link TopicInfo} for the given topic, creating it (with its built-in partitioner) on first use. + */ + protected TopicInfo topicInfoFor(String topic) { + return topicInfoMap.computeIfAbsent(topic, + k -> new TopicInfo(createBuiltInPartitioner(logContext, k, batchSize, partitionerRackAware, rack))); + } + + /** + * Update the built-in partitioner for a successful append and return the result. Shared by the + * full and incremental strategies at each point a record is appended. + */ + protected RecordAppendResult updatePartitionInfoOnAppend(RecordAppendResult appendResult, TopicInfo topicInfo, + BuiltInPartitioner.StickyPartitionInfo partitionInfo, + Deque dq, Cluster cluster) { + // If the queue has incomplete batches we disable switch (see comments in updatePartitionInfo). + boolean enableSwitch = allBatchesFull(dq); + topicInfo.builtInPartitioner.updatePartitionInfo(partitionInfo, appendResult.appendedBytes, cluster, enableSwitch); + return appendResult; + } + /** * Append a new batch to the queue * @@ -1045,8 +1059,7 @@ public Deque getDeque(TopicPartition tp) { * Get the deque for the given topic-partition, creating it if necessary. */ private Deque getOrCreateDeque(TopicPartition tp) { - TopicInfo topicInfo = topicInfoMap.computeIfAbsent(tp.topic(), - k -> new TopicInfo(createBuiltInPartitioner(logContext, k, batchSize, partitionerRackAware, rack))); + TopicInfo topicInfo = topicInfoFor(tp.topic()); return topicInfo.batches.computeIfAbsent(tp.partition(), k -> new ArrayDeque<>()); } From 058642e2cf5046292c28c72f37165b5fde954787 Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Wed, 22 Jul 2026 14:29:27 -0400 Subject: [PATCH 42/61] pool with allocationMode --- .../kafka/clients/producer/KafkaProducer.java | 4 +- .../producer/internals/BufferPool.java | 29 ++++++++++++++ .../BufferPoolChunkAllocationTest.java | 38 +++++++++---------- .../ChunkedByteBufferOutputStreamTest.java | 6 +-- .../ChunkedRecordAccumulatorTest.java | 14 +++---- 5 files changed, 60 insertions(+), 31 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/KafkaProducer.java b/clients/src/main/java/org/apache/kafka/clients/producer/KafkaProducer.java index a64eedb94e18c..a660302ae3ff8 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/KafkaProducer.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/KafkaProducer.java @@ -487,7 +487,7 @@ public KafkaProducer(Properties properties, Serializer keySerializer, Seriali PRODUCER_METRIC_GROUP_NAME, time, transactionManager, - new BufferPool(this.totalMemorySize, ChunkedRecordAccumulator.CHUNK_SIZE, metrics, time, PRODUCER_METRIC_GROUP_NAME)); + new BufferPool(this.totalMemorySize, ChunkedRecordAccumulator.CHUNK_SIZE, metrics, time, PRODUCER_METRIC_GROUP_NAME, BufferPool.AllocationMode.CHUNKED)); } else { this.accumulator = new RecordAccumulator(logContext, batchSize, @@ -501,7 +501,7 @@ public KafkaProducer(Properties properties, Serializer keySerializer, Seriali PRODUCER_METRIC_GROUP_NAME, time, transactionManager, - new BufferPool(this.totalMemorySize, batchSize, metrics, time, PRODUCER_METRIC_GROUP_NAME)); + new BufferPool(this.totalMemorySize, batchSize, metrics, time, PRODUCER_METRIC_GROUP_NAME, BufferPool.AllocationMode.SINGLE)); } this.errors = this.metrics.sensor("errors"); diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/BufferPool.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/BufferPool.java index c2b2a1ec840a9..ba4c523b4fd5f 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/BufferPool.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/BufferPool.java @@ -49,6 +49,13 @@ public class BufferPool { static final String WAIT_TIME_SENSOR_NAME = "bufferpool-wait-time"; + /** + * Which allocation method a pool serves: {@link #SINGLE} accepts only {@link #allocate} (the full + * strategy), {@link #CHUNKED} only {@link #allocateChunks} (the incremental strategy). Fixed at + * construction so the two are never mixed on the same pool. + */ + public enum AllocationMode { SINGLE, CHUNKED } + private final long totalMemory; private final int poolableSize; /** Lock held for any read or write of {@link #free}, {@link #waiters}, {@link #nonPooledAvailableMemory}, or {@link #closed}. */ @@ -63,6 +70,7 @@ public class BufferPool { protected final Time time; private final Sensor waitTime; protected boolean closed; + private final AllocationMode allocationMode; /** * Create a new buffer pool @@ -74,6 +82,21 @@ public class BufferPool { * @param metricGrpName logical group name for metrics */ public BufferPool(long memory, int poolableSize, Metrics metrics, Time time, String metricGrpName) { + this(memory, poolableSize, metrics, time, metricGrpName, AllocationMode.SINGLE); + } + + /** + * Create a new buffer pool that serves the given {@link AllocationMode}. + * + * @param memory The maximum amount of memory that this buffer pool can allocate + * @param poolableSize The buffer size to cache in the free list rather than deallocating + * @param metrics instance of Metrics + * @param time time instance + * @param metricGrpName logical group name for metrics + * @param allocationMode which allocation method this pool serves ({@link #allocate} vs {@link #allocateChunks}) + */ + public BufferPool(long memory, int poolableSize, Metrics metrics, Time time, String metricGrpName, AllocationMode allocationMode) { + this.allocationMode = allocationMode; this.poolableSize = poolableSize; this.lock = new ReentrantLock(); this.free = new ArrayDeque<>(); @@ -111,6 +134,9 @@ public BufferPool(long memory, int poolableSize, Metrics metrics, Time time, Str * forever) */ public ByteBuffer allocate(int size, long maxTimeToBlockMs) throws InterruptedException { + if (allocationMode != AllocationMode.SINGLE) + throw new IllegalStateException("allocate() is not supported on a " + allocationMode + + " buffer pool; use allocateChunks()"); if (size > this.totalMemory) throw new IllegalArgumentException("Attempt to allocate " + size + " bytes, but there is a hard limit of " @@ -249,6 +275,9 @@ private long awaitMemory(Condition moreMemory, long remainingTimeToBlockNs, * @throws KafkaException if the pool is closed during the wait */ public List allocateChunks(int totalSize, long maxTimeToBlockMs) throws InterruptedException { + if (allocationMode != AllocationMode.CHUNKED) + throw new IllegalStateException("allocateChunks() is not supported on a " + allocationMode + + " buffer pool; use allocate()"); if (totalSize <= 0) throw new IllegalArgumentException("totalSize must be positive: " + totalSize); diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/BufferPoolChunkAllocationTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/BufferPoolChunkAllocationTest.java index c71688e9301a7..9ec3bd20bc1dd 100644 --- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/BufferPoolChunkAllocationTest.java +++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/BufferPoolChunkAllocationTest.java @@ -50,7 +50,7 @@ public void teardown() { private BufferPool pool(long totalMemory, int chunkSize) { String metricGroup = "producer-metrics"; - return new BufferPool(totalMemory, chunkSize, metrics, time, metricGroup); + return new BufferPool(totalMemory, chunkSize, metrics, time, metricGroup, BufferPool.AllocationMode.CHUNKED); } /** Single-chunk request returns a list of one buffer at the chunk size. */ @@ -129,7 +129,7 @@ public void testRollbackOnPartialFailure() throws Exception { long total = 2 * chunkSize; // only 2 chunks worth of memory BufferPool p = pool(total, chunkSize); // Reserve one chunk so the pool has only 1 left. - ByteBuffer held = p.allocate(chunkSize, 100); + ByteBuffer held = p.allocateChunks(chunkSize, 100).get(0); // Request 2 chunks with a zero deadline — first chunk fits, second can't, must throw. assertThrows(BufferExhaustedException.class, () -> p.allocateChunks(2 * chunkSize, 0)); @@ -152,9 +152,9 @@ public void testNoPartialHoldsDuringWait() throws Exception { long total = 4L * chunkSize; BufferPool p = pool(total, chunkSize); // Drain the pool down to 1 chunk's worth so a 2-chunk request can't be satisfied immediately. - ByteBuffer h1 = p.allocate(chunkSize, 100); - ByteBuffer h2 = p.allocate(chunkSize, 100); - ByteBuffer h3 = p.allocate(chunkSize, 100); + ByteBuffer h1 = p.allocateChunks(chunkSize, 100).get(0); + ByteBuffer h2 = p.allocateChunks(chunkSize, 100).get(0); + ByteBuffer h3 = p.allocateChunks(chunkSize, 100).get(0); long available = p.availableMemory(); assertEquals(chunkSize, available, "pool should have exactly 1 chunk free"); @@ -190,15 +190,15 @@ public void testNoPartialHoldsDuringWait() throws Exception { * K-chunk request occupies a single waiter slot, not K of them. */ @Test - public void testFifoFairnessAcrossChunkedAndSingleChunkRequests() throws Exception { + public void testFifoFairnessAcrossMultiChunkAndSingleChunkRequests() throws Exception { int chunkSize = 64; long total = 4L * chunkSize; BufferPool p = pool(total, chunkSize); // Drain the pool entirely so any new request must wait. - ByteBuffer h1 = p.allocate(chunkSize, 100); - ByteBuffer h2 = p.allocate(chunkSize, 100); - ByteBuffer h3 = p.allocate(chunkSize, 100); - ByteBuffer h4 = p.allocate(chunkSize, 100); + ByteBuffer h1 = p.allocateChunks(chunkSize, 100).get(0); + ByteBuffer h2 = p.allocateChunks(chunkSize, 100).get(0); + ByteBuffer h3 = p.allocateChunks(chunkSize, 100).get(0); + ByteBuffer h4 = p.allocateChunks(chunkSize, 100).get(0); assertEquals(0, p.availableMemory()); // T_multi enters first, requesting 2 chunks; T_single enters second, requesting 1 chunk. @@ -217,7 +217,7 @@ public void testFifoFairnessAcrossChunkedAndSingleChunkRequests() throws Excepti }, "multi"); Thread tSingle = new Thread(() -> { try { - ByteBuffer got = p.allocate(chunkSize, 10_000); + ByteBuffer got = p.allocateChunks(chunkSize, 10_000).get(0); singleCompletionMs.set(System.nanoTime()); p.deallocate(got); } catch (Throwable th) { @@ -270,9 +270,9 @@ public void testSlowPathDoesNotDoubleRefundPolledChunksOnTimeout() throws Except BufferPool p = pool(total, chunkSize); // Drain so the next allocateChunks must wait. - ByteBuffer h1 = p.allocate(chunkSize, 100); - ByteBuffer h2 = p.allocate(chunkSize, 100); - ByteBuffer h3 = p.allocate(chunkSize, 100); + ByteBuffer h1 = p.allocateChunks(chunkSize, 100).get(0); + ByteBuffer h2 = p.allocateChunks(chunkSize, 100).get(0); + ByteBuffer h3 = p.allocateChunks(chunkSize, 100).get(0); assertEquals(0, p.availableMemory()); AtomicReference err = new AtomicReference<>(); @@ -325,9 +325,9 @@ public void testSlowPathSuccessDoesNotCorruptAccounting() throws Exception { BufferPool p = pool(total, chunkSize); // Drain so a 2-chunk request must wait. - ByteBuffer h1 = p.allocate(chunkSize, 100); - ByteBuffer h2 = p.allocate(chunkSize, 100); - ByteBuffer h3 = p.allocate(chunkSize, 100); + ByteBuffer h1 = p.allocateChunks(chunkSize, 100).get(0); + ByteBuffer h2 = p.allocateChunks(chunkSize, 100).get(0); + ByteBuffer h3 = p.allocateChunks(chunkSize, 100).get(0); assertEquals(0, p.availableMemory()); AtomicReference err = new AtomicReference<>(); @@ -373,8 +373,8 @@ public void testCloseDuringAtomicWait() throws Exception { long total = 2L * chunkSize; BufferPool p = pool(total, chunkSize); // Drain so the next request must wait. - ByteBuffer h1 = p.allocate(chunkSize, 100); - ByteBuffer h2 = p.allocate(chunkSize, 100); + ByteBuffer h1 = p.allocateChunks(chunkSize, 100).get(0); + ByteBuffer h2 = p.allocateChunks(chunkSize, 100).get(0); assertEquals(0, p.availableMemory()); AtomicReference err = new AtomicReference<>(); diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java index 2f554a2bda135..dfb6731068541 100644 --- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java +++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java @@ -44,7 +44,7 @@ public void teardown() { private BufferPool pool(long total, int chunkSize) { String metricGroup = "test"; - return new BufferPool(total, chunkSize, metrics, time, metricGroup); + return new BufferPool(total, chunkSize, metrics, time, metricGroup, BufferPool.AllocationMode.CHUNKED); } private List chunks(BufferPool pool, int chunkSize, int count) throws InterruptedException { @@ -190,7 +190,7 @@ public void testAddBuffersExtendsStream() throws Exception { assertEquals(0, stream.remaining()); // Extend and write more — must land in the new chunk. - stream.addBuffers(Collections.singletonList(p.allocate(chunkSize, 100))); + stream.addBuffers(Collections.singletonList(p.allocateChunks(chunkSize, 100).get(0))); assertEquals(chunkSize, stream.remaining()); byte[] more = new byte[]{9, 9, 9}; @@ -223,7 +223,7 @@ public void testDeallocateReturnsAllChunks() throws Exception { List initial = chunks(p, chunkSize, 3); try (ChunkedByteBufferOutputStream stream = new ChunkedByteBufferOutputStream(initial, chunkSize, p)) { // Also attach one extra chunk so deallocate must handle the added-buffer case too. - stream.addBuffers(Collections.singletonList(p.allocate(chunkSize, 100))); + stream.addBuffers(Collections.singletonList(p.allocateChunks(chunkSize, 100).get(0))); assertEquals(total - 4L * chunkSize, p.availableMemory()); stream.deallocate(); diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java index b9e4d3ed29b09..2e863ba5b95b9 100644 --- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java +++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java @@ -82,7 +82,7 @@ public void teardown() { } private ChunkedRecordAccumulator newAccumulator(int batchSize, int chunkSize, long totalMemory, Compression compression) { - BufferPool pool = new BufferPool(totalMemory, chunkSize, metrics, time, "producer-metrics"); + BufferPool pool = new BufferPool(totalMemory, chunkSize, metrics, time, "producer-metrics", BufferPool.AllocationMode.CHUNKED); return new ChunkedRecordAccumulator(logContext, batchSize, compression, /* lingerMs */ 0, /* retryBackoffMs */ 0L, /* retryBackoffMaxMs */ 0L, /* deliveryTimeoutMs */ 3200, metrics, "producer-metrics", time, @@ -161,7 +161,7 @@ public void testMidBatchExtensionGrowsExistingBatch() throws Exception { private BufferPool poolMockingConcurrentChunkAllocation(int chunkSize, long totalMemory, AtomicReference injectAppendOnce, byte[] injectedValue) { - return new BufferPool(totalMemory, chunkSize, metrics, time, "producer-metrics") { + return new BufferPool(totalMemory, chunkSize, metrics, time, "producer-metrics", BufferPool.AllocationMode.CHUNKED) { @Override public List allocateChunks(int totalSize, long maxTimeToBlockMs) throws InterruptedException { List chunks = super.allocateChunks(totalSize, maxTimeToBlockMs); @@ -234,7 +234,7 @@ public void testConcurrentExtensionRaceLoserStartsNewBatch() throws Exception { final AtomicReference injectAppendOnce = new AtomicReference<>(); final AtomicInteger chunkDeallocations = new AtomicInteger(); - BufferPool pool = new BufferPool(64L * chunkSize, chunkSize, metrics, time, "producer-metrics") { + BufferPool pool = new BufferPool(64L * chunkSize, chunkSize, metrics, time, "producer-metrics", BufferPool.AllocationMode.CHUNKED) { @Override public List allocateChunks(int totalSize, long maxTimeToBlockMs) throws InterruptedException { List chunks = super.allocateChunks(totalSize, maxTimeToBlockMs); @@ -299,7 +299,7 @@ public void testPartitionSwitchRefundsHeldExtensionChunks() throws Exception { AtomicBoolean extensionAllocated = new AtomicBoolean(false); AtomicInteger chunkDeallocations = new AtomicInteger(0); - BufferPool pool = new BufferPool(totalMemory, chunkSize, metrics, time, "producer-metrics") { + BufferPool pool = new BufferPool(totalMemory, chunkSize, metrics, time, "producer-metrics", BufferPool.AllocationMode.CHUNKED) { @Override public List allocateChunks(int totalSize, long maxTimeToBlockMs) throws InterruptedException { List chunks = super.allocateChunks(totalSize, maxTimeToBlockMs); @@ -364,7 +364,7 @@ protected boolean partitionChanged(String topic, TopicInfo topicInfo, public void testInflightExpirationReturnsAllChunksToPool() throws Exception { int chunkSize = 128; long totalMemory = 32L * chunkSize; - BufferPool pool = new BufferPool(totalMemory, chunkSize, metrics, time, "producer-metrics"); + BufferPool pool = new BufferPool(totalMemory, chunkSize, metrics, time, "producer-metrics", BufferPool.AllocationMode.CHUNKED); ChunkedRecordAccumulator accum = new ChunkedRecordAccumulator(logContext, 8192, Compression.NONE, /* lingerMs */ 0, /* retryBackoffMs */ 0L, /* retryBackoffMaxMs */ 0L, /* deliveryTimeoutMs */ 3200, metrics, "producer-metrics", time, @@ -416,7 +416,7 @@ public void testExhaustedExtensionFallsBackToBlockingNewBatchPath() throws Excep AtomicInteger closeForAppendsCalls = new AtomicInteger(); List closeCallsAtAlloc = new ArrayList<>(); - BufferPool pool = new BufferPool(16L * chunkSize, chunkSize, metrics, time, "producer-metrics") { + BufferPool pool = new BufferPool(16L * chunkSize, chunkSize, metrics, time, "producer-metrics", BufferPool.AllocationMode.CHUNKED) { @Override public List allocateChunks(int totalSize, long maxTimeToBlockMs) throws InterruptedException { allocTimeouts.add(maxTimeToBlockMs); @@ -585,7 +585,7 @@ public void testCumulativeAccountsForBatchHeaderOnce() throws Exception { int chunkSize = 256; int batchSize = 8192; long totalMemory = 64L * chunkSize; - BufferPool pool = new BufferPool(totalMemory, chunkSize, metrics, time, "producer-metrics"); + BufferPool pool = new BufferPool(totalMemory, chunkSize, metrics, time, "producer-metrics", BufferPool.AllocationMode.CHUNKED); ChunkedRecordAccumulator accum = new ChunkedRecordAccumulator(logContext, batchSize, Compression.NONE, /* lingerMs */ 0, /* retryBackoffMs */ 0L, /* retryBackoffMaxMs */ 0L, /* deliveryTimeoutMs */ 3200, metrics, "producer-metrics", time, From 69828a46153774e925b653ee2e273544216a3900 Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Wed, 22 Jul 2026 15:11:43 -0400 Subject: [PATCH 43/61] addressing comments --- .../internals/ChunkedProducerBatch.java | 4 ++ .../internals/ChunkedRecordAccumulator.java | 51 ++++++++++++++----- .../producer/internals/RecordAccumulator.java | 28 ++++++---- .../record/internal/AbstractRecords.java | 16 ++---- .../record/internal/MemoryRecordsBuilder.java | 13 ++--- 5 files changed, 67 insertions(+), 45 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedProducerBatch.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedProducerBatch.java index 33a3b20268154..f32ea107b95eb 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedProducerBatch.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedProducerBatch.java @@ -97,6 +97,10 @@ protected void deallocateBuffer(BufferPool pool) { * layer may still be reading the pooled one — a chunked batch's inflight bytes live in the * separate flattened buffer (see {@link ChunkedByteBufferOutputStream#buffer()}), so it is * safe to return the actual chunks to the pool here. + *

+ * TODO (KAFKA-20580): review when removing the flatten. + * Once we send directly from the chunks, the chunks themselves hold the + * inflight bytes so would be unsafe to return them here. */ @Override protected void deallocateInflightBuffer(BufferPool pool) { diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java index 563f43e74d85d..1c9720967c22f 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java @@ -115,8 +115,9 @@ public RecordAppendResult append(String topic, TopicInfo topicInfo = topicInfoFor(topic); appendsInProgress.incrementAndGet(); - ChunkedByteBufferOutputStream bufferStream = null; - int newBatchSize = 0; // set when bufferStream is allocated; reused as the write-limit basis + // The buffer stream allocated to back a new batch (sized for its first record), paired + // with that size. Set and cleared together across retries; null when none is held. + NewBatchBuffer newBatch = null; List extensionChunks = null; if (headers == null) headers = Record.EMPTY_HEADERS; try { @@ -173,14 +174,14 @@ public RecordAppendResult append(String topic, continue; } nowMs = time.milliseconds(); - } else if (appendResult.needsNewBatch() && bufferStream == null) { + } else if (appendResult.needsNewBatch() && newBatch == null) { // The open batch is done (e.g., full, closed) so start a new one. // Block on the pool for enough chunks to fit this record, sized with the // same cumulative estimator used mid-batch (header + record bytes for NONE, // ratio-adjusted when compressed) so the two stay consistent. int recordUncompressed = AbstractRecords.recordSizeUpperBound( RecordBatch.CURRENT_MAGIC_VALUE, compression.type(), key, value, headers); - newBatchSize = MemoryRecordsBuilder.estimatedBytesWritten( + int newBatchSize = MemoryRecordsBuilder.estimatedBytesWritten( RecordBatch.CURRENT_MAGIC_VALUE, compression.type(), CompressionRatioEstimator.estimation(topic, compression.type()), recordUncompressed); @@ -196,7 +197,9 @@ public RecordAppendResult append(String topic, throw e; } nowMs = time.milliseconds(); - bufferStream = new ChunkedByteBufferOutputStream(initialChunks, chunkedFree.poolableSize(), chunkedFree); + newBatch = new NewBatchBuffer( + new ChunkedByteBufferOutputStream(initialChunks, chunkedFree.poolableSize(), chunkedFree), + newBatchSize); } synchronized (dq) { @@ -241,14 +244,13 @@ public RecordAppendResult append(String topic, // needsNewBatch path: extensionChunks == null here implies needsNewBatch, // so bufferStream was allocated (this iteration or carried from a prior one). - if (bufferStream == null) + if (newBatch == null) throw new IllegalStateException("needsNewBatch path reached without an allocated buffer stream"); // Reuse the new-batch size estimate as the write-limit basis (equals the record's // uncompressed upper bound without compression). TODO: review once compression lands. - final int firstRecordSize = newBatchSize; - final ChunkedByteBufferOutputStream batchStream = bufferStream; + final NewBatchBuffer pending = newBatch; appendResult = appendNewBatch(topic, effectivePartition, dq, timestamp, key, value, headers, callbacks, - () -> chunkedRecordsBuilder(batchStream, firstRecordSize), nowMs); + () -> chunkedRecordsBuilder(pending.stream, pending.size), nowMs); if (appendResult.needsNewBatch()) throw new IllegalStateException("appendNewBatch must not return a needsNewBatch result"); if (appendResult.needsBufferExtension()) { @@ -256,18 +258,18 @@ public RecordAppendResult append(String topic, // than start a new one (detected by appendNewBatch's in-lock tryAppend). // Our bufferStream was sized for a fresh batch — release it and loop so // the extension path allocates exactly the gap-sized chunks. - bufferStream.deallocate(); - bufferStream = null; + newBatch.stream.deallocate(); + newBatch = null; continue; } if (appendResult.newBatchCreated) - bufferStream = null; + newBatch = null; return updatePartitionInfoOnAppend(appendResult, topicInfo, partitionInfo, dq, cluster); } } } finally { - if (bufferStream != null) - bufferStream.deallocate(); + if (newBatch != null) + newBatch.stream.deallocate(); deallocateExtensionChunks(extensionChunks); appendsInProgress.decrementAndGet(); } @@ -291,6 +293,11 @@ private void deallocateExtensionChunks(List extensionChunks) { * attempting the append; the caller allocates chunks outside the deque lock, attaches * them, and retries. Otherwise defers to the parent, which appends or returns * {@link RecordAppendResult#NEEDS_NEW_BATCH}. + * + * @return a {@link RecordAppendResult#needsExtension(int) needsBufferExtension} result when the open batch is + * within its batch-size limit but its chunks lack capacity (the append is not attempted); otherwise the + * parent implementation's outcome: an {@code appended} result ({@link RecordAppendResult#appended()}) or + * {@link RecordAppendResult#NEEDS_NEW_BATCH}. */ @Override protected RecordAppendResult tryAppend(long timestamp, byte[] key, byte[] value, Header[] headers, @@ -328,4 +335,20 @@ private MemoryRecordsBuilder chunkedRecordsBuilder(ChunkedByteBufferOutputStream RecordBatch.NO_PRODUCER_EPOCH, RecordBatch.NO_SEQUENCE, false, false, RecordBatch.NO_PARTITION_LEADER_EPOCH, writeLimit); } + + /** + * A buffer stream allocated to back a new batch, sized to fit the batch's first record, + * paired with that size. That same size is also used as the batch's write-limit basis (see + * {@link #chunkedRecordsBuilder}). The two are always set and cleared together in {@link #append}, + * including across retries. + */ + private static final class NewBatchBuffer { + final ChunkedByteBufferOutputStream stream; + final int size; + + NewBatchBuffer(ChunkedByteBufferOutputStream stream, int size) { + this.stream = stream; + this.size = size; + } + } } diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/RecordAccumulator.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/RecordAccumulator.java index 316943bbfbb97..975c48ca04272 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/RecordAccumulator.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/RecordAccumulator.java @@ -401,8 +401,8 @@ protected RecordAppendResult updatePartitionInfoOnAppend(RecordAppendResult appe * @return the append result, which is never {@code needsNewBatch}. It is either {@code appended} * — the record was appended, whether to a batch another thread created concurrently or to * the batch this method creates — or, in the incremental strategy, {@code needsBufferExtension}: - * a concurrent appender created an extendable open batch, so the caller releases its - * pre-allocated buffer and retries via the extension path. + * a concurrent appender created an extendable open batch, but the new record doesn't fit in, + * so the caller releases its pre-allocated buffer and retries via the extension path. */ protected RecordAppendResult appendNewBatch(String topic, int partition, @@ -420,8 +420,8 @@ protected RecordAppendResult appendNewBatch(String topic, if (!appendResult.needsNewBatch()) { // Propagate without creating a new batch: either another thread already made us a batch // (success), or — incremental strategy — a concurrent appender created an extendable open batch - // (needsBufferExtension), so the caller releases its pre-allocated buffer and retries via - // the extension path. + // that the new record doesn't fit into (needsBufferExtension), so the caller releases its + // pre-allocated buffer and retries via the extension path. return appendResult; } @@ -453,16 +453,22 @@ protected boolean allBatchesFull(Deque deque) { return last == null || last.isFull(); } - /** - * Try to append to a ProducerBatch. + /** + * Try to append to a ProducerBatch. + *

+ * If it is full (or absent), we return {@link RecordAppendResult#NEEDS_NEW_BATCH} and a new batch is created. + * We also close the batch for record appends to free up resources like compression buffers. The batch will be + * fully closed (ie. the record batch headers will be written and memory records built) in one of the following + * cases (whichever comes first): right before send, if it is expired, or when the producer is closed. * - * If it is full (or absent), we return {@link RecordAppendResult#NEEDS_NEW_BATCH} and a new batch is created. - * We also close the batch for record appends to free up resources like compression buffers. The batch will be - * fully closed (ie. the record batch headers will be written and memory records built) in one of the following - * cases (whichever comes first): right before send, if it is expired, or when the producer is closed. + * @return one of two outcomes: an {@code appended} result ({@link RecordAppendResult#appended()}) when the + * record was appended to the open batch, or {@link RecordAppendResult#NEEDS_NEW_BATCH} when there is + * no open batch that can take it (full or absent). The incremental strategy overrides this and may + * additionally return a {@link RecordAppendResult#needsExtension(int) needsBufferExtension} result + * when the open batch is within its batch-size limit but its chunks lack capacity for the record. */ protected RecordAppendResult tryAppend(long timestamp, byte[] key, byte[] value, Header[] headers, - Callback callback, Deque deque, long nowMs) { + Callback callback, Deque deque, long nowMs) { if (closed) throw new KafkaException("Producer closed while send in progress"); ProducerBatch last = deque.peekLast(); diff --git a/clients/src/main/java/org/apache/kafka/common/record/internal/AbstractRecords.java b/clients/src/main/java/org/apache/kafka/common/record/internal/AbstractRecords.java index 4e6e32d75a556..bb94cb4ad13c7 100644 --- a/clients/src/main/java/org/apache/kafka/common/record/internal/AbstractRecords.java +++ b/clients/src/main/java/org/apache/kafka/common/record/internal/AbstractRecords.java @@ -149,20 +149,14 @@ else if (compressionType != CompressionType.NONE) */ public static int recordSizeUpperBound(byte magic, CompressionType compressionType, byte[] key, byte[] value, Header[] headers) { - return recordSizeUpperBound(magic, compressionType, Utils.wrapNullable(key), Utils.wrapNullable(value), headers); - } - - /** - * @see #recordSizeUpperBound(byte, CompressionType, byte[], byte[], Header[]) - */ - private static int recordSizeUpperBound(byte magic, CompressionType compressionType, ByteBuffer key, - ByteBuffer value, Header[] headers) { + ByteBuffer keyBuffer = Utils.wrapNullable(key); + ByteBuffer valueBuffer = Utils.wrapNullable(value); if (magic >= RecordBatch.MAGIC_VALUE_V2) - return DefaultRecord.recordSizeUpperBound(key, value, headers); + return DefaultRecord.recordSizeUpperBound(keyBuffer, valueBuffer, headers); else if (compressionType != CompressionType.NONE) - return Records.LOG_OVERHEAD + LegacyRecord.recordOverhead(magic) + LegacyRecord.recordSize(magic, key, value); + return Records.LOG_OVERHEAD + LegacyRecord.recordOverhead(magic) + LegacyRecord.recordSize(magic, keyBuffer, valueBuffer); else - return Records.LOG_OVERHEAD + LegacyRecord.recordSize(magic, key, value); + return Records.LOG_OVERHEAD + LegacyRecord.recordSize(magic, keyBuffer, valueBuffer); } /** diff --git a/clients/src/main/java/org/apache/kafka/common/record/internal/MemoryRecordsBuilder.java b/clients/src/main/java/org/apache/kafka/common/record/internal/MemoryRecordsBuilder.java index 6998314c02df3..946449ec0ea3f 100644 --- a/clients/src/main/java/org/apache/kafka/common/record/internal/MemoryRecordsBuilder.java +++ b/clients/src/main/java/org/apache/kafka/common/record/internal/MemoryRecordsBuilder.java @@ -851,18 +851,13 @@ public static int estimatedBytesWritten(byte magic, CompressionType compressionT * incremental strategy to size mid-batch chunk extensions. */ public int estimatedBytesWrittenAfter(byte[] key, byte[] value, Header[] headers) { - return estimatedBytesWrittenAfter(wrapNullable(key), wrapNullable(value), headers); - } - - /** - * @see #estimatedBytesWrittenAfter(byte[], byte[], Header[]) - */ - private int estimatedBytesWrittenAfter(ByteBuffer key, ByteBuffer value, Header[] headers) { + ByteBuffer keyBuffer = wrapNullable(key); + ByteBuffer valueBuffer = wrapNullable(value); final int recordSize; if (magic < RecordBatch.MAGIC_VALUE_V2) { - recordSize = Records.LOG_OVERHEAD + LegacyRecord.recordSize(magic, key, value); + recordSize = Records.LOG_OVERHEAD + LegacyRecord.recordSize(magic, keyBuffer, valueBuffer); } else { - recordSize = DefaultRecord.recordSizeUpperBound(key, value, headers); + recordSize = DefaultRecord.recordSizeUpperBound(keyBuffer, valueBuffer, headers); } return estimatedBytesWritten(magic, compression.type(), estimatedCompressionRatio, uncompressedRecordsSizeInBytes + recordSize); From 87074bf85bc1299cdd5cbbd17ecbca79c54badac Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Wed, 22 Jul 2026 16:11:53 -0400 Subject: [PATCH 44/61] more updates in tests --- .../BufferPoolChunkAllocationTest.java | 103 +++++++++--------- .../ChunkedRecordAccumulatorTest.java | 42 ++++--- 2 files changed, 82 insertions(+), 63 deletions(-) diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/BufferPoolChunkAllocationTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/BufferPoolChunkAllocationTest.java index 9ec3bd20bc1dd..bd6da6d0b1641 100644 --- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/BufferPoolChunkAllocationTest.java +++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/BufferPoolChunkAllocationTest.java @@ -43,6 +43,11 @@ public class BufferPoolChunkAllocationTest { private final MockTime time = new MockTime(); private final Metrics metrics = new Metrics(time); + // Safety ceiling for block-until-available requests that the test unblocks itself (by + // deallocating or closing). The waiter is always signaled first, so a passing test returns in + // milliseconds; this only bounds how long a broken test can block before failing. + private static final long MAX_BLOCK_TIME_MS = 2_000; + @AfterEach public void teardown() { metrics.close(); @@ -149,12 +154,11 @@ public void testRollbackOnPartialFailure() throws Exception { @Test public void testNoPartialHoldsDuringWait() throws Exception { int chunkSize = 64; - long total = 4L * chunkSize; + // Hold 2 of 3 chunks so exactly 1 is free — a 2-chunk request can't be satisfied immediately. + long total = 3L * chunkSize; BufferPool p = pool(total, chunkSize); - // Drain the pool down to 1 chunk's worth so a 2-chunk request can't be satisfied immediately. ByteBuffer h1 = p.allocateChunks(chunkSize, 100).get(0); ByteBuffer h2 = p.allocateChunks(chunkSize, 100).get(0); - ByteBuffer h3 = p.allocateChunks(chunkSize, 100).get(0); long available = p.availableMemory(); assertEquals(chunkSize, available, "pool should have exactly 1 chunk free"); @@ -162,7 +166,7 @@ public void testNoPartialHoldsDuringWait() throws Exception { AtomicReference err = new AtomicReference<>(); Thread t = new Thread(() -> { try { - p.allocateChunks(2 * chunkSize, 5_000); + p.allocateChunks(2 * chunkSize, MAX_BLOCK_TIME_MS); } catch (Throwable th) { err.set(th); } @@ -176,12 +180,12 @@ public void testNoPartialHoldsDuringWait() throws Exception { assertEquals(chunkSize, p.availableMemory(), "pool memory must reflect no partial holds during the atomic wait"); - // Free enough chunks for the waiter to complete. + // The pool still has 1 chunk free (asserted above) and the waiter never consumed it, so + // freeing just one more chunk reaches the 2 it needs and unblocks it. p.deallocate(h1); - p.deallocate(h2); - t.join(5_000); + t.join(); assertNull(err.get(), "waiter unexpectedly threw: " + err.get()); - p.deallocate(h3); + p.deallocate(h2); } /** @@ -192,23 +196,26 @@ public void testNoPartialHoldsDuringWait() throws Exception { @Test public void testFifoFairnessAcrossMultiChunkAndSingleChunkRequests() throws Exception { int chunkSize = 64; - long total = 4L * chunkSize; + // Exactly the chunks the two waiters need: 2 for the multi request, 1 for the single. + long total = 3L * chunkSize; BufferPool p = pool(total, chunkSize); // Drain the pool entirely so any new request must wait. ByteBuffer h1 = p.allocateChunks(chunkSize, 100).get(0); ByteBuffer h2 = p.allocateChunks(chunkSize, 100).get(0); ByteBuffer h3 = p.allocateChunks(chunkSize, 100).get(0); - ByteBuffer h4 = p.allocateChunks(chunkSize, 100).get(0); assertEquals(0, p.availableMemory()); // T_multi enters first, requesting 2 chunks; T_single enters second, requesting 1 chunk. AtomicReference multiCompletionMs = new AtomicReference<>(); AtomicReference singleCompletionMs = new AtomicReference<>(); + AtomicReference> multiChunks = new AtomicReference<>(); + AtomicReference singleChunk = new AtomicReference<>(); CountDownLatch multiStarted = new CountDownLatch(1); Thread tMulti = new Thread(() -> { try { multiStarted.countDown(); - List got = p.allocateChunks(2 * chunkSize, 10_000); + List got = p.allocateChunks(2 * chunkSize, MAX_BLOCK_TIME_MS); + multiChunks.set(got); multiCompletionMs.set(System.nanoTime()); for (ByteBuffer b : got) p.deallocate(b); } catch (Throwable th) { @@ -217,7 +224,8 @@ public void testFifoFairnessAcrossMultiChunkAndSingleChunkRequests() throws Exce }, "multi"); Thread tSingle = new Thread(() -> { try { - ByteBuffer got = p.allocateChunks(chunkSize, 10_000).get(0); + ByteBuffer got = p.allocateChunks(chunkSize, MAX_BLOCK_TIME_MS).get(0); + singleChunk.set(got); singleCompletionMs.set(System.nanoTime()); p.deallocate(got); } catch (Throwable th) { @@ -239,32 +247,33 @@ public void testFifoFairnessAcrossMultiChunkAndSingleChunkRequests() throws Exce p.deallocate(h2); // Both freed — the FIFO leader (multi) should claim both before the single-chunk waiter // gets any. The multi completes first. - tMulti.join(5_000); + tMulti.join(); // Now free another chunk for the single-chunk waiter. p.deallocate(h3); - tSingle.join(5_000); + tSingle.join(); assertNotNull(multiCompletionMs.get(), "multi did not complete"); assertNotNull(singleCompletionMs.get(), "single did not complete"); assertTrue(multiCompletionMs.get() != -1L && singleCompletionMs.get() != -1L, "both threads must complete without exception"); + // Each waiter must actually have received its chunks. + assertNotNull(multiChunks.get(), "multi did not receive its chunks"); + assertEquals(2, multiChunks.get().size(), "multi must receive exactly 2 chunks"); + assertNotNull(singleChunk.get(), "single did not receive its chunk"); assertTrue(multiCompletionMs.get() <= singleCompletionMs.get(), "FIFO violated: single (joined later) completed at " + singleCompletionMs.get() + " before multi at " + multiCompletionMs.get()); - p.deallocate(h4); } /** - * Slow-path rollback must not double-refund pooled chunks. When the waiter polls some - * chunks from the free list, then times out on a subsequent iteration, the inner finally - * refunds {@code accumulated} bytes to {@code nonPooledAvailableMemory} and the outer - * catch re-inserts the polled chunks into {@code free}. If {@code accumulated} includes - * bytes that came from polled chunks, the same memory is refunded twice — the pool - * over-reports {@code availableMemory()} and subsequent allocations can exceed the - * configured {@code buffer.memory} limit. + * A chunk request that can't be satisfied immediately must wait for memory, taking chunks as + * they are freed. If it then times out before acquiring all it needs, each chunk it already + * took must be returned to the pool exactly once. A double refund would make the pool + * over-report {@code availableMemory()}, letting later allocations exceed the configured + * {@code buffer.memory} limit. */ @Test - public void testSlowPathDoesNotDoubleRefundPolledChunksOnTimeout() throws Exception { + public void testWaitingRequestDoesNotDoubleRefundChunksOnTimeout() throws Exception { int chunkSize = 64; long total = 3L * chunkSize; BufferPool p = pool(total, chunkSize); @@ -278,10 +287,8 @@ public void testSlowPathDoesNotDoubleRefundPolledChunksOnTimeout() throws Except AtomicReference err = new AtomicReference<>(); Thread t = new Thread(() -> { try { - // Asks for 3 chunks with a finite deadline. After we deallocate 2 chunks below - // the waiter will poll them, find itself still 1 short, and time out on the - // next await iteration — with 2 chunks held in `pooled` and 2*chunkSize in - // `accumulated`. + // Asks for 3 chunks with a finite deadline. We free 2 chunks below, so the waiter + // takes those but is still 1 short and times out waiting for the 3rd. p.allocateChunks(3 * chunkSize, 500); } catch (Throwable th) { err.set(th); @@ -291,35 +298,32 @@ public void testSlowPathDoesNotDoubleRefundPolledChunksOnTimeout() throws Except TestUtils.waitForCondition(() -> p.queued() == 1, "waiter should be parked on the pool's queue"); - // Hand the waiter 2 chunks — it polls them into `pooled` then awaits again for the 3rd. + // Free 2 chunks — the waiter takes both, then waits again for the 3rd it never gets. p.deallocate(h1); p.deallocate(h2); - // Give the waiter a moment to wake, poll, and re-enter await before the timeout fires. + // Give the waiter a moment to wake, take the 2 chunks, and re-enter the wait before it times out. Thread.sleep(100); - t.join(5_000); + t.join(); assertInstanceOf(BufferExhaustedException.class, err.get(), "expected BufferExhaustedException, got " + err.get()); - // After the throw, exactly 2 chunkSizes of physical memory should be back in the pool - // (h1 + h2). h3 is still held outside the pool. The bug double-refunds: it credits - // 2*chunkSize to `nonPooledAvailableMemory` (inner finally) AND re-inserts h1, h2 into - // `free` (outer catch), so availableMemory() reports 4*chunkSize. + // After the timeout, exactly the 2 freed chunks (h1 + h2) must be back in the pool, each + // returned once; h3 is still held outside the pool. A double refund would over-report + // availableMemory() (e.g. 4*chunkSize instead of 2). assertEquals(2L * chunkSize, p.availableMemory(), - "pool over-reports availableMemory due to double-refund of pooled chunks on rollback"); + "pool over-reports availableMemory if polled chunks are refunded twice on rollback"); p.deallocate(h3); } /** - * Slow-path success must not corrupt the pool's accounting. When the waiter polls chunks - * from the free list to satisfy its request, the finally block runs with the inner loop's - * "success" reset still in effect — it must not refund or decrement {@code nonPool}. - * Pool chunks transferred to the caller via {@code pooled} are accounted for by removing - * them from {@code free}, not by deducting their bytes from {@code nonPool}. + * A chunk request that must wait for memory and then completes by taking chunks as they are + * freed must account for those chunks exactly once — as memory now held by the caller — so the + * pool neither loses nor over-reports capacity. */ @Test - public void testSlowPathSuccessDoesNotCorruptAccounting() throws Exception { + public void testWaitingRequestDoesNotCorruptAccountingOnSuccess() throws Exception { int chunkSize = 64; long total = 3L * chunkSize; BufferPool p = pool(total, chunkSize); @@ -334,7 +338,7 @@ public void testSlowPathSuccessDoesNotCorruptAccounting() throws Exception { AtomicReference> got = new AtomicReference<>(); Thread t = new Thread(() -> { try { - got.set(p.allocateChunks(2 * chunkSize, 10_000)); + got.set(p.allocateChunks(2 * chunkSize, MAX_BLOCK_TIME_MS)); } catch (Throwable th) { err.set(th); } @@ -343,17 +347,16 @@ public void testSlowPathSuccessDoesNotCorruptAccounting() throws Exception { TestUtils.waitForCondition(() -> p.queued() == 1, "waiter should be parked on the pool's queue"); - // Deallocate 2 chunks → waiter wakes, polls both, accumulated reaches memoryRequired, - // exits normally. The success path's finally must NOT refund or decrement nonPool. + // Free 2 chunks → the waiter wakes, takes both, and completes normally. p.deallocate(h1); p.deallocate(h2); - t.join(5_000); + t.join(); assertNull(err.get(), "waiter unexpectedly threw: " + err.get()); - // After success: waiter owns the 2 chunks (returned via `got`), the test still holds h3. - // The pool has lent out everything it had — availableMemory must be exactly 0. + // After success the waiter owns the 2 chunks (returned via `got`) and the test still holds + // h3, so the pool has lent out everything it had — availableMemory must be exactly 0. assertEquals(0, p.availableMemory(), - "slow-path success corrupted accounting (likely a phantom nonPool decrement in the finally)"); + "a completed waiting request must leave the pool with no available memory (accounting not corrupted)"); // Returning all the held buffers must fully restore the pool. for (ByteBuffer chunk : got.get()) @@ -380,7 +383,7 @@ public void testCloseDuringAtomicWait() throws Exception { AtomicReference err = new AtomicReference<>(); Thread t = new Thread(() -> { try { - p.allocateChunks(2 * chunkSize, 10_000); + p.allocateChunks(2 * chunkSize, MAX_BLOCK_TIME_MS); } catch (Throwable th) { err.set(th); } @@ -390,7 +393,7 @@ public void testCloseDuringAtomicWait() throws Exception { // Close the pool: signals all waiters; the waiter must throw KafkaException. p.close(); - t.join(5_000); + t.join(); assertInstanceOf(KafkaException.class, err.get(), "expected KafkaException, got " + err.get()); p.deallocate(h1); p.deallocate(h2); diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java index 2e863ba5b95b9..1b688f30f397d 100644 --- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java +++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java @@ -159,8 +159,8 @@ public void testMidBatchExtensionGrowsExistingBatch() throws Exception { * appender racing the same batch. */ private BufferPool poolMockingConcurrentChunkAllocation(int chunkSize, long totalMemory, - AtomicReference injectAppendOnce, - byte[] injectedValue) { + AtomicReference injectAppendOnce, + byte[] injectedValue) { return new BufferPool(totalMemory, chunkSize, metrics, time, "producer-metrics", BufferPool.AllocationMode.CHUNKED) { @Override public List allocateChunks(int totalSize, long maxTimeToBlockMs) throws InterruptedException { @@ -271,7 +271,6 @@ public void deallocate(ByteBuffer buffer) { Deque dq = batchesFor(accum, tp1); assertEquals(2, dq.size(), "The losing record must roll to a new batch, not exceed batch.size"); assertNotNull(dq.peekFirst()); - assertNotNull(dq.peekLast()); assertEquals(2, dq.peekFirst().recordCount, "First batch should have the initial record plus the race winner"); assertNotNull(dq.peekLast()); assertEquals(1, dq.peekLast().recordCount, "Second batch should have the loser record"); @@ -465,12 +464,10 @@ public void closeForRecordAppends() { Deque dq = batchesFor(accum, tp1); assertEquals(2, dq.size(), "closed batch + new batch expected"); - assertNotNull(dq.peekFirst()); - assertNotNull(dq.peekLast()); // The original batch is far below its writeLimit, so isFull() can only be true via // the closed append stream — i.e., it was closed for appends on the failed extension. - assertTrue(dq.peekFirst().isFull(), "original batch must be closed for appends"); assertNotNull(dq.peekFirst()); + assertTrue(dq.peekFirst().isFull(), "original batch must be closed for appends"); assertEquals(1, dq.peekFirst().recordCount); assertNotNull(dq.peekLast()); assertEquals(1, dq.peekLast().recordCount); @@ -498,11 +495,14 @@ public void testBufferExhaustedNotDoubleCountedAcrossExtensionAndNewBatch() thro maxBlockTimeMs, time.milliseconds(), cluster); assertEquals(0.0, (double) exhausted.metricValue()); - // Second record overflows the batch's chunk, so extension attempt fails (pool empty, - // recovered by closing the batch), then the new-batch acquire fails too (maxTimeToBlock - // = 0). The two failed acquires must count as a single dropped record. + // Second record overflows the batch's chunk. The extension attempt (always non-blocking) + // fails first — pool empty, recovered by closing the batch — then the new-batch acquire + // blocks up to max.block.ms and also fails since the pool stays empty. A small block time + // keeps the test fast (nothing frees memory during the wait). Both failed acquires must + // count as a single dropped record. + long newBatchBlockMs = 50L; assertThrows(BufferExhaustedException.class, () -> accum.append(topic, partition1, 0L, key, - new byte[150], Record.EMPTY_HEADERS, null, 0L, time.milliseconds(), cluster)); + new byte[150], Record.EMPTY_HEADERS, null, newBatchBlockMs, time.milliseconds(), cluster)); assertEquals(1.0, (double) exhausted.metricValue(), "the dropped record must be counted once, not once per failed acquire"); } finally { @@ -516,8 +516,14 @@ public void testBufferExhaustedNotDoubleCountedAcrossExtensionAndNewBatch() thro @Test public void testBatchCloseDoesNotDeallocateChunksPrematurely() throws Exception { int chunkSize = 256; - ChunkedRecordAccumulator accum = newAccumulator(8192, chunkSize, 32L * chunkSize, Compression.NONE); + long totalMemory = 32L * chunkSize; + BufferPool pool = new BufferPool(totalMemory, chunkSize, metrics, time, "producer-metrics", BufferPool.AllocationMode.CHUNKED); + ChunkedRecordAccumulator accum = new ChunkedRecordAccumulator(logContext, 8192, Compression.NONE, + /* lingerMs */ 0, /* retryBackoffMs */ 0L, /* retryBackoffMaxMs */ 0L, + /* deliveryTimeoutMs */ 3200, metrics, "producer-metrics", time, + /* transactionManager */ null, pool); + // A small record fits in a single chunk, so exactly one (data-bearing) chunk is reserved. byte[] value = new byte[64]; accum.append(topic, partition1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, time.milliseconds(), cluster); @@ -525,9 +531,13 @@ public void testBatchCloseDoesNotDeallocateChunksPrematurely() throws Exception Deque dq = batchesFor(accum, tp1); ProducerBatch batch = dq.peekFirst(); assertNotNull(batch); + assertEquals(totalMemory - chunkSize, pool.availableMemory(), "append must reserve exactly one chunk"); // Matches the call sites in RecordAccumulator.drain and Sender. batch.close(); + // close() must not return the chunk to the pool: available memory is unchanged, and the + // built record set is still readable because its bytes still live in the batch's chunk. + assertEquals(totalMemory - chunkSize, pool.availableMemory(), "close() must not return chunks to the pool"); MemoryRecords records = batch.records(); assertTrue(records.sizeInBytes() > 0, "batch must produce a non-empty record set after close; chunks were deallocated prematurely"); @@ -541,6 +551,10 @@ public void testBatchCloseDoesNotDeallocateChunksPrematurely() throws Exception } assertEquals(1, count, "expected exactly 1 record after close"); + // Completion (deallocate) is what returns the chunk to the pool. + accum.deallocate(batch); + assertEquals(totalMemory, pool.availableMemory(), "deallocate must return the chunk to the pool"); + accum.close(); } @@ -566,8 +580,10 @@ public void testExtensionTracksCumulativeBatchSize() throws Exception { assertNotNull(batch); assertEquals(6, batch.recordCount); - // batch.close() flattens chunks and writes the header. If chunks under-allocated, this - // throws (insufficient capacity in the underlying buffer). + // Finalize the batch: close() writes the record-batch header and builds the record set. + // Chunk capacity for every record is ensured at append time (the accumulator attaches + // extension chunks before appending, and an append without capacity would throw), so the + // chunks already hold the whole batch by the time it is built. batch.close(); MemoryRecords records = batch.records(); int actualSize = records.sizeInBytes(); From 8e9070b49c8f60c59a487f6f89ad15c533e043e3 Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Wed, 22 Jul 2026 16:23:30 -0400 Subject: [PATCH 45/61] assert stream closed on buffer call --- .../ChunkedByteBufferOutputStream.java | 7 ++++- .../ChunkedByteBufferOutputStreamTest.java | 26 ++++++++++++------- 2 files changed, 22 insertions(+), 11 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java index 4202b40c44448..a76ea1b0c1c61 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java @@ -166,16 +166,21 @@ public void addBuffers(List newChunks) { } /** - * Returns the written bytes as a {@link ByteBuffer}. + * Returns the written bytes as a {@link ByteBuffer}. Must be called only after the stream is + * {@link #close() closed for appends}. *

* Currently the chunks are flattened into a single new buffer, built once and cached so repeat * calls return the same instance, which callers such as * {@code MemoryRecordsBuilder#writeDefaultBatchHeader} rely on when they write the batch header * directly into the returned buffer. + * + * @throws IllegalStateException if the stream has not been closed for appends */ @Override public ByteBuffer buffer() { ensureNotDeallocated(); + if (!closed) + throw new IllegalStateException("buffer() must not be called before the stream is closed for appends"); if (flattenedBuffer == null) flattenedBuffer = flatten(); return flattenedBuffer; diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java index dfb6731068541..2f33bd577cd30 100644 --- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java +++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java @@ -133,6 +133,7 @@ public void testSingleChunkWriteRoundtrip() throws Exception { byte[] payload = new byte[]{1, 2, 3, 4, 5}; stream.write(payload, 0, payload.length); + stream.close(); ByteBuffer flat = stream.buffer(); flat.flip(); byte[] out = new byte[flat.remaining()]; @@ -153,6 +154,7 @@ public void testWriteAcrossMultipleChunks() throws Exception { for (int i = 0; i < payload.length; i++) payload[i] = (byte) i; stream.write(payload, 0, payload.length); + stream.close(); ByteBuffer flat = stream.buffer(); flat.flip(); byte[] out = new byte[flat.remaining()]; @@ -233,8 +235,8 @@ public void testDeallocateReturnsAllChunks() throws Exception { /** * The fully-unused chunks are returned to the pool on {@link ChunkedByteBufferOutputStream#close()} - * (the stream is closed for appends), not as a side effect of reading {@link - * ChunkedByteBufferOutputStream#buffer()}. Data-bearing chunks stay reserved until + * (the stream is closed for appends). Reading {@link ChunkedByteBufferOutputStream#buffer()} + * afterwards has no chunk-releasing side effect; the data-bearing chunks stay reserved until * {@link ChunkedByteBufferOutputStream#deallocate()}, which must not return the already-released * chunks a second time. */ @@ -249,14 +251,9 @@ public void testUnusedChunksReleasedOnCloseNotOnBuffer() throws Exception { stream.write(payload, 0, payload.length); assertEquals(total - 3L * chunkSize, p.availableMemory()); - // reading buffer() on a still-open stream must not release chunks. - ByteBuffer built = stream.buffer(); - assertEquals(total - 3L * chunkSize, p.availableMemory(), - "buffer() must not release chunks on a still-open stream"); - built.flip(); - byte[] out = new byte[built.remaining()]; - built.get(out); - assertArrayEquals(payload, out); + // buffer() is only valid once the stream is closed for appends. + assertThrows(IllegalStateException.class, stream::buffer, + "buffer() must not be called before the stream is closed"); // close() (appends done) releases the two unused chunks; a second close() is a no-op. stream.close(); @@ -265,6 +262,15 @@ public void testUnusedChunksReleasedOnCloseNotOnBuffer() throws Exception { stream.close(); assertEquals(total - chunkSize, p.availableMemory()); + // Reading buffer() after close must not release the remaining data-bearing chunk. + ByteBuffer built = stream.buffer(); + assertEquals(total - chunkSize, p.availableMemory(), + "buffer() must not release chunks"); + built.flip(); + byte[] out = new byte[built.remaining()]; + built.get(out); + assertArrayEquals(payload, out); + // Completion-time deallocate returns only the remaining data-bearing chunk (no double free). stream.deallocate(); assertEquals(total, p.availableMemory(), From 79caa2f7bfa729dec5931d41c9eacf09be4e2bed Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Wed, 22 Jul 2026 16:27:46 -0400 Subject: [PATCH 46/61] validate chunk capacity on addBuffers & test --- .../internals/ChunkedByteBufferOutputStream.java | 13 +++++++++++-- .../ChunkedByteBufferOutputStreamTest.java | 13 +++++++++++++ 2 files changed, 24 insertions(+), 2 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java index a76ea1b0c1c61..63606839dbe1c 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java @@ -77,12 +77,20 @@ public ChunkedByteBufferOutputStream(List initialChunks, int chunkSi private static ByteBuffer validatedFirstChunk(List initialChunks, int chunkSize) { if (initialChunks == null || initialChunks.isEmpty()) throw new IllegalArgumentException("initialChunks must be non-empty"); - for (ByteBuffer chunk : initialChunks) { + validateChunkCapacities(initialChunks, chunkSize); + return initialChunks.get(0); + } + + /** + * Validates that every chunk's capacity equals {@code chunkSize}, which the stream's capacity + * bookkeeping relies on. + */ + private static void validateChunkCapacities(List chunks, int chunkSize) { + for (ByteBuffer chunk : chunks) { if (chunk.capacity() != chunkSize) throw new IllegalArgumentException("each chunk must have capacity " + chunkSize + ", but found a chunk of capacity " + chunk.capacity()); } - return initialChunks.get(0); } @Override @@ -162,6 +170,7 @@ private void advanceToNextChunk() { public void addBuffers(List newChunks) { ensureNotDeallocated(); ensureWritable(); + validateChunkCapacities(newChunks, chunkSize); chunks.addAll(newChunks); } diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java index 2f33bd577cd30..a836e9e0c5a32 100644 --- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java +++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java @@ -203,6 +203,19 @@ public void testAddBuffersExtendsStream() throws Exception { } } + @Test + public void testAddBuffersRejectsWrongSizeChunks() throws Exception { + int chunkSize = 8; + BufferPool p = pool(64, chunkSize); + try (ChunkedByteBufferOutputStream stream = new ChunkedByteBufferOutputStream(chunks(p, chunkSize, 1), chunkSize, p)) { + // A chunk whose capacity doesn't match chunkSize violates the contract, same as the constructor. + List wrongSize = Collections.singletonList(ByteBuffer.allocate(chunkSize + 1)); + assertThrows(IllegalArgumentException.class, () -> stream.addBuffers(wrongSize)); + + stream.deallocate(); + } + } + @Test public void testPositionWalksAcrossChunks() throws Exception { int chunkSize = 4; From b1a48d99d783eea6dd6071d6ccc0cf0e633abb35 Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Wed, 22 Jul 2026 16:38:28 -0400 Subject: [PATCH 47/61] reuse estimator --- .../internals/ChunkedRecordAccumulator.java | 20 ++++++++----------- .../record/internal/AbstractRecords.java | 16 --------------- .../record/internal/MemoryRecordsBuilder.java | 9 +++------ 3 files changed, 11 insertions(+), 34 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java index 1c9720967c22f..c90c120eab76f 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java @@ -27,7 +27,6 @@ import org.apache.kafka.common.metrics.Metrics; import org.apache.kafka.common.record.TimestampType; import org.apache.kafka.common.record.internal.AbstractRecords; -import org.apache.kafka.common.record.internal.CompressionRatioEstimator; import org.apache.kafka.common.record.internal.CompressionType; import org.apache.kafka.common.record.internal.MemoryRecordsBuilder; import org.apache.kafka.common.record.internal.Record; @@ -175,16 +174,13 @@ public RecordAppendResult append(String topic, } nowMs = time.milliseconds(); } else if (appendResult.needsNewBatch() && newBatch == null) { - // The open batch is done (e.g., full, closed) so start a new one. - // Block on the pool for enough chunks to fit this record, sized with the - // same cumulative estimator used mid-batch (header + record bytes for NONE, - // ratio-adjusted when compressed) so the two stay consistent. - int recordUncompressed = AbstractRecords.recordSizeUpperBound( + // The open batch is done (e.g., full, closed) so start a new one. Size it for + // this first record with the same estimator the full strategy uses + // (RecordAccumulator.append), but reserve only enough for the record rather than + // a whole batch.size. + // TODO: review when compression is supported. + int newBatchSize = AbstractRecords.estimateSizeInBytesUpperBound( RecordBatch.CURRENT_MAGIC_VALUE, compression.type(), key, value, headers); - int newBatchSize = MemoryRecordsBuilder.estimatedBytesWritten( - RecordBatch.CURRENT_MAGIC_VALUE, compression.type(), - CompressionRatioEstimator.estimation(topic, compression.type()), - recordUncompressed); log.trace("Allocating {} byte chunked buffer ({} byte chunks) for topic {} partition {} with remaining timeout {}ms", newBatchSize, chunkedFree.poolableSize(), topic, effectivePartition, maxTimeToBlock); List initialChunks; @@ -246,8 +242,8 @@ public RecordAppendResult append(String topic, // so bufferStream was allocated (this iteration or carried from a prior one). if (newBatch == null) throw new IllegalStateException("needsNewBatch path reached without an allocated buffer stream"); - // Reuse the new-batch size estimate as the write-limit basis (equals the record's - // uncompressed upper bound without compression). TODO: review once compression lands. + // Reuse the new-batch size estimate as the write-limit basis. + // TODO: review when compression is supported. final NewBatchBuffer pending = newBatch; appendResult = appendNewBatch(topic, effectivePartition, dq, timestamp, key, value, headers, callbacks, () -> chunkedRecordsBuilder(pending.stream, pending.size), nowMs); diff --git a/clients/src/main/java/org/apache/kafka/common/record/internal/AbstractRecords.java b/clients/src/main/java/org/apache/kafka/common/record/internal/AbstractRecords.java index bb94cb4ad13c7..d0704d18ec37a 100644 --- a/clients/src/main/java/org/apache/kafka/common/record/internal/AbstractRecords.java +++ b/clients/src/main/java/org/apache/kafka/common/record/internal/AbstractRecords.java @@ -143,22 +143,6 @@ else if (compressionType != CompressionType.NONE) return Records.LOG_OVERHEAD + LegacyRecord.recordSize(magic, key, value); } - /** - * Upper-bound on the bytes a single record contributes to a batch, excluding the batch header. - * Used by the incremental strategy, which counts the header once, separately. - */ - public static int recordSizeUpperBound(byte magic, CompressionType compressionType, byte[] key, - byte[] value, Header[] headers) { - ByteBuffer keyBuffer = Utils.wrapNullable(key); - ByteBuffer valueBuffer = Utils.wrapNullable(value); - if (magic >= RecordBatch.MAGIC_VALUE_V2) - return DefaultRecord.recordSizeUpperBound(keyBuffer, valueBuffer, headers); - else if (compressionType != CompressionType.NONE) - return Records.LOG_OVERHEAD + LegacyRecord.recordOverhead(magic) + LegacyRecord.recordSize(magic, keyBuffer, valueBuffer); - else - return Records.LOG_OVERHEAD + LegacyRecord.recordSize(magic, keyBuffer, valueBuffer); - } - /** * Return the size of the record batch header. * diff --git a/clients/src/main/java/org/apache/kafka/common/record/internal/MemoryRecordsBuilder.java b/clients/src/main/java/org/apache/kafka/common/record/internal/MemoryRecordsBuilder.java index 946449ec0ea3f..fc1fd73f4f91b 100644 --- a/clients/src/main/java/org/apache/kafka/common/record/internal/MemoryRecordsBuilder.java +++ b/clients/src/main/java/org/apache/kafka/common/record/internal/MemoryRecordsBuilder.java @@ -830,13 +830,10 @@ private int estimatedBytesWritten() { * Returns the projected number of bytes the builder would write for * {@code uncompressedRecordsSizeInBytes} of records under the given magic, compression type, * and ratio: exact for uncompressed, a ratio-aware estimate for compressed. - *

- * Used by the incremental strategy to size chunk reservations (first record, and mid-batch via - * {@link #estimatedBytesWrittenAfter}). */ - public static int estimatedBytesWritten(byte magic, CompressionType compressionType, - float compressionRatio, - int uncompressedRecordsSizeInBytes) { + private static int estimatedBytesWritten(byte magic, CompressionType compressionType, + float compressionRatio, + int uncompressedRecordsSizeInBytes) { int batchHeaderSizeInBytes = AbstractRecords.recordBatchHeaderSizeInBytes(magic, compressionType); if (compressionType == CompressionType.NONE) { return batchHeaderSizeInBytes + uncompressedRecordsSizeInBytes; From e721a231adc3a1d394341f5a0068612eec68dbc6 Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Fri, 24 Jul 2026 16:31:10 -0400 Subject: [PATCH 48/61] renames --- .../kafka/clients/producer/KafkaProducer.java | 4 ++-- .../producer/internals/BufferPool.java | 20 +++++++++---------- .../internals/ChunkedRecordAccumulator.java | 10 +++++----- .../BufferPoolChunkAllocationTest.java | 4 ++-- .../ChunkedByteBufferOutputStreamTest.java | 2 +- .../ChunkedRecordAccumulatorTest.java | 16 +++++++-------- 6 files changed, 28 insertions(+), 28 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/KafkaProducer.java b/clients/src/main/java/org/apache/kafka/clients/producer/KafkaProducer.java index a660302ae3ff8..6912ce0c56d51 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/KafkaProducer.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/KafkaProducer.java @@ -487,7 +487,7 @@ public KafkaProducer(Properties properties, Serializer keySerializer, Seriali PRODUCER_METRIC_GROUP_NAME, time, transactionManager, - new BufferPool(this.totalMemorySize, ChunkedRecordAccumulator.CHUNK_SIZE, metrics, time, PRODUCER_METRIC_GROUP_NAME, BufferPool.AllocationMode.CHUNKED)); + new BufferPool(this.totalMemorySize, ChunkedRecordAccumulator.CHUNK_SIZE, metrics, time, PRODUCER_METRIC_GROUP_NAME, BufferPool.AllocationMode.INCREMENTAL)); } else { this.accumulator = new RecordAccumulator(logContext, batchSize, @@ -501,7 +501,7 @@ public KafkaProducer(Properties properties, Serializer keySerializer, Seriali PRODUCER_METRIC_GROUP_NAME, time, transactionManager, - new BufferPool(this.totalMemorySize, batchSize, metrics, time, PRODUCER_METRIC_GROUP_NAME, BufferPool.AllocationMode.SINGLE)); + new BufferPool(this.totalMemorySize, batchSize, metrics, time, PRODUCER_METRIC_GROUP_NAME, BufferPool.AllocationMode.FULL)); } this.errors = this.metrics.sensor("errors"); diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/BufferPool.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/BufferPool.java index ba4c523b4fd5f..320b63b4c64da 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/BufferPool.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/BufferPool.java @@ -50,11 +50,11 @@ public class BufferPool { static final String WAIT_TIME_SENSOR_NAME = "bufferpool-wait-time"; /** - * Which allocation method a pool serves: {@link #SINGLE} accepts only {@link #allocate} (the full - * strategy), {@link #CHUNKED} only {@link #allocateChunks} (the incremental strategy). Fixed at + * Which allocation method a pool serves: {@link #FULL} accepts only {@link #allocate} (the full + * strategy), {@link #INCREMENTAL} only {@link #allocateChunks} (the incremental strategy). Fixed at * construction so the two are never mixed on the same pool. */ - public enum AllocationMode { SINGLE, CHUNKED } + public enum AllocationMode { FULL, INCREMENTAL } private final long totalMemory; private final int poolableSize; @@ -82,7 +82,7 @@ public enum AllocationMode { SINGLE, CHUNKED } * @param metricGrpName logical group name for metrics */ public BufferPool(long memory, int poolableSize, Metrics metrics, Time time, String metricGrpName) { - this(memory, poolableSize, metrics, time, metricGrpName, AllocationMode.SINGLE); + this(memory, poolableSize, metrics, time, metricGrpName, AllocationMode.FULL); } /** @@ -134,9 +134,9 @@ public BufferPool(long memory, int poolableSize, Metrics metrics, Time time, Str * forever) */ public ByteBuffer allocate(int size, long maxTimeToBlockMs) throws InterruptedException { - if (allocationMode != AllocationMode.SINGLE) - throw new IllegalStateException("allocate() is not supported on a " + allocationMode - + " buffer pool; use allocateChunks()"); + if (allocationMode != AllocationMode.FULL) + throw new IllegalStateException("allocate() is not supported in " + allocationMode + + " allocation mode; use allocateChunks()"); if (size > this.totalMemory) throw new IllegalArgumentException("Attempt to allocate " + size + " bytes, but there is a hard limit of " @@ -275,9 +275,9 @@ private long awaitMemory(Condition moreMemory, long remainingTimeToBlockNs, * @throws KafkaException if the pool is closed during the wait */ public List allocateChunks(int totalSize, long maxTimeToBlockMs) throws InterruptedException { - if (allocationMode != AllocationMode.CHUNKED) - throw new IllegalStateException("allocateChunks() is not supported on a " + allocationMode - + " buffer pool; use allocate()"); + if (allocationMode != AllocationMode.INCREMENTAL) + throw new IllegalStateException("allocateChunks() is not supported in " + allocationMode + + " allocation mode; use allocate()"); if (totalSize <= 0) throw new IllegalArgumentException("totalSize must be positive: " + totalSize); diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java index c90c120eab76f..f62266633c0ca 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java @@ -244,9 +244,9 @@ public RecordAppendResult append(String topic, throw new IllegalStateException("needsNewBatch path reached without an allocated buffer stream"); // Reuse the new-batch size estimate as the write-limit basis. // TODO: review when compression is supported. - final NewBatchBuffer pending = newBatch; + final NewBatchBuffer pendingNewBatch = newBatch; appendResult = appendNewBatch(topic, effectivePartition, dq, timestamp, key, value, headers, callbacks, - () -> chunkedRecordsBuilder(pending.stream, pending.size), nowMs); + () -> chunkedRecordsBuilder(pendingNewBatch.stream, pendingNewBatch.firstAppendSize), nowMs); if (appendResult.needsNewBatch()) throw new IllegalStateException("appendNewBatch must not return a needsNewBatch result"); if (appendResult.needsBufferExtension()) { @@ -340,11 +340,11 @@ private MemoryRecordsBuilder chunkedRecordsBuilder(ChunkedByteBufferOutputStream */ private static final class NewBatchBuffer { final ChunkedByteBufferOutputStream stream; - final int size; + final int firstAppendSize; - NewBatchBuffer(ChunkedByteBufferOutputStream stream, int size) { + NewBatchBuffer(ChunkedByteBufferOutputStream stream, int firstAppendSize) { this.stream = stream; - this.size = size; + this.firstAppendSize = firstAppendSize; } } } diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/BufferPoolChunkAllocationTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/BufferPoolChunkAllocationTest.java index bd6da6d0b1641..53e77b55a440e 100644 --- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/BufferPoolChunkAllocationTest.java +++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/BufferPoolChunkAllocationTest.java @@ -55,7 +55,7 @@ public void teardown() { private BufferPool pool(long totalMemory, int chunkSize) { String metricGroup = "producer-metrics"; - return new BufferPool(totalMemory, chunkSize, metrics, time, metricGroup, BufferPool.AllocationMode.CHUNKED); + return new BufferPool(totalMemory, chunkSize, metrics, time, metricGroup, BufferPool.AllocationMode.INCREMENTAL); } /** Single-chunk request returns a list of one buffer at the chunk size. */ @@ -342,7 +342,7 @@ public void testWaitingRequestDoesNotCorruptAccountingOnSuccess() throws Excepti } catch (Throwable th) { err.set(th); } - }, "slow-success"); + }, "chunked-waiter"); t.start(); TestUtils.waitForCondition(() -> p.queued() == 1, "waiter should be parked on the pool's queue"); diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java index a836e9e0c5a32..568afe1a117dc 100644 --- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java +++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java @@ -44,7 +44,7 @@ public void teardown() { private BufferPool pool(long total, int chunkSize) { String metricGroup = "test"; - return new BufferPool(total, chunkSize, metrics, time, metricGroup, BufferPool.AllocationMode.CHUNKED); + return new BufferPool(total, chunkSize, metrics, time, metricGroup, BufferPool.AllocationMode.INCREMENTAL); } private List chunks(BufferPool pool, int chunkSize, int count) throws InterruptedException { diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java index 1b688f30f397d..2c2c0d7b4162f 100644 --- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java +++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java @@ -82,7 +82,7 @@ public void teardown() { } private ChunkedRecordAccumulator newAccumulator(int batchSize, int chunkSize, long totalMemory, Compression compression) { - BufferPool pool = new BufferPool(totalMemory, chunkSize, metrics, time, "producer-metrics", BufferPool.AllocationMode.CHUNKED); + BufferPool pool = new BufferPool(totalMemory, chunkSize, metrics, time, "producer-metrics", BufferPool.AllocationMode.INCREMENTAL); return new ChunkedRecordAccumulator(logContext, batchSize, compression, /* lingerMs */ 0, /* retryBackoffMs */ 0L, /* retryBackoffMaxMs */ 0L, /* deliveryTimeoutMs */ 3200, metrics, "producer-metrics", time, @@ -161,7 +161,7 @@ public void testMidBatchExtensionGrowsExistingBatch() throws Exception { private BufferPool poolMockingConcurrentChunkAllocation(int chunkSize, long totalMemory, AtomicReference injectAppendOnce, byte[] injectedValue) { - return new BufferPool(totalMemory, chunkSize, metrics, time, "producer-metrics", BufferPool.AllocationMode.CHUNKED) { + return new BufferPool(totalMemory, chunkSize, metrics, time, "producer-metrics", BufferPool.AllocationMode.INCREMENTAL) { @Override public List allocateChunks(int totalSize, long maxTimeToBlockMs) throws InterruptedException { List chunks = super.allocateChunks(totalSize, maxTimeToBlockMs); @@ -234,7 +234,7 @@ public void testConcurrentExtensionRaceLoserStartsNewBatch() throws Exception { final AtomicReference injectAppendOnce = new AtomicReference<>(); final AtomicInteger chunkDeallocations = new AtomicInteger(); - BufferPool pool = new BufferPool(64L * chunkSize, chunkSize, metrics, time, "producer-metrics", BufferPool.AllocationMode.CHUNKED) { + BufferPool pool = new BufferPool(64L * chunkSize, chunkSize, metrics, time, "producer-metrics", BufferPool.AllocationMode.INCREMENTAL) { @Override public List allocateChunks(int totalSize, long maxTimeToBlockMs) throws InterruptedException { List chunks = super.allocateChunks(totalSize, maxTimeToBlockMs); @@ -298,7 +298,7 @@ public void testPartitionSwitchRefundsHeldExtensionChunks() throws Exception { AtomicBoolean extensionAllocated = new AtomicBoolean(false); AtomicInteger chunkDeallocations = new AtomicInteger(0); - BufferPool pool = new BufferPool(totalMemory, chunkSize, metrics, time, "producer-metrics", BufferPool.AllocationMode.CHUNKED) { + BufferPool pool = new BufferPool(totalMemory, chunkSize, metrics, time, "producer-metrics", BufferPool.AllocationMode.INCREMENTAL) { @Override public List allocateChunks(int totalSize, long maxTimeToBlockMs) throws InterruptedException { List chunks = super.allocateChunks(totalSize, maxTimeToBlockMs); @@ -363,7 +363,7 @@ protected boolean partitionChanged(String topic, TopicInfo topicInfo, public void testInflightExpirationReturnsAllChunksToPool() throws Exception { int chunkSize = 128; long totalMemory = 32L * chunkSize; - BufferPool pool = new BufferPool(totalMemory, chunkSize, metrics, time, "producer-metrics", BufferPool.AllocationMode.CHUNKED); + BufferPool pool = new BufferPool(totalMemory, chunkSize, metrics, time, "producer-metrics", BufferPool.AllocationMode.INCREMENTAL); ChunkedRecordAccumulator accum = new ChunkedRecordAccumulator(logContext, 8192, Compression.NONE, /* lingerMs */ 0, /* retryBackoffMs */ 0L, /* retryBackoffMaxMs */ 0L, /* deliveryTimeoutMs */ 3200, metrics, "producer-metrics", time, @@ -415,7 +415,7 @@ public void testExhaustedExtensionFallsBackToBlockingNewBatchPath() throws Excep AtomicInteger closeForAppendsCalls = new AtomicInteger(); List closeCallsAtAlloc = new ArrayList<>(); - BufferPool pool = new BufferPool(16L * chunkSize, chunkSize, metrics, time, "producer-metrics", BufferPool.AllocationMode.CHUNKED) { + BufferPool pool = new BufferPool(16L * chunkSize, chunkSize, metrics, time, "producer-metrics", BufferPool.AllocationMode.INCREMENTAL) { @Override public List allocateChunks(int totalSize, long maxTimeToBlockMs) throws InterruptedException { allocTimeouts.add(maxTimeToBlockMs); @@ -517,7 +517,7 @@ public void testBufferExhaustedNotDoubleCountedAcrossExtensionAndNewBatch() thro public void testBatchCloseDoesNotDeallocateChunksPrematurely() throws Exception { int chunkSize = 256; long totalMemory = 32L * chunkSize; - BufferPool pool = new BufferPool(totalMemory, chunkSize, metrics, time, "producer-metrics", BufferPool.AllocationMode.CHUNKED); + BufferPool pool = new BufferPool(totalMemory, chunkSize, metrics, time, "producer-metrics", BufferPool.AllocationMode.INCREMENTAL); ChunkedRecordAccumulator accum = new ChunkedRecordAccumulator(logContext, 8192, Compression.NONE, /* lingerMs */ 0, /* retryBackoffMs */ 0L, /* retryBackoffMaxMs */ 0L, /* deliveryTimeoutMs */ 3200, metrics, "producer-metrics", time, @@ -601,7 +601,7 @@ public void testCumulativeAccountsForBatchHeaderOnce() throws Exception { int chunkSize = 256; int batchSize = 8192; long totalMemory = 64L * chunkSize; - BufferPool pool = new BufferPool(totalMemory, chunkSize, metrics, time, "producer-metrics", BufferPool.AllocationMode.CHUNKED); + BufferPool pool = new BufferPool(totalMemory, chunkSize, metrics, time, "producer-metrics", BufferPool.AllocationMode.INCREMENTAL); ChunkedRecordAccumulator accum = new ChunkedRecordAccumulator(logContext, batchSize, Compression.NONE, /* lingerMs */ 0, /* retryBackoffMs */ 0L, /* retryBackoffMaxMs */ 0L, /* deliveryTimeoutMs */ 3200, metrics, "producer-metrics", time, From 1de3d7657abb4daaaec63d7118607adfee87184e Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Fri, 24 Jul 2026 17:02:54 -0400 Subject: [PATCH 49/61] improve test timeouts --- .../producer/internals/BufferPoolChunkAllocationTest.java | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/BufferPoolChunkAllocationTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/BufferPoolChunkAllocationTest.java index 53e77b55a440e..13186efa31dd7 100644 --- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/BufferPoolChunkAllocationTest.java +++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/BufferPoolChunkAllocationTest.java @@ -241,9 +241,8 @@ public void testFifoFairnessAcrossMultiChunkAndSingleChunkRequests() throws Exce // Wait for tSingle to also be parked. TestUtils.waitForCondition(() -> p.queued() == 2, "single-chunk waiter joined after the multi-chunk one"); - // Now free chunks one at a time, allowing the multi-chunk request to accumulate. + // Free the two chunks the multi request needs. p.deallocate(h1); - Thread.sleep(50); p.deallocate(h2); // Both freed — the FIFO leader (multi) should claim both before the single-chunk waiter // gets any. The multi completes first. @@ -302,8 +301,9 @@ public void testWaitingRequestDoesNotDoubleRefundChunksOnTimeout() throws Except p.deallocate(h1); p.deallocate(h2); - // Give the waiter a moment to wake, take the 2 chunks, and re-enter the wait before it times out. - Thread.sleep(100); + // Rely on availableMemory being 0 to confirm the waiter has taken both freed chunks. + TestUtils.waitForCondition(() -> p.availableMemory() == 0, + "waiter should have taken both freed chunks"); t.join(); assertInstanceOf(BufferExhaustedException.class, err.get(), "expected BufferExhaustedException, got " + err.get()); From db14370228abbf5f5f506b3004f46cc908cb49fd Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Wed, 29 Jul 2026 09:45:32 -0400 Subject: [PATCH 50/61] log warn if fallback --- .../org/apache/kafka/clients/producer/KafkaProducer.java | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/KafkaProducer.java b/clients/src/main/java/org/apache/kafka/clients/producer/KafkaProducer.java index 6912ce0c56d51..9230c2e2841df 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/KafkaProducer.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/KafkaProducer.java @@ -467,6 +467,14 @@ public KafkaProducer(Properties properties, Serializer keySerializer, Seriali // (batch.size >= CHUNK_SIZE). Below that, a batch can't fill even one chunk, so chunking // would over-reserve and the producer falls back to the full strategy instead. boolean useIncremental = incremental && batchSize >= ChunkedRecordAccumulator.CHUNK_SIZE; + if (incremental && !useIncremental) { + log.warn("Ignoring {}={} and falling back to {}: {} is {} bytes, below the {} byte chunk size, " + + "so a batch cannot fill a single chunk.", + ProducerConfig.BUFFER_MEMORY_ALLOCATION_STRATEGY_CONFIG, + ProducerConfig.BUFFER_MEMORY_ALLOCATION_STRATEGY_INCREMENTAL, + ProducerConfig.BUFFER_MEMORY_ALLOCATION_STRATEGY_FULL, + ProducerConfig.BATCH_SIZE_CONFIG, batchSize, ChunkedRecordAccumulator.CHUNK_SIZE); + } // The chunked path does not support compression yet (TODO: KAFKA-20579) if (useIncremental && compression.type() != CompressionType.NONE) { throw new ConfigException("The " + ProducerConfig.BUFFER_MEMORY_ALLOCATION_STRATEGY_INCREMENTAL From ad3fa0d7cb1916f5c416554bf5c04043d264346f Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Wed, 29 Jul 2026 10:02:08 -0400 Subject: [PATCH 51/61] avoid waste when advancing chunks --- .../ChunkedByteBufferOutputStream.java | 23 +++++++++++-------- 1 file changed, 14 insertions(+), 9 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java index 63606839dbe1c..975da27c841a8 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java @@ -97,7 +97,7 @@ private static void validateChunkCapacities(List chunks, int chunkSi public void write(int b) { ensureNotDeallocated(); ensureWritable(); - ensureChunkCapacity(1); + advanceWhileCurrentChunkFull(); currentChunk.put((byte) b); } @@ -106,7 +106,7 @@ public void write(byte[] bytes, int off, int len) { ensureNotDeallocated(); ensureWritable(); while (len > 0) { - ensureChunkCapacity(1); + advanceWhileCurrentChunkFull(); int toWrite = Math.min(len, currentChunk.remaining()); currentChunk.put(bytes, off, toWrite); off += toWrite; @@ -119,7 +119,7 @@ public void write(ByteBuffer sourceBuffer) { ensureNotDeallocated(); ensureWritable(); while (sourceBuffer.hasRemaining()) { - ensureChunkCapacity(1); + advanceWhileCurrentChunkFull(); int toWrite = Math.min(sourceBuffer.remaining(), currentChunk.remaining()); int oldLimit = sourceBuffer.limit(); sourceBuffer.limit(sourceBuffer.position() + toWrite); @@ -145,8 +145,11 @@ private void ensureNotDeallocated() { throw new IllegalStateException("operation not allowed after the stream has been deallocated"); } - private void ensureChunkCapacity(int needed) { - while (currentChunk.remaining() < needed) { + /** + * Makes room for the next write by advancing past the chunks that are already full. + */ + private void advanceWhileCurrentChunkFull() { + while (!currentChunk.hasRemaining()) { advanceToNextChunk(); } } @@ -322,10 +325,12 @@ public int initialCapacity() { @Override public void ensureRemaining(int remainingBytesRequired) { ensureNotDeallocated(); - // A single call can guarantee at most `chunkSize` of space (the stream advances one chunk - // at a time). Callers needing more attach chunks via addBuffers first. write(byte[]) loops - // across chunks, so contiguous capacity isn't required. - ensureChunkCapacity(Math.min(remainingBytesRequired, chunkSize)); + // A single write can be split across several chunks, so the required bytes needn't be + // contiguous: only the total free space matters. Advancing here would waste the tail of the + // current chunk, so writes advance lazily and this only validates. TODO: KAFKA-20579, grow here. + if (remainingBytesRequired > remaining()) + throw new IllegalStateException("required " + remainingBytesRequired + + " bytes but only " + remaining() + " remaining across the attached chunks"); } /** From 0cc4b838817b205e1f24bf524736273ca1cb9490 Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Wed, 29 Jul 2026 10:12:39 -0400 Subject: [PATCH 52/61] validation & doc --- .../producer/internals/ChunkedProducerBatch.java | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedProducerBatch.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedProducerBatch.java index f32ea107b95eb..a95d8d2e726e3 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedProducerBatch.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedProducerBatch.java @@ -37,6 +37,10 @@ public class ChunkedProducerBatch extends ProducerBatch { public ChunkedProducerBatch(TopicPartition tp, MemoryRecordsBuilder recordsBuilder, long createdMs) { super(tp, recordsBuilder, createdMs); + if (!(recordsBuilder.bufferStream() instanceof ChunkedByteBufferOutputStream)) + throw new IllegalArgumentException("recordsBuilder must be an instance of " + + ChunkedByteBufferOutputStream.class.getSimpleName() + ", but found " + + recordsBuilder.bufferStream().getClass().getName()); } /** @@ -58,13 +62,13 @@ int extensionBytesNeeded(long timestamp, byte[] key, byte[] value, Header[] head } /** - * Appends the record after checking there's chunk capacity for it. - * At this point it's expected that the allocated chunks have - * capacity for the record because the accumulator never routes an - * append here without ensuring chunk capacity first: the first-record path pre-sizes the - * stream for the record and the mid-batch path attaches extension chunks before retrying. + * Appends the record, first checking that the stream has the capacity to hold it. This batch + * never acquires memory itself: the capacity is arranged before the append, by the accumulator + * attaching chunks (see {@link #extensionBytesNeeded} and {@link #addBuffers}). * - * @throws IllegalStateException if the attached chunks lack capacity for the record + * @return the record's future, or null if the batch is at its batch-size limit + * @throws IllegalStateException if the batch could still take the record but the stream lacks + * the capacity to hold it */ @Override public FutureRecordMetadata tryAppend(long timestamp, byte[] key, byte[] value, Header[] headers, Callback callback, long now) { From af1e8eed965468c594fa89ffa0871fe357c726a3 Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Wed, 29 Jul 2026 10:18:40 -0400 Subject: [PATCH 53/61] misc updates --- .../producer/internals/BufferPool.java | 22 +++++++++---------- gradle/spotbugs-exclude.xml | 1 - 2 files changed, 11 insertions(+), 12 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/BufferPool.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/BufferPool.java index 320b63b4c64da..312e0bd962f52 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/BufferPool.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/BufferPool.java @@ -59,17 +59,17 @@ public enum AllocationMode { FULL, INCREMENTAL } private final long totalMemory; private final int poolableSize; /** Lock held for any read or write of {@link #free}, {@link #waiters}, {@link #nonPooledAvailableMemory}, or {@link #closed}. */ - protected final ReentrantLock lock; + private final ReentrantLock lock; /** Pooled buffers of capacity {@link #poolableSize}, available for reuse. Guarded by {@link #lock}. */ - protected final Deque free; + private final Deque free; /** FIFO queue of pending allocation requests; the longest-waiting thread is woken first. Guarded by {@link #lock}. */ - protected final Deque waiters; + private final Deque waiters; /** Total available memory is the sum of nonPooledAvailableMemory and the number of byte buffers in free * poolableSize. Guarded by {@link #lock}. */ - protected long nonPooledAvailableMemory; + private long nonPooledAvailableMemory; private final Metrics metrics; - protected final Time time; + private final Time time; private final Sensor waitTime; - protected boolean closed; + private boolean closed; private final AllocationMode allocationMode; /** @@ -408,10 +408,10 @@ protected void recordWaitTime(long timeNs) { /** * Record that a record send was dropped because the buffer pool was exhausted. Shared by the - * full strategy (allocate) and the incremental strategy (ChunkedRecordAccumulator), + * full strategy (allocate) and the incremental strategy ({@link ChunkedRecordAccumulator}), * so both update the same buffer-exhausted metrics. */ - protected void recordBufferExhausted() { + void recordBufferExhausted() { this.metrics.sensor("buffer-exhausted-records").record(); } @@ -419,7 +419,7 @@ protected void recordBufferExhausted() { * Wake the longest-waiting thread if any memory (pooled or non-pooled) is available. * Must be called with {@link #lock} held. No-op if no waiters or no memory is free. */ - protected void signalNextWaiterIfMemoryAvailable() { + private void signalNextWaiterIfMemoryAvailable() { if (!(this.nonPooledAvailableMemory == 0 && this.free.isEmpty()) && !this.waiters.isEmpty()) this.waiters.peekFirst().signal(); } @@ -445,7 +445,7 @@ private ByteBuffer safeAllocateByteBuffer(int size) { * waiter. Acquires {@link #lock} internally. Used by callers * that reserve memory and then need to roll back the reservation (e.g., upon errors). */ - protected void releaseReservedBytes(long bytes) { + private void releaseReservedBytes(long bytes) { this.lock.lock(); try { this.nonPooledAvailableMemory += bytes; @@ -465,7 +465,7 @@ protected ByteBuffer allocateByteBuffer(int size) { * Attempt to ensure we have at least the requested number of bytes of memory for allocation by deallocating pooled * buffers (if needed). Must be called with {@link #lock} held. */ - protected void freeUp(int size) { + private void freeUp(int size) { while (!this.free.isEmpty() && this.nonPooledAvailableMemory < size) this.nonPooledAvailableMemory += this.free.pollLast().capacity(); } diff --git a/gradle/spotbugs-exclude.xml b/gradle/spotbugs-exclude.xml index c7bf9a6464f20..00b51f8b71dee 100644 --- a/gradle/spotbugs-exclude.xml +++ b/gradle/spotbugs-exclude.xml @@ -564,7 +564,6 @@ For a detailed description of spotbugs bug categories, see https://spotbugs.read --> - From 5d80f9622f636026cc732a01afeb7f1c6725c5cf Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Wed, 29 Jul 2026 10:24:09 -0400 Subject: [PATCH 54/61] more tests --- .../ChunkedByteBufferOutputStreamTest.java | 20 +++++++++++++++++++ .../ChunkedRecordAccumulatorTest.java | 9 +++++++++ 2 files changed, 29 insertions(+) diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java index 568afe1a117dc..db39b8279d858 100644 --- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java +++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStreamTest.java @@ -181,6 +181,26 @@ public void testRemainingSumsFreeBytesAcrossChunks() throws Exception { } } + @Test + public void testEnsureRemainingDoesNotWasteCurrentChunk() throws Exception { + int chunkSize = 8; + BufferPool p = pool(64, chunkSize); + try (ChunkedByteBufferOutputStream stream = new ChunkedByteBufferOutputStream(chunks(p, chunkSize, 2), chunkSize, p)) { + stream.write(new byte[]{1, 2, 3}, 0, 3); + + // Requesting more than the current chunk's free bytes (5) but no more than the total + // free bytes across chunks (13) must not skip ahead: the bytes left in the current chunk + // stay writable. + stream.ensureRemaining(chunkSize + 1); + assertEquals(2 * chunkSize - 3, stream.remaining()); + stream.write(new byte[]{4, 5}, 0, 2); + assertEquals(5, stream.position()); + assertEquals(2 * chunkSize - 5, stream.remaining()); + + stream.deallocate(); + } + } + @Test public void testAddBuffersExtendsStream() throws Exception { int chunkSize = 8; diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java index 2c2c0d7b4162f..67d67a5bc72e6 100644 --- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java +++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java @@ -25,6 +25,7 @@ import org.apache.kafka.common.metrics.KafkaMetric; import org.apache.kafka.common.metrics.Metrics; import org.apache.kafka.common.protocol.Errors; +import org.apache.kafka.common.record.TimestampType; import org.apache.kafka.common.record.internal.MemoryRecords; import org.apache.kafka.common.record.internal.MemoryRecordsBuilder; import org.apache.kafka.common.record.internal.Record; @@ -626,4 +627,12 @@ public void testCumulativeAccountsForBatchHeaderOnce() throws Exception { accum.close(); } + + @Test + public void testChunkedBatchRejectsNonChunkedStream() { + MemoryRecordsBuilder plainBuilder = MemoryRecords.builder(ByteBuffer.allocate(256), + RecordBatch.CURRENT_MAGIC_VALUE, Compression.NONE, TimestampType.CREATE_TIME, 0L); + assertThrows(IllegalArgumentException.class, + () -> new ChunkedProducerBatch(tp1, plainBuilder, time.milliseconds())); + } } From 3554d74e729e53b98ee71230fad7d21a68d8936c Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Wed, 29 Jul 2026 10:48:07 -0400 Subject: [PATCH 55/61] track remainingTimeToBlock --- .../internals/ChunkedRecordAccumulator.java | 68 ++++++++++++------- 1 file changed, 42 insertions(+), 26 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java index f62266633c0ca..4a2479531e07b 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java @@ -118,6 +118,11 @@ public RecordAppendResult append(String topic, // with that size. Set and cleared together across retries; null when none is held. NewBatchBuffer newBatch = null; List 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; if (headers == null) headers = Record.EMPTY_HEADERS; try { while (true) { @@ -149,27 +154,10 @@ public RecordAppendResult append(String topic, } if (appendResult.needsBufferExtension()) { - // Mid-batch extension: the open batch can still take this record so grow it in - // place. The acquire is non-blocking to fail fast when the pool is exhausted: - // close the batch and let the record block once on the new-batch path. - // A blocking call would lead to the same outcome, but would block once here - // and still need a second blocking call to start a new batch anyways - // (a first blocking call here would make all open batches drainable, including this one, - // so most probably our batch would be gone/drained by the time memory is returned to the pool, - // and we would need a new batch for our record anyways). - try { - extensionChunks = chunkedFree.allocateChunks(appendResult.extensionBytesNeeded, 0L); - } catch (BufferExhaustedException e) { - log.trace("Pool exhausted while extending batch for topic {} partition {}; closing existing batch", - topic, effectivePartition); - synchronized (dq) { - ProducerBatch last = dq.peekLast(); - if (last != null && last.isWritable()) { - last.closeForRecordAppends(); - } - } - // Continue to the next iteration that should block to start a new batch (needsNewBatch), - // given that this one has been closed for appends. + extensionChunks = allocateExtensionChunks(appendResult.extensionBytesNeeded, dq, topic, effectivePartition); + if (extensionChunks == null) { + // Pool exhausted, so no writable batch is left to extend: retry, normally + // landing on the blocking new-batch path (needsNewBatch). continue; } nowMs = time.milliseconds(); @@ -182,17 +170,21 @@ public RecordAppendResult append(String topic, int newBatchSize = AbstractRecords.estimateSizeInBytesUpperBound( RecordBatch.CURRENT_MAGIC_VALUE, compression.type(), key, value, headers); log.trace("Allocating {} byte chunked buffer ({} byte chunks) for topic {} partition {} with remaining timeout {}ms", - newBatchSize, chunkedFree.poolableSize(), topic, effectivePartition, maxTimeToBlock); + newBatchSize, chunkedFree.poolableSize(), topic, effectivePartition, remainingTimeToBlock); List initialChunks; + long allocationStartMs = time.milliseconds(); try { - initialChunks = chunkedFree.allocateChunks(newBatchSize, maxTimeToBlock); + initialChunks = chunkedFree.allocateChunks(newBatchSize, remainingTimeToBlock); } catch (BufferExhaustedException e) { - // The blocking new-batch acquire was not able to get memory within - // max.block.ms. Record it in the buffer-exhausted metrics. + // The blocking new-batch acquire was not able to get memory within what is + // left of max.block.ms. Record it in the buffer-exhausted metrics. chunkedFree.recordBufferExhausted(); throw e; + } finally { + // Update the remaining time to block. + nowMs = time.milliseconds(); + remainingTimeToBlock = Math.max(0L, remainingTimeToBlock - (nowMs - allocationStartMs)); } - nowMs = time.milliseconds(); newBatch = new NewBatchBuffer( new ChunkedByteBufferOutputStream(initialChunks, chunkedFree.poolableSize(), chunkedFree), newBatchSize); @@ -271,6 +263,30 @@ public RecordAppendResult append(String topic, } } + /** + * 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. + * + * @return the chunks, or null if the pool was exhausted + */ + private List allocateExtensionChunks(int extensionBytesNeeded, Deque dq, + String topic, int partition) throws InterruptedException { + try { + return chunkedFree.allocateChunks(extensionBytesNeeded, 0L); + } catch (BufferExhaustedException e) { + log.trace("Pool exhausted while extending batch for topic {} partition {}; closing existing batch", + topic, partition); + synchronized (dq) { + ProducerBatch last = dq.peekLast(); + if (last != null && last.isWritable()) { + last.closeForRecordAppends(); + } + } + return null; + } + } + /** * Return any held extension chunks to the pool. No-op when none are held. */ From 198fe38a77dabf908a81b32085940125dd256f33 Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Wed, 29 Jul 2026 16:10:01 -0400 Subject: [PATCH 56/61] addressing comments --- .../producer/internals/BufferPool.java | 8 +++++++ .../internals/ChunkedProducerBatch.java | 22 +++++++++++++------ .../internals/ChunkedRecordAccumulator.java | 6 ++++- 3 files changed, 28 insertions(+), 8 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/BufferPool.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/BufferPool.java index 312e0bd962f52..3d0df7b54ce30 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/BufferPool.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/BufferPool.java @@ -555,6 +555,14 @@ public long totalMemory() { return this.totalMemory; } + /** + * Which allocation method this pool serves, so callers that only use one of them can validate + * the pool they were given up front. + */ + AllocationMode allocationMode() { + return this.allocationMode; + } + // package-private method used only for testing Deque waiters() { return this.waiters; diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedProducerBatch.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedProducerBatch.java index a95d8d2e726e3..ffa0d47896fd5 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedProducerBatch.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedProducerBatch.java @@ -50,6 +50,8 @@ public ChunkedProducerBatch(TopicPartition tp, MemoryRecordsBuilder recordsBuild * is pre-sized for the first record). Positive when the record is within the batch-size limit * but the attached chunks lack capacity — the accumulator then allocates exactly the missing * bytes (rounded up to whole chunks) and attaches them via {@link #addBuffers} before retrying. + *

+ * TODO (KAFKA-20859): improve by reusing size calculation. */ int extensionBytesNeeded(long timestamp, byte[] key, byte[] value, Header[] headers) { if (!recordsBuilder.hasRoomFor(timestamp, key, value, headers)) @@ -62,20 +64,26 @@ int extensionBytesNeeded(long timestamp, byte[] key, byte[] value, Header[] head } /** - * Appends the record, first checking that the stream has the capacity to hold it. This batch - * never acquires memory itself: the capacity is arranged before the append, by the accumulator - * attaching chunks (see {@link #extensionBytesNeeded} and {@link #addBuffers}). + * Appends the record. This batch never acquires memory itself: the capacity is arranged before + * the append, by the accumulator attaching chunks (see {@link #extensionBytesNeeded} and + * {@link #addBuffers}). + *

+ * The capacity is verified here only for the batch's first record, which is the one append no + * caller checks: {@code RecordAccumulator.appendNewBatch} creates the batch and appends straight + * into it, relying on the stream having been pre-sized. Every later append arrives through + * {@code ChunkedRecordAccumulator.tryAppend}, which evaluates {@link #extensionBytesNeeded} + * itself and attaches chunks before retrying, so repeating the check for those would size the + * record a second time on every append for no added safety. * * @return the record's future, or null if the batch is at its batch-size limit - * @throws IllegalStateException if the batch could still take the record but the stream lacks - * the capacity to hold it + * @throws IllegalStateException if the stream was not pre-sized to hold the batch's first record */ @Override public FutureRecordMetadata tryAppend(long timestamp, byte[] key, byte[] value, Header[] headers, Callback callback, long now) { - if (extensionBytesNeeded(timestamp, key, value, headers) != 0) + if (recordCount == 0 && extensionBytesNeeded(timestamp, key, value, headers) != 0) throw new IllegalStateException( "Unexpected append to a chunked batch whose chunks lack capacity for the record; " + - "the accumulator should have attached extension chunks first"); + "the stream should have been pre-sized for the batch's first record"); return super.tryAppend(timestamp, key, value, headers, callback, now); } diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java index 4a2479531e07b..9434e317f3dbf 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java @@ -75,6 +75,10 @@ public ChunkedRecordAccumulator(LogContext logContext, BufferPool bufferPool) { super(logContext, batchSize, compression, lingerMs, retryBackoffMs, retryBackoffMaxMs, deliveryTimeoutMs, partitionerConfig, metrics, metricGrpName, time, transactionManager, bufferPool); + if (bufferPool.allocationMode() != BufferPool.AllocationMode.INCREMENTAL) + throw new IllegalArgumentException("bufferPool must serve " + + BufferPool.AllocationMode.INCREMENTAL + " allocation, but serves " + + bufferPool.allocationMode()); // TODO: drop this once the incremental strategy supports compressed data (with the // mid-record growth fallback for compressor overshoot). if (compression.type() != CompressionType.NONE) @@ -183,7 +187,7 @@ public RecordAppendResult append(String topic, } finally { // Update the remaining time to block. nowMs = time.milliseconds(); - remainingTimeToBlock = Math.max(0L, remainingTimeToBlock - (nowMs - allocationStartMs)); + remainingTimeToBlock = Math.max(0L, remainingTimeToBlock - Math.max(0L, nowMs - allocationStartMs)); } newBatch = new NewBatchBuffer( new ChunkedByteBufferOutputStream(initialChunks, chunkedFree.poolableSize(), chunkedFree), From 4b80262c3d45749d335f9395de782779d5c48400 Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Wed, 29 Jul 2026 16:16:30 -0400 Subject: [PATCH 57/61] update test comments --- .../internals/BufferPoolChunkAllocationTest.java | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/BufferPoolChunkAllocationTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/BufferPoolChunkAllocationTest.java index 13186efa31dd7..2a6d7d70db9a7 100644 --- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/BufferPoolChunkAllocationTest.java +++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/BufferPoolChunkAllocationTest.java @@ -125,22 +125,23 @@ public void testRejectsNonPositiveRequest() { } /** - * When a multi-chunk request can't be fully satisfied within the deadline, chunks already - * acquired during the call must be returned to the pool. + * A request that cannot be satisfied immediately and has no time to wait takes nothing: it + * blocks on the wait queue before acquiring anything, so the timeout leaves pool memory + * untouched (no roll back needed). */ @Test - public void testRollbackOnPartialFailure() throws Exception { + public void testImmediateTimeoutAcquiresNothing() throws Exception { int chunkSize = 64; long total = 2 * chunkSize; // only 2 chunks worth of memory BufferPool p = pool(total, chunkSize); // Reserve one chunk so the pool has only 1 left. ByteBuffer held = p.allocateChunks(chunkSize, 100).get(0); - // Request 2 chunks with a zero deadline — first chunk fits, second can't, must throw. + // Request 2 chunks with a zero deadline. Only 1 chunk's worth is free, so the request goes + // to the wait queue and times out on its first wait, before taking anything. assertThrows(BufferExhaustedException.class, () -> p.allocateChunks(2 * chunkSize, 0)); - // Available memory must reflect only the chunk we deliberately hold; the request's - // first-chunk acquisition was rolled back. + // Available memory reflects only the chunk we deliberately hold. assertEquals(total - chunkSize, p.availableMemory()); p.deallocate(held); From d6a2fd624318adbbd8971db05437687290c1f8bd Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Wed, 29 Jul 2026 17:19:18 -0400 Subject: [PATCH 58/61] fix replacing isWritable with extensionBytesNeeded & cleanup --- .../internals/ChunkedProducerBatch.java | 6 ++--- .../internals/ChunkedRecordAccumulator.java | 25 ++++++++++--------- .../producer/internals/ProducerBatch.java | 4 --- 3 files changed, 16 insertions(+), 19 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedProducerBatch.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedProducerBatch.java index ffa0d47896fd5..6f4ffe2de4bab 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedProducerBatch.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedProducerBatch.java @@ -91,9 +91,9 @@ public FutureRecordMetadata tryAppend(long timestamp, byte[] key, byte[] value, * Attach pre-allocated chunks to this batch's stream so the next {@code tryAppend} can * spill into them. Ownership of the chunks transfers to the stream. *

- * Concurrent appenders may over-attach (each sizes its own gap against the same remaining - * capacity), attaching more than the batch-size limit can use; chunks left fully unused are - * returned to the pool when the batch closes. + * The accumulator calls this only while {@link #extensionBytesNeeded} is positive, under the same + * deque lock, so the chunks are still needed when attached. Any never written to are returned to + * the pool when the batch closes for appends. */ void addBuffers(List chunks) { stream().addBuffers(chunks); diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java index 9434e317f3dbf..8d9d2bb8945d7 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java @@ -207,16 +207,16 @@ public RecordAppendResult append(String topic, if (extensionChunks != null) { ProducerBatch last = dq.peekLast(); - // The off-lock allocateChunks window allows the open batch we checked to be - // drained and replaced — possibly by a split batch (a plain - // ProducerBatch), which can't take extension chunks. Only attach to a - // writable chunked batch; otherwise refund the chunks and re-evaluate. - if (last instanceof ChunkedProducerBatch && last.isWritable()) { - // Attach the chunks we sized off-lock without re-checking whether a - // concurrent appender already grew the batch enough. Optimizes for the - // uncontended case; under concurrency two appenders can each attach their - // own gap (temporary over-allocation, one could have been enough for both). - // The unused chunks returns to the pool when the batch closes for appends. + // The batch may have changed while allocateChunks was off-lock: drained and + // replaced, closed for appends, filled to its limit, or already grown by a + // concurrent appender. The instanceof excludes a replacement that cannot take + // chunks at all (a split batch is a plain ProducerBatch), and a positive + // extensionBytesNeeded means the batch still needs chunks for this record, so + // together they attach only when doing so is both safe and useful: attaching + // to a batch closed for appends would throw, and to one that no longer needs + // them would merely hold chunks it cannot use. + if (last instanceof ChunkedProducerBatch + && ((ChunkedProducerBatch) last).extensionBytesNeeded(timestamp, key, value, headers) > 0) { ((ChunkedProducerBatch) last).addBuffers(extensionChunks); extensionChunks = null; RecordAppendResult retryResult = tryAppend(timestamp, key, value, headers, callbacks, dq, nowMs); @@ -228,7 +228,7 @@ public RecordAppendResult append(String topic, // right: needsBufferExtension with a fresh gap, or needsNewBatch continue; } - // The open batch is gone, closed, or non-chunked (e.g., a split batch). Return chunks to pool. + // The batch no longer needs these chunks, or cannot take them. Return them to the pool. deallocateExtensionChunks(extensionChunks); extensionChunks = null; continue; @@ -283,7 +283,8 @@ private List allocateExtensionChunks(int extensionBytesNeeded, Deque topic, partition); synchronized (dq) { ProducerBatch last = dq.peekLast(); - if (last != null && last.isWritable()) { + // No need to check whether it is still open: closeForRecordAppends is idempotent. + if (last != null) { last.closeForRecordAppends(); } } diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ProducerBatch.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ProducerBatch.java index bad4c81ccc545..95e8874a1b167 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ProducerBatch.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ProducerBatch.java @@ -567,10 +567,6 @@ public int initialCapacity() { return recordsBuilder.initialCapacity(); } - public boolean isWritable() { - return !recordsBuilder.isClosed(); - } - public byte magic() { return recordsBuilder.magic(); } From 64f09dbf7c1d2d5a71cb392af51179f51d0facd7 Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Wed, 29 Jul 2026 17:20:39 -0400 Subject: [PATCH 59/61] add tests & cleanup --- .../ChunkedRecordAccumulatorTest.java | 65 +++++++++++++++++++ .../internals/RecordAccumulatorTest.java | 4 +- 2 files changed, 67 insertions(+), 2 deletions(-) diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java index 67d67a5bc72e6..ebfbbfa1734fc 100644 --- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java +++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java @@ -635,4 +635,69 @@ public void testChunkedBatchRejectsNonChunkedStream() { assertThrows(IllegalArgumentException.class, () -> new ChunkedProducerBatch(tp1, plainBuilder, time.milliseconds())); } + + /** + * A batch closed for appends while the extension chunks were acquired off-lock must not be + * attached to: the chunks go back to the pool and the record goes to a new batch. + */ + @Test + public void testBatchClosedForAppendsDuringAllocationIsNotExtended() throws Exception { + final int chunkSize = 256; + final AtomicBoolean closeOnce = new AtomicBoolean(true); + final AtomicReference> dqRef = new AtomicReference<>(); + + BufferPool pool = new BufferPool(64L * chunkSize, chunkSize, metrics, time, "producer-metrics", + BufferPool.AllocationMode.INCREMENTAL) { + @Override + public List allocateChunks(int totalSize, long maxTimeToBlockMs) throws InterruptedException { + List chunks = super.allocateChunks(totalSize, maxTimeToBlockMs); + Deque dq = dqRef.get(); + // Mock a concurrent appender that found the batch full: RecordAccumulator.tryAppend + // calls closeForRecordAppends() whenever last.tryAppend returns null. + if (dq != null && closeOnce.getAndSet(false)) { + synchronized (dq) { + ProducerBatch last = dq.peekLast(); + if (last != null) + last.closeForRecordAppends(); + } + } + return chunks; + } + }; + ChunkedRecordAccumulator accum = new ChunkedRecordAccumulator(logContext, 8192, Compression.NONE, + /* lingerMs */ 0, /* retryBackoffMs */ 0L, /* retryBackoffMaxMs */ 0L, + /* deliveryTimeoutMs */ 3200, metrics, "producer-metrics", time, + /* transactionManager */ null, pool); + try { + // First record opens the batch with a single chunk. + accum.append(topic, partition1, 0L, key, new byte[32], Record.EMPTY_HEADERS, null, + maxBlockTimeMs, time.milliseconds(), cluster); + Deque dq = batchesFor(accum, tp1); + assertEquals(1, dq.size()); + dqRef.set(dq); + + long availableBeforeSecond = pool.availableMemory(); + + // Second record needs an extension; the batch is closed for appends mid-window. + RecordAccumulator.RecordAppendResult result = accum.append(topic, partition1, 0L, key, + new byte[300], Record.EMPTY_HEADERS, null, maxBlockTimeMs, time.milliseconds(), cluster); + + assertTrue(result.newBatchCreated, "record must land in a new batch, not the closed one"); + assertEquals(2, dq.size(), "closed batch + new batch expected"); + assertNotNull(dq.peekFirst()); + assertTrue(dq.peekFirst().isFull(), "the original batch must still be closed for appends"); + assertEquals(1, dq.peekFirst().recordCount, "no record may have been added to the closed batch"); + assertNotNull(dq.peekLast()); + assertEquals(1, dq.peekLast().recordCount); + + // The extension chunks acquired for the closed batch must have been refunded, so the only + // memory still held beyond the first batch is what the new batch needed. + assertTrue(pool.availableMemory() < availableBeforeSecond, + "the new batch should hold pool memory"); + assertEquals(0, pool.availableMemory() % chunkSize, + "pool accounting must stay chunk-aligned after the refund"); + } finally { + accum.close(); + } + } } diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/RecordAccumulatorTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/RecordAccumulatorTest.java index 7890a5474b515..386fa75636ad7 100644 --- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/RecordAccumulatorTest.java +++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/RecordAccumulatorTest.java @@ -240,7 +240,7 @@ public void testFull() throws Exception { assertEquals(1, partitionBatches.size()); ProducerBatch batch = partitionBatches.peekFirst(); - assertTrue(batch.isWritable()); + assertFalse(batch.isFull(), "batch should still accept appends"); assertEquals(0, accum.ready(metadataCache, now).readyNodes.size(), "No partitions should be ready."); } @@ -250,7 +250,7 @@ public void testFull() throws Exception { Deque partitionBatches = accum.getDeque(tp1); assertEquals(2, partitionBatches.size()); Iterator partitionBatchesIterator = partitionBatches.iterator(); - assertTrue(partitionBatchesIterator.next().isWritable()); + assertTrue(partitionBatchesIterator.next().isFull(), "the first batch should no longer accept appends"); assertEquals(Collections.singleton(node1), accum.ready(metadataCache, time.milliseconds()).readyNodes, "Our partition's leader should be ready"); List batches = accum.drain(metadataCache, Collections.singleton(node1), Integer.MAX_VALUE, 0).get(node1.id()); From 61377cfcf2c12a68d67c797944b045864af333ea Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Thu, 30 Jul 2026 07:44:26 -0400 Subject: [PATCH 60/61] addressing comments --- .../internals/ChunkedRecordAccumulator.java | 11 ++++--- .../ChunkedRecordAccumulatorTest.java | 29 ++++++++++++++----- 2 files changed, 26 insertions(+), 14 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java index 8d9d2bb8945d7..6c92447e5988d 100644 --- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java +++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java @@ -209,12 +209,11 @@ public RecordAppendResult append(String topic, ProducerBatch last = dq.peekLast(); // The batch may have changed while allocateChunks was off-lock: drained and // replaced, closed for appends, filled to its limit, or already grown by a - // concurrent appender. The instanceof excludes a replacement that cannot take - // chunks at all (a split batch is a plain ProducerBatch), and a positive - // extensionBytesNeeded means the batch still needs chunks for this record, so - // together they attach only when doing so is both safe and useful: attaching - // to a batch closed for appends would throw, and to one that no longer needs - // them would merely hold chunks it cannot use. + // concurrent appender. extensionBytesNeeded is 0 whenever attaching would be + // wrong, so it serves as both tests at once: the batch still needs chunks for + // this record, and it is still open (attaching to a stream closed for appends + // would throw). The instanceof covers the one case it cannot: a replacement + // that is a plain ProducerBatch (a split batch), which takes no chunks at all. if (last instanceof ChunkedProducerBatch && ((ChunkedProducerBatch) last).extensionBytesNeeded(timestamp, key, value, headers) > 0) { ((ChunkedProducerBatch) last).addBuffers(extensionChunks); diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java index ebfbbfa1734fc..46d4d007467c7 100644 --- a/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java +++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulatorTest.java @@ -41,9 +41,11 @@ import java.util.ArrayList; import java.util.Collections; import java.util.Deque; +import java.util.IdentityHashMap; import java.util.List; import java.util.Map; import java.util.Optional; +import java.util.Set; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; @@ -51,6 +53,7 @@ import java.util.stream.Stream; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -645,12 +648,19 @@ public void testBatchClosedForAppendsDuringAllocationIsNotExtended() throws Exce final int chunkSize = 256; final AtomicBoolean closeOnce = new AtomicBoolean(true); final AtomicReference> dqRef = new AtomicReference<>(); + // Chunks handed to the extension path, and every chunk handed back to the pool. Tracked by + // identity so the assertions can name the exact buffers rather than infer from a total. + final List extensionChunks = new ArrayList<>(); + final Set returnedToPool = Collections.newSetFromMap(new IdentityHashMap<>()); BufferPool pool = new BufferPool(64L * chunkSize, chunkSize, metrics, time, "producer-metrics", BufferPool.AllocationMode.INCREMENTAL) { @Override public List allocateChunks(int totalSize, long maxTimeToBlockMs) throws InterruptedException { List chunks = super.allocateChunks(totalSize, maxTimeToBlockMs); + // The extension acquire is the only non-blocking one (see allocateExtensionChunks). + if (maxTimeToBlockMs == 0L) + extensionChunks.addAll(chunks); Deque dq = dqRef.get(); // Mock a concurrent appender that found the batch full: RecordAccumulator.tryAppend // calls closeForRecordAppends() whenever last.tryAppend returns null. @@ -663,6 +673,12 @@ public List allocateChunks(int totalSize, long maxTimeToBlockMs) thr } return chunks; } + + @Override + public void deallocate(ByteBuffer buffer, int size) { + returnedToPool.add(buffer); + super.deallocate(buffer, size); + } }; ChunkedRecordAccumulator accum = new ChunkedRecordAccumulator(logContext, 8192, Compression.NONE, /* lingerMs */ 0, /* retryBackoffMs */ 0L, /* retryBackoffMaxMs */ 0L, @@ -676,8 +692,6 @@ public List allocateChunks(int totalSize, long maxTimeToBlockMs) thr assertEquals(1, dq.size()); dqRef.set(dq); - long availableBeforeSecond = pool.availableMemory(); - // Second record needs an extension; the batch is closed for appends mid-window. RecordAccumulator.RecordAppendResult result = accum.append(topic, partition1, 0L, key, new byte[300], Record.EMPTY_HEADERS, null, maxBlockTimeMs, time.milliseconds(), cluster); @@ -690,12 +704,11 @@ public List allocateChunks(int totalSize, long maxTimeToBlockMs) thr assertNotNull(dq.peekLast()); assertEquals(1, dq.peekLast().recordCount); - // The extension chunks acquired for the closed batch must have been refunded, so the only - // memory still held beyond the first batch is what the new batch needed. - assertTrue(pool.availableMemory() < availableBeforeSecond, - "the new batch should hold pool memory"); - assertEquals(0, pool.availableMemory() % chunkSize, - "pool accounting must stay chunk-aligned after the refund"); + // Every chunk the extension path acquired for the now-closed batch must have gone back to + // the pool, rather than being attached to it or dropped. + assertFalse(extensionChunks.isEmpty(), "the extension path should have acquired chunks"); + assertTrue(returnedToPool.containsAll(extensionChunks), + "every chunk acquired to extend the closed batch must be returned to the pool"); } finally { accum.close(); } From 5f710b8637db7a1a4c22fdf8009089830b4d5893 Mon Sep 17 00:00:00 2001 From: Lianet Magrans Date: Thu, 30 Jul 2026 08:36:50 -0400 Subject: [PATCH 61/61] reuse existing batchHeaderSize --- .../record/internal/MemoryRecordsBuilder.java | 21 +++++++------------ 1 file changed, 8 insertions(+), 13 deletions(-) diff --git a/clients/src/main/java/org/apache/kafka/common/record/internal/MemoryRecordsBuilder.java b/clients/src/main/java/org/apache/kafka/common/record/internal/MemoryRecordsBuilder.java index fc1fd73f4f91b..1b3ca24c43ae2 100644 --- a/clients/src/main/java/org/apache/kafka/common/record/internal/MemoryRecordsBuilder.java +++ b/clients/src/main/java/org/apache/kafka/common/record/internal/MemoryRecordsBuilder.java @@ -823,22 +823,18 @@ private void ensureOpenForRecordBatchWrite() { * @return The estimated number of bytes written */ private int estimatedBytesWritten() { - return estimatedBytesWritten(magic, compression.type(), estimatedCompressionRatio, uncompressedRecordsSizeInBytes); + return estimatedBytesWritten(uncompressedRecordsSizeInBytes); } /** - * Returns the projected number of bytes the builder would write for - * {@code uncompressedRecordsSizeInBytes} of records under the given magic, compression type, - * and ratio: exact for uncompressed, a ratio-aware estimate for compressed. + * Returns the projected number of bytes the builder would write for the given uncompressed + * record bytes: exact for uncompressed, a ratio-aware estimate for compressed. */ - private static int estimatedBytesWritten(byte magic, CompressionType compressionType, - float compressionRatio, - int uncompressedRecordsSizeInBytes) { - int batchHeaderSizeInBytes = AbstractRecords.recordBatchHeaderSizeInBytes(magic, compressionType); - if (compressionType == CompressionType.NONE) { - return batchHeaderSizeInBytes + uncompressedRecordsSizeInBytes; + private int estimatedBytesWritten(int uncompressedSize) { + if (compression.type() == CompressionType.NONE) { + return batchHeaderSizeInBytes + uncompressedSize; } else { - return batchHeaderSizeInBytes + (int) (uncompressedRecordsSizeInBytes * compressionRatio * COMPRESSION_RATE_ESTIMATION_FACTOR); + return batchHeaderSizeInBytes + (int) (uncompressedSize * estimatedCompressionRatio * COMPRESSION_RATE_ESTIMATION_FACTOR); } } @@ -856,8 +852,7 @@ public int estimatedBytesWrittenAfter(byte[] key, byte[] value, Header[] headers } else { recordSize = DefaultRecord.recordSizeUpperBound(keyBuffer, valueBuffer, headers); } - return estimatedBytesWritten(magic, compression.type(), estimatedCompressionRatio, - uncompressedRecordsSizeInBytes + recordSize); + return estimatedBytesWritten(uncompressedRecordsSizeInBytes + recordSize); } /**