junrao commented on code in PR #22654:
URL: https://github.com/apache/kafka/pull/22654#discussion_r3494619726
##########
clients/src/main/java/org/apache/kafka/clients/producer/ProducerConfig.java:
##########
@@ -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,
Review Comment:
Should we mark this as internal to prevent it from leaking into 4.4 before
the feature is fully implemented?
##########
clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java:
##########
@@ -0,0 +1,310 @@
+/*
+ * 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.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;
+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 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}.
+ * <p>
+ * See {@link #append} and {@link #tryAppend} for how batches are created and
grown.
+ * <p>
+ * 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,
+ 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);
+ // 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;
+ }
+
+ 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<ByteBuffer> extensionChunks = null;
+ int extensionBytes;
+ 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<ProducerBatch> dq =
topicInfo.batches.computeIfAbsent(effectivePartition, k -> new ArrayDeque<>());
+ 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) {
+ 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.
+ try {
+ extensionChunks =
chunkedFree.allocateChunks(extensionBytes, 0L);
Review Comment:
Why do we need to do a special non-blocking call? In the existing logic, if
an allocation request is blocked, all existing ProducerBatches become
immediately drainable.
##########
clients/src/main/java/org/apache/kafka/clients/producer/internals/RecordAccumulator.java:
##########
@@ -375,29 +379,35 @@ 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 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
*/
- private RecordAppendResult appendNewBatch(String topic,
+ protected RecordAppendResult appendNewBatch(String topic,
int partition,
Deque<ProducerBatch> dq,
long timestamp,
byte[] key,
byte[] value,
Header[] headers,
AppendCallbacks callbacks,
- ByteBuffer buffer,
+ Supplier<MemoryRecordsBuilder>
recordsBuilderSupplier,
long nowMs) {
assert partition != RecordMetadata.UNKNOWN_PARTITION;
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 open batch
+ // (needsBufferExtension), so the caller releases its
pre-allocated buffer and retries via
+ // the extension path.
return appendResult;
}
- MemoryRecordsBuilder recordsBuilder = recordsBuilder(buffer);
- ProducerBatch batch = new ProducerBatch(new TopicPartition(topic,
partition), recordsBuilder, nowMs);
+ MemoryRecordsBuilder recordsBuilder = recordsBuilderSupplier.get();
+ ProducerBatch batch = createProducerBatch(new TopicPartition(topic,
partition), recordsBuilder, nowMs);
Review Comment:
In this next line, should we also assert that the return value is not
appendResult.needsBufferExtension?
##########
clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java:
##########
@@ -0,0 +1,310 @@
+/*
+ * 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.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;
+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 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}.
+ * <p>
+ * See {@link #append} and {@link #tryAppend} for how batches are created and
grown.
+ * <p>
+ * 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,
+ 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);
+ // 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;
+ }
+
+ 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<ByteBuffer> extensionChunks = null;
+ int extensionBytes;
+ 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<ProducerBatch> dq =
topicInfo.batches.computeIfAbsent(effectivePartition, k -> new ArrayDeque<>());
+ 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) {
+ 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.
+ 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();
+ }
+ }
+ 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.
+ 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<ByteBuffer> 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)) {
+ // The partition switched while we allocated extension
chunks off-lock. They
+ // 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;
+ }
+ continue;
+ }
+
+ 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()) {
+ ((ChunkedProducerBatch)
last).addBuffers(extensionChunks);
+ extensionChunks = null;
+ RecordAppendResult retryResult =
tryAppend(timestamp, key, value, headers, callbacks, dq, nowMs);
+ if (retryResult != null &&
!retryResult.needsBufferExtension) {
+ 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.
+ 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);
+ extensionChunks = null;
+ continue;
+ }
+
+ // 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 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.
+ bufferStream.deallocate();
+ bufferStream = null;
+ continue;
+ }
+ 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();
+ }
+ }
+
+ /**
+ * Try to append to a ProducerBatch, with mid-batch chunk extension
support.
+ * <p>
+ * 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.
+ */
+ @Override
+ protected RecordAppendResult tryAppend(long timestamp, byte[] key, byte[]
value, Header[] headers,
Review Comment:
It's a bit awkward to have a return value of null and
RecordAppendResult.needsExtension. Could we introduce a non-null value to
indicate the batch is full?
##########
clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java:
##########
@@ -0,0 +1,262 @@
+/*
+ * 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. Chunks are supplied by the caller (initial chunks via
the constructor,
+ * additional chunks via {@link #addBuffers(List)}).
+ * <p>
+ * Current/temporary behavior:
+ * <ul>
+ * <li>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).</li>
+ * <li>{@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).</li>
+ * </ul>
+ */
+public class ChunkedByteBufferOutputStream extends ByteBufferOutputStream {
+
+ private final List<ByteBuffer> 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 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
+ */
+ public ChunkedByteBufferOutputStream(List<ByteBuffer> initialChunks, int
chunkSize, BufferPool pool) {
+ super(validatedFirstChunk(initialChunks, chunkSize));
+ this.chunkSize = chunkSize;
+ this.pool = pool;
+ this.chunks = new ArrayList<>(initialChunks);
+ this.currentChunk = this.chunks.get(0);
+ this.currentChunkIndex = 0;
+ 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<ByteBuffer>
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);
+ 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.
+ */
+ private void advanceToNextChunk() {
+ if (currentChunkIndex + 1 >= chunks.size()) {
+ // 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);
+ }
+
+ /**
+ * 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<ByteBuffer> newChunks) {
+ chunks.addAll(newChunks);
+ }
+
+ @Override
+ 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();
+ }
+ flattenedBuffer = ByteBuffer.allocate(totalSize);
+ for (ByteBuffer chunk : chunks) {
Review Comment:
We only need to iterate up to currentChunk. Ditto in position().
##########
clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java:
##########
@@ -0,0 +1,197 @@
+/*
+ * 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()}).
+ * <p>
+ * 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 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
+ * @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<ByteBuffer> 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;
+
+ // Chunks taken from the free list. The remaining bytes are reserved
against
+ // nonPooledAvailableMemory and materialized as raw allocations after
the lock is released.
+ List<ByteBuffer> 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 <= memoryRequired <= totalMemory
(validated above), so the int cast is safe.
+ 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.
+ 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;
+ }
+ }
+ // Clear the rollback tracker.
+ 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;
Review Comment:
Could we return accumulated and pooled chunks in the same place? For
example, we can set a flag like allocationCompleted to replace `accumulated =
0`. Then we can free both accumulated and pooled chunks if the flag is not set.
##########
clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java:
##########
@@ -0,0 +1,310 @@
+/*
+ * 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.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;
+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 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}.
+ * <p>
+ * See {@link #append} and {@link #tryAppend} for how batches are created and
grown.
+ * <p>
+ * 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,
+ 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);
+ // 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;
+ }
+
+ 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<ByteBuffer> extensionChunks = null;
+ int extensionBytes;
+ 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<ProducerBatch> dq =
topicInfo.batches.computeIfAbsent(effectivePartition, k -> new ArrayDeque<>());
+ 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) {
+ 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.
+ 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();
+ }
+ }
+ continue;
Review Comment:
It may take a bit of time for the closed batches to be drained. If we
continue here, it seems that the client will just busy-loop until some batches
are drained and some free space becomes available in buffer pool?
##########
clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java:
##########
@@ -0,0 +1,310 @@
+/*
+ * 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.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;
+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 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}.
+ * <p>
+ * See {@link #append} and {@link #tryAppend} for how batches are created and
grown.
+ * <p>
+ * 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,
+ 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);
+ // 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;
+ }
+
+ 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<ByteBuffer> extensionChunks = null;
+ int extensionBytes;
+ 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<ProducerBatch> dq =
topicInfo.batches.computeIfAbsent(effectivePartition, k -> new ArrayDeque<>());
+ 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) {
+ 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.
+ 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();
+ }
+ }
+ 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.
+ 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<ByteBuffer> 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)) {
+ // The partition switched while we allocated extension
chunks off-lock. They
+ // 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;
+ }
+ continue;
+ }
+
+ 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()) {
+ ((ChunkedProducerBatch)
last).addBuffers(extensionChunks);
Review Comment:
I guess two concurrent clients could add buffers exceeding the batch size?
Those buffers won't be used, but can only be freed after the batch is drained.
##########
clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java:
##########
@@ -0,0 +1,262 @@
+/*
+ * 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. Chunks are supplied by the caller (initial chunks via
the constructor,
+ * additional chunks via {@link #addBuffers(List)}).
+ * <p>
+ * Current/temporary behavior:
+ * <ul>
+ * <li>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).</li>
+ * <li>{@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).</li>
+ * </ul>
+ */
+public class ChunkedByteBufferOutputStream extends ByteBufferOutputStream {
+
+ private final List<ByteBuffer> 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 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
+ */
+ public ChunkedByteBufferOutputStream(List<ByteBuffer> initialChunks, int
chunkSize, BufferPool pool) {
+ super(validatedFirstChunk(initialChunks, chunkSize));
+ this.chunkSize = chunkSize;
+ this.pool = pool;
+ this.chunks = new ArrayList<>(initialChunks);
+ this.currentChunk = this.chunks.get(0);
+ this.currentChunkIndex = 0;
+ 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<ByteBuffer>
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);
+ 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) {
Review Comment:
Hmm, it seems this only works if `needed` is 1. If `needed` is larger than
1, it doesn't iterate the remaining chunks to ensure there is enough bytes.
##########
clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java:
##########
@@ -0,0 +1,310 @@
+/*
+ * 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.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;
+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 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}.
+ * <p>
+ * See {@link #append} and {@link #tryAppend} for how batches are created and
grown.
+ * <p>
+ * 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,
+ 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);
+ // 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;
+ }
+
+ 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<ByteBuffer> extensionChunks = null;
+ int extensionBytes;
+ 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<ProducerBatch> dq =
topicInfo.batches.computeIfAbsent(effectivePartition, k -> new ArrayDeque<>());
+ 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) {
+ 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.
+ 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();
+ }
+ }
+ 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.
+ 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<ByteBuffer> 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)) {
+ // The partition switched while we allocated extension
chunks off-lock. They
+ // 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;
+ }
+ continue;
+ }
+
+ 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()) {
+ ((ChunkedProducerBatch)
last).addBuffers(extensionChunks);
+ extensionChunks = null;
+ RecordAppendResult retryResult =
tryAppend(timestamp, key, value, headers, callbacks, dq, nowMs);
+ if (retryResult != null &&
!retryResult.needsBufferExtension) {
+ boolean enableSwitch = allBatchesFull(dq);
+
topicInfo.builtInPartitioner.updatePartitionInfo(partitionInfo,
retryResult.appendedBytes, cluster, enableSwitch);
+ return retryResult;
+ }
+ // needsBufferExtension: a concurrent appender
consumed capacity our
Review Comment:
What happens to extensionChunks? They have been added to the batch, but
won't be used. So, are they only freed when the batch is drained for sending?
##########
clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java:
##########
@@ -0,0 +1,310 @@
+/*
+ * 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.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;
+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 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}.
+ * <p>
+ * See {@link #append} and {@link #tryAppend} for how batches are created and
grown.
+ * <p>
+ * 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,
+ 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);
+ // 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;
+ }
+
+ 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<ByteBuffer> extensionChunks = null;
+ int extensionBytes;
+ 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<ProducerBatch> dq =
topicInfo.batches.computeIfAbsent(effectivePartition, k -> new ArrayDeque<>());
+ 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) {
+ 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.
+ 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();
+ }
+ }
+ 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.
+ 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<ByteBuffer> 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)) {
+ // The partition switched while we allocated extension
chunks off-lock. They
+ // 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;
+ }
+ continue;
+ }
+
+ 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()) {
+ ((ChunkedProducerBatch)
last).addBuffers(extensionChunks);
+ extensionChunks = null;
+ RecordAppendResult retryResult =
tryAppend(timestamp, key, value, headers, callbacks, dq, nowMs);
+ if (retryResult != null &&
!retryResult.needsBufferExtension) {
+ 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.
+ 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);
+ extensionChunks = null;
+ continue;
+ }
+
+ // 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) {
Review Comment:
Could appendResult.needsBufferExtension be true? We only set
appendResult.needsBufferExtension to true in
ChunkedRecordAccumulator.tryAppend(), which is not called by appendNewBatch().
##########
clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java:
##########
@@ -0,0 +1,197 @@
+/*
+ * 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()}).
+ * <p>
+ * 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 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
+ * @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<ByteBuffer> 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;
+
+ // Chunks taken from the free list. The remaining bytes are reserved
against
+ // nonPooledAvailableMemory and materialized as raw allocations after
the lock is released.
+ List<ByteBuffer> 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 <= memoryRequired <= totalMemory
(validated above), so the int cast is safe.
+ 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.
+ 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);
Review Comment:
This may be ok, but it's a bit weird to free up the chunks only to be
reallocated again. Here is an alternative that doesn't require a freeup() call.
```
// Reuse pooled chunks first. If a reused chunk covers a slot we 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 here
accumulated -= chunkSize;
this.nonPooledAvailableMemory += chunkSize; // refund →
available to other waiters
}
}
// Reserve raw memory for any still-uncovered chunks, in whole chunks.
while (pooled.size() + (int)(accumulated / chunkSize) < numChunks
&& this.nonPooledAvailableMemory >= chunkSize) {
this.nonPooledAvailableMemory -= chunkSize;
accumulated += chunkSize;
}
```
##########
clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java:
##########
@@ -0,0 +1,197 @@
+/*
+ * 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()}).
+ * <p>
+ * 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 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
+ * @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<ByteBuffer> 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;
+
+ // Chunks taken from the free list. The remaining bytes are reserved
against
+ // nonPooledAvailableMemory and materialized as raw allocations after
the lock is released.
+ List<ByteBuffer> 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 <= memoryRequired <= totalMemory
(validated above), so the int cast is safe.
+ 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.
+ while (pooled.size() < numChunks
+ && (long) (pooled.size() + 1) * chunkSize +
accumulated <= memoryRequired
Review Comment:
The first condition is redundant, given the second one.
```
(pooled+1)*chunkSize + accumulated ≤ memoryRequired
⇒ (pooled+1)*chunkSize ≤ memoryRequired − accumulated ≤ memoryRequired
= numChunks*chunkSize
⇒ pooled+1 ≤ numChunks
⇒ pooled < numChunks
```
##########
clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java:
##########
@@ -0,0 +1,197 @@
+/*
+ * 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()}).
+ * <p>
+ * 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 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
+ * @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<ByteBuffer> 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;
+
+ // Chunks taken from the free list. The remaining bytes are reserved
against
+ // nonPooledAvailableMemory and materialized as raw allocations after
the lock is released.
+ List<ByteBuffer> 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 <= memoryRequired <= totalMemory
(validated above), so the int cast is safe.
+ freeUp((int) remainingBytes);
Review Comment:
This is a no-op since all pooled chunks have been used if we reach here.
##########
clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java:
##########
@@ -0,0 +1,197 @@
+/*
+ * 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()}).
+ * <p>
+ * 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 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
+ * @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<ByteBuffer> 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;
+
+ // Chunks taken from the free list. The remaining bytes are reserved
against
+ // nonPooledAvailableMemory and materialized as raw allocations after
the lock is released.
+ List<ByteBuffer> 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 <= memoryRequired <= totalMemory
(validated above), so the int cast is safe.
Review Comment:
Why is remainingBytes guaranteed to be an int? memoryRequired could be
larger than int and pooled.size() initially could be 0.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
