lianetm commented on code in PR #22654:
URL: https://github.com/apache/kafka/pull/22654#discussion_r3590439749
##########
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}.
Review Comment:
done
##########
clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java:
##########
@@ -0,0 +1,325 @@
+/*
+ * 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;
+ 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<>());
+ RecordAppendResult appendResult;
+ synchronized (dq) {
+ if (partitionChanged(topic, topicInfo, partitionInfo, dq,
nowMs, cluster))
+ continue;
+
+ // The tryAppend checks the open batch (dq.peekLast()) for
chunk capacity:
+ // a needsBufferExtension result means it is within its
batch-size limit
+ // but its chunks lack capacity for this record, so it
will allocate the gap
+ // outside the deque lock. A needsNewBatch result means
there is no open batch
+ // (full or absent), so it will fall through to the
first-record (new batch) path.
+ appendResult = tryAppend(timestamp, key, value, headers,
callbacks, dq, nowMs);
+ if (appendResult.appended()) {
+ boolean enableSwitch = allBatchesFull(dq);
+
topicInfo.builtInPartitioner.updatePartitionInfo(partitionInfo,
appendResult.appendedBytes, cluster, enableSwitch);
+ return appendResult;
+ }
+ }
+
+ if (appendResult.needsBufferExtension) {
+ // Mid-batch extension: the open batch can still take this
record so grow it in
+ // place. The acquire is non-blocking to fail fast when
the pool is exhausted:
+ // close the batch and let the record block once on the
new-batch path.
+ // A blocking call would lead to the same outcome, but
would block once here
+ // and still need a second blocking call to start a new
batch anyways
+ // (a first blocking call here would make all open batches
drainable, including this one,
+ // so most probably our batch would be gone/drained by the
time memory is returned to the pool,
+ // and we would need a new batch for our record anyways).
+ try {
+ extensionChunks =
chunkedFree.allocateChunks(appendResult.extensionBytesNeeded, 0L);
+ } catch (BufferExhaustedException e) {
+ log.trace("Pool exhausted while extending batch for
topic {} partition {}; closing existing batch",
+ topic, effectivePartition);
+ synchronized (dq) {
+ ProducerBatch last = dq.peekLast();
+ if (last != null && last.isWritable()) {
+ last.closeForRecordAppends();
+ }
+ }
+ // Continue to the next iteration that should block to
start a new batch (needsNewBatch),
+ // given that this one has been closed for appends.
+ continue;
+ }
+ nowMs = time.milliseconds();
+ } else if (appendResult.needsNewBatch && bufferStream == null)
{
+ // The open batch is done (e.g., full, closed) so start a
new one.
+ // Block on the pool for enough chunks to fit this record,
sized with the
+ // same cumulative estimator used mid-batch (header +
record bytes for NONE,
+ // ratio-adjusted when compressed) so the two stay
consistent.
+ int recordUncompressed =
AbstractRecords.recordSizeUpperBound(
+ RecordBatch.CURRENT_MAGIC_VALUE,
compression.type(), key, value, headers);
+ int size = MemoryRecordsBuilder.estimatedBytesWritten(
+ 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;
+ try {
+ initialChunks = chunkedFree.allocateChunks(size,
maxTimeToBlock);
+ } catch (BufferExhaustedException e) {
+ // The blocking new-batch acquire was not able to get
memory within
+ // max.block.ms. Record it in the buffer-exhausted
metrics.
+ chunkedFree.recordBufferExhausted();
+ throw e;
+ }
+ nowMs = time.milliseconds();
+ bufferStream = new
ChunkedByteBufferOutputStream(initialChunks, chunkedFree.poolableSize(),
chunkedFree);
+ }
+
+ 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) {
Review Comment:
makes sense, done
--
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]
