m1a2st commented on code in PR #22654: URL: https://github.com/apache/kafka/pull/22654#discussion_r3566407827
########## clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java: ########## @@ -0,0 +1,287 @@ +/* + * 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) { + while (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; + } + // Written bytes only live in chunks up to currentChunk; later chunks are untouched. + int lastDataChunk = Math.min(currentChunkIndex, chunks.size() - 1); + int totalSize = 0; + for (int i = 0; i <= lastDataChunk; i++) { + totalSize += chunks.get(i).position(); + } + flattenedBuffer = ByteBuffer.allocate(totalSize); + for (int i = 0; i <= lastDataChunk; i++) { + ByteBuffer chunk = chunks.get(i); + int chunkPos = chunk.position(); + chunk.flip(); + flattenedBuffer.put(chunk); + chunk.limit(chunk.capacity()); + chunk.position(chunkPos); + } + dirty = false; + releaseUnusedChunks(); Review Comment: `buffer()` is a public override of `ByteBufferOutputStream.buffer()`, so its contract appears to be that of a simple accessor. However, this implementation hides a mutating side effect: it calls `releaseUnusedChunks()`, which returns chunks to the pool and truncates the chunk list. That makes the behavior somewhat surprising. Any future caller that invokes `buffer()` while the batch is still writable—for example, for metrics collection, a debug/peek path—would inadvertently release chunks that have not yet been written. Could we make the release operation explicit instead of embedding it as a side effect of a getter? ########## 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) { + 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.appended()) { + boolean enableSwitch = allBatchesFull(dq); + topicInfo.builtInPartitioner.updatePartitionInfo(partitionInfo, retryResult.appendedBytes, cluster, enableSwitch); + return retryResult; + } + // Still not appended: concurrent appenders filled the batch, + // so the extension we attached is no longer enough. + // Loop so the next iteration routes the record + // right: needsBufferExtension with a fresh gap, or needsNewBatch + continue; + } + // The open batch is gone, closed, or non-chunked (e.g., a split batch). Return chunks to pool. + for (ByteBuffer chunk : extensionChunks) + chunkedFree.deallocate(chunk); + extensionChunks = null; + continue; + } + + // needsNewBatch path: extensionChunks == null here implies needsNewBatch, + // so bufferStream was allocated (this iteration or carried from a prior one). + assert bufferStream != null; + int firstRecordSize = AbstractRecords.estimateSizeInBytesUpperBound( + RecordBatch.CURRENT_MAGIC_VALUE, compression.type(), key, value, headers); + final ChunkedByteBufferOutputStream batchStream = bufferStream; + 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, which appends or returns + * {@link RecordAppendResult#NEEDS_NEW_BATCH}. + */ + @Override + protected RecordAppendResult tryAppend(long timestamp, byte[] key, byte[] value, Header[] headers, + Callback callback, Deque<ProducerBatch> deque, long nowMs) { + if (closed) + throw new KafkaException("Producer closed while send in progress"); + ProducerBatch last = deque.peekLast(); + // Split batches in an incremental deque are plain ProducerBatch (heap-backed, grow-on-demand) + // and never need chunk extension, so the check only applies to chunked batches. + if (last instanceof ChunkedProducerBatch) { + int extensionBytes = ((ChunkedProducerBatch) last).extensionBytesNeeded(timestamp, key, value, headers); + if (extensionBytes > 0) + return RecordAppendResult.needsExtension(extensionBytes); + } + return super.tryAppend(timestamp, key, value, headers, callback, deque, nowMs); Review Comment: A successful mid-batch append currently computes the record's serialized size three times, compared to just once on the non-chunked path. Tracing `ChunkedRecordAccumulator.tryAppend → ChunkedProducerBatch.extensionBytesNeeded → super.tryAppend`: 1. `extensionBytesNeeded` → `recordsBuilder.hasRoomFor(...)` → `DefaultRecord.sizeInBytes(...)` — walks the key, value, and headers. 2. `extensionBytesNeeded` → `recordsBuilder.estimatedBytesWrittenAfter(...)` → `DefaultRecord.recordSizeUpperBound(...)` — walks them again. 3. `super.tryAppend` → `ProducerBatch.tryAppend` → `hasRoomFor(...)` → `DefaultRecord.sizeInBytes(...)` — walks them a third time. Steps (1) and (3) perform the same `hasRoomFor` check: the accumulator verifies capacity, and `ProducerBatch.tryAppend` immediately verifies it again. Is there a way to compute the record size once and reuse it for both the room check and the extension-gap calculation? -- 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]
