junrao commented on code in PR #22654:
URL: https://github.com/apache/kafka/pull/22654#discussion_r3606283544
##########
clients/src/main/java/org/apache/kafka/clients/producer/internals/RecordAccumulator.java:
##########
@@ -1251,23 +1270,83 @@ public PartitionerConfig() {
}
/*
- * Metadata about a record just appended to the record accumulator
+ * Result of an attempt to append a record to the accumulator. Carries
exactly one of three
+ * mutually-exclusive outcomes: the record was appended ({@link
RecordAppendResult#appended()}, {@code future} is set),
+ * the open batch needs more chunk capacity first ({@link
RecordAppendResult#needsBufferExtension()}),
+ * or a new batch must be created for the record ({@link
RecordAppendResult#needsNewBatch()}).
*/
public static final class RecordAppendResult {
+ /**
+ * The three mutually-exclusive outcomes of an append attempt.
+ */
+ public enum Outcome { APPENDED, NEEDS_BUFFER_EXTENSION,
NEEDS_NEW_BATCH }
+
+ public final Outcome outcome;
Review Comment:
Should this be private?
##########
clients/src/main/java/org/apache/kafka/clients/producer/internals/RecordAccumulator.java:
##########
@@ -375,46 +379,60 @@ 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
+ * @return the append result, which never has {@code needsNewBatch=true}:
the method either
+ * propagates a non-{@code needsNewBatch} result from an open
batch created concurrently
+ * (a success, or — incremental strategy — a {@code
needsBufferExtension} signal), or it
Review Comment:
This comment is not very clear to me. `a success` should return `appended`,
not `needsBufferExtension`, right?
##########
clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java:
##########
@@ -0,0 +1,328 @@
+/*
+ * 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.internals.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;
+ // Set once the content has been read for send via buffer().
+ private boolean finalized;
+ // Single-buffer view produced by flatten() and cached here so repeat
buffer() calls
+ // return the same instance. To be removed once scatter-gather
(KAFKA-20580) is implemented.
+ private ByteBuffer flattenedBuffer;
+
+ /**
+ * 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;
+ }
+
+ /**
+ * 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) {
+ ensureNotDeallocated();
+ ensureWritable();
+ ensureChunkCapacity(1);
+ currentChunk.put((byte) b);
+ }
+
+ @Override
+ public void write(byte[] bytes, int off, int len) {
+ ensureNotDeallocated();
+ ensureWritable();
+ while (len > 0) {
+ ensureChunkCapacity(1);
+ int toWrite = Math.min(len, currentChunk.remaining());
+ currentChunk.put(bytes, off, toWrite);
+ off += toWrite;
+ len -= toWrite;
+ }
+ }
+
+ @Override
+ public void write(ByteBuffer sourceBuffer) {
+ ensureNotDeallocated();
+ ensureWritable();
+ 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);
+ }
+ }
+
+ /**
+ * Guards against writes after the stream has been finalized with a call
to {@link #buffer()}.
+ */
+ private void ensureWritable() {
+ if (finalized)
+ throw new IllegalStateException("cannot write after buffer() has
been called");
+ }
+
+ /**
+ * Guards against any use after {@link #deallocate()} has returned the
chunks.
+ */
+ private void ensureNotDeallocated() {
+ if (currentChunk == null)
+ throw new IllegalStateException("operation not allowed after the
stream has been deallocated");
+ }
+
+ private void ensureChunkCapacity(int needed) {
+ while (currentChunk.remaining() < needed) {
+ advanceToNextChunk();
+ }
+ }
+
+ /**
+ * 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) {
+ ensureNotDeallocated();
+ chunks.addAll(newChunks);
+ }
+
+ /**
+ * Returns the written bytes as a single contiguous buffer. Calling this
finalizes the stream: no
+ * further writes are allowed (see {@link #ensureWritable()}) and later
calls return the same
+ * instance, which callers such as {@code
MemoryRecordsBuilder#writeDefaultBatchHeader} rely on
+ * when they write the batch header directly into the returned buffer.
+ */
+ @Override
+ public ByteBuffer buffer() {
+ ensureNotDeallocated();
+ finalized = true;
+ if (flattenedBuffer == null)
+ flattenedBuffer = flatten();
+ return flattenedBuffer;
+ }
+
+ /**
+ * Flattens the written bytes across the data-bearing chunks into a single
new buffer (an extra
+ * copy). This will be removed once scatter-gather send (KAFKA-20580) is
implemented.
+ */
+ private ByteBuffer flatten() {
+ // Written bytes only live in chunks up to currentChunk, later chunks
are untouched.
+ int totalSize = 0;
+ for (int i = 0; i <= currentChunkIndex; i++) {
+ totalSize += chunks.get(i).position();
+ }
+ ByteBuffer flattened = ByteBuffer.allocate(totalSize);
+ for (int i = 0; i <= currentChunkIndex; i++) {
+ ByteBuffer chunk = chunks.get(i);
+ int chunkPos = chunk.position();
+ chunk.flip();
+ flattened.put(chunk);
+ chunk.limit(chunk.capacity());
+ chunk.position(chunkPos);
+ }
+ return flattened;
+ }
+
+ /**
+ * Releases the fully-unused chunks, given that the stream is closed for
appends.
+ */
+ @Override
+ public void close() {
+ releaseUnusedChunks();
+ }
+
+ /**
+ * Return the fully-unused chunks to the pool. The data-bearing chunks are
+ * kept until batch completion ({@link #deallocate()}), as they hold the
in-flight data.
+ */
+ private void releaseUnusedChunks() {
+ if (currentChunk == null) // already deallocated; nothing attached
+ return;
+ List<ByteBuffer> unused = chunks.subList(currentChunkIndex + 1,
chunks.size());
+ if (pool != null) {
+ for (ByteBuffer chunk : unused)
+ pool.deallocate(chunk);
+ }
+ // Remove the released chunks from `chunks`, so they are
+ // not deallocated again on batch completion.
+ unused.clear();
+ }
+
+ /**
+ * Total bytes written across all chunks.
+ */
+ @Override
+ public int position() {
+ ensureNotDeallocated();
+ // Written bytes only live in chunks up to currentChunk, later chunks
are untouched.
+ int total = 0;
+ for (int i = 0; i <= currentChunkIndex; i++) {
+ total += chunks.get(i).position();
+ }
+ return total;
+ }
+
+ /**
+ * Sets the write position, walking across pre-supplied chunks if the
requested position
+ * exceeds the first chunk's capacity. Only valid before any write.
+ */
+ @Override
+ public void position(int position) {
+ ensureNotDeallocated();
+ if (currentChunkIndex != 0 || currentChunk.position() != 0) {
+ throw new IllegalStateException("position() can only be called
before any writes");
+ }
+ int remaining = position;
+ int idx = 0;
+ while (remaining > 0 && idx < chunks.size()) {
+ ByteBuffer chunk = chunks.get(idx);
+ int take = Math.min(remaining, chunk.capacity());
+ chunk.position(take);
+ remaining -= take;
+ if (remaining > 0)
+ idx++;
+ }
+ if (remaining > 0) {
+ throw new IllegalArgumentException("position " + position
+ + " exceeds total pre-allocated capacity");
+ }
+ currentChunkIndex = idx;
+ currentChunk = chunks.get(idx);
+ }
+
+ /**
+ * Total capacity across all attached chunks (written + free).
+ * Every chunk has the same size, so this equals {@code position() +
remaining()} without walking the list.
+ */
+ int attachedCapacity() {
+ ensureNotDeallocated();
+ return chunks.size() * chunkSize;
+ }
+
+ /**
+ * Total bytes available across the current chunk and every queued
(not-yet-active) chunk.
+ */
+ @Override
+ public int remaining() {
+ ensureNotDeallocated();
+ int total = currentChunk.remaining();
+ for (int i = currentChunkIndex + 1; i < chunks.size(); i++)
+ total += chunks.get(i).remaining();
+ return total;
+ }
+
+ @Override
+ public int limit() {
+ return Integer.MAX_VALUE;
+ }
+
+ @Override
+ public int initialCapacity() {
+ return chunkSize;
+ }
+
+ @Override
+ public void ensureRemaining(int remainingBytesRequired) {
+ ensureNotDeallocated();
+ // A single call can guarantee at most `chunkSize` of space (the
stream advances one chunk
+ // at a time). Callers needing more attach chunks via addBuffers
first. write(byte[]) loops
+ // across chunks, so contiguous capacity isn't required.
+ ensureChunkCapacity(Math.min(remainingBytesRequired, chunkSize));
+ }
+
+ /**
+ * Returns all pool-allocated chunks to the buffer pool. Called at batch
completion.
+ */
+ public void deallocate(BufferPool pool) {
+ if (pool != null) {
+ for (ByteBuffer chunk : chunks) {
+ pool.deallocate(chunk);
+ }
+ }
+ chunks.clear();
+ currentChunk = null;
+ currentChunkIndex = 0;
Review Comment:
0 is a valid index. Should we set it to -1?
##########
clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java:
##########
@@ -0,0 +1,328 @@
+/*
+ * 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.internals.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;
+ // Set once the content has been read for send via buffer().
+ private boolean finalized;
+ // Single-buffer view produced by flatten() and cached here so repeat
buffer() calls
+ // return the same instance. To be removed once scatter-gather
(KAFKA-20580) is implemented.
+ private ByteBuffer flattenedBuffer;
+
+ /**
+ * 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;
+ }
+
+ /**
+ * 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) {
+ ensureNotDeallocated();
+ ensureWritable();
+ ensureChunkCapacity(1);
+ currentChunk.put((byte) b);
+ }
+
+ @Override
+ public void write(byte[] bytes, int off, int len) {
+ ensureNotDeallocated();
+ ensureWritable();
+ while (len > 0) {
+ ensureChunkCapacity(1);
+ int toWrite = Math.min(len, currentChunk.remaining());
+ currentChunk.put(bytes, off, toWrite);
+ off += toWrite;
+ len -= toWrite;
+ }
+ }
+
+ @Override
+ public void write(ByteBuffer sourceBuffer) {
+ ensureNotDeallocated();
+ ensureWritable();
+ 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);
+ }
+ }
+
+ /**
+ * Guards against writes after the stream has been finalized with a call
to {@link #buffer()}.
+ */
+ private void ensureWritable() {
+ if (finalized)
+ throw new IllegalStateException("cannot write after buffer() has
been called");
+ }
+
+ /**
+ * Guards against any use after {@link #deallocate()} has returned the
chunks.
+ */
+ private void ensureNotDeallocated() {
+ if (currentChunk == null)
+ throw new IllegalStateException("operation not allowed after the
stream has been deallocated");
+ }
+
+ private void ensureChunkCapacity(int needed) {
+ while (currentChunk.remaining() < needed) {
+ advanceToNextChunk();
+ }
+ }
+
+ /**
+ * 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) {
+ ensureNotDeallocated();
+ chunks.addAll(newChunks);
Review Comment:
Should we call ensureWritable() here too?
##########
clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java:
##########
@@ -0,0 +1,328 @@
+/*
+ * 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.internals.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;
+ // Set once the content has been read for send via buffer().
+ private boolean finalized;
+ // Single-buffer view produced by flatten() and cached here so repeat
buffer() calls
+ // return the same instance. To be removed once scatter-gather
(KAFKA-20580) is implemented.
+ private ByteBuffer flattenedBuffer;
+
+ /**
+ * 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;
+ }
+
+ /**
+ * 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) {
+ ensureNotDeallocated();
+ ensureWritable();
+ ensureChunkCapacity(1);
+ currentChunk.put((byte) b);
+ }
+
+ @Override
+ public void write(byte[] bytes, int off, int len) {
+ ensureNotDeallocated();
+ ensureWritable();
+ while (len > 0) {
+ ensureChunkCapacity(1);
+ int toWrite = Math.min(len, currentChunk.remaining());
+ currentChunk.put(bytes, off, toWrite);
+ off += toWrite;
+ len -= toWrite;
+ }
+ }
+
+ @Override
+ public void write(ByteBuffer sourceBuffer) {
+ ensureNotDeallocated();
+ ensureWritable();
+ 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);
+ }
+ }
+
+ /**
+ * Guards against writes after the stream has been finalized with a call
to {@link #buffer()}.
+ */
+ private void ensureWritable() {
+ if (finalized)
+ throw new IllegalStateException("cannot write after buffer() has
been called");
+ }
+
+ /**
+ * Guards against any use after {@link #deallocate()} has returned the
chunks.
+ */
+ private void ensureNotDeallocated() {
+ if (currentChunk == null)
+ throw new IllegalStateException("operation not allowed after the
stream has been deallocated");
+ }
+
+ private void ensureChunkCapacity(int needed) {
+ while (currentChunk.remaining() < needed) {
+ advanceToNextChunk();
+ }
+ }
+
+ /**
+ * 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) {
+ ensureNotDeallocated();
+ chunks.addAll(newChunks);
+ }
+
+ /**
+ * Returns the written bytes as a single contiguous buffer. Calling this
finalizes the stream: no
+ * further writes are allowed (see {@link #ensureWritable()}) and later
calls return the same
+ * instance, which callers such as {@code
MemoryRecordsBuilder#writeDefaultBatchHeader} rely on
+ * when they write the batch header directly into the returned buffer.
+ */
+ @Override
+ public ByteBuffer buffer() {
+ ensureNotDeallocated();
+ finalized = true;
Review Comment:
It's probably better to set this when close() is called. We can then rename
it to closed.
##########
clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java:
##########
@@ -0,0 +1,328 @@
+/*
+ * 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.internals.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;
+ // Set once the content has been read for send via buffer().
+ private boolean finalized;
+ // Single-buffer view produced by flatten() and cached here so repeat
buffer() calls
+ // return the same instance. To be removed once scatter-gather
(KAFKA-20580) is implemented.
+ private ByteBuffer flattenedBuffer;
+
+ /**
+ * 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;
+ }
+
+ /**
+ * 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) {
+ ensureNotDeallocated();
+ ensureWritable();
+ ensureChunkCapacity(1);
+ currentChunk.put((byte) b);
+ }
+
+ @Override
+ public void write(byte[] bytes, int off, int len) {
+ ensureNotDeallocated();
+ ensureWritable();
+ while (len > 0) {
+ ensureChunkCapacity(1);
+ int toWrite = Math.min(len, currentChunk.remaining());
+ currentChunk.put(bytes, off, toWrite);
+ off += toWrite;
+ len -= toWrite;
+ }
+ }
+
+ @Override
+ public void write(ByteBuffer sourceBuffer) {
+ ensureNotDeallocated();
+ ensureWritable();
+ 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);
+ }
+ }
+
+ /**
+ * Guards against writes after the stream has been finalized with a call
to {@link #buffer()}.
+ */
+ private void ensureWritable() {
+ if (finalized)
+ throw new IllegalStateException("cannot write after buffer() has
been called");
+ }
+
+ /**
+ * Guards against any use after {@link #deallocate()} has returned the
chunks.
+ */
+ private void ensureNotDeallocated() {
+ if (currentChunk == null)
+ throw new IllegalStateException("operation not allowed after the
stream has been deallocated");
+ }
+
+ private void ensureChunkCapacity(int needed) {
+ while (currentChunk.remaining() < needed) {
+ advanceToNextChunk();
+ }
+ }
+
+ /**
+ * 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) {
+ ensureNotDeallocated();
+ chunks.addAll(newChunks);
+ }
+
+ /**
+ * Returns the written bytes as a single contiguous buffer. Calling this
finalizes the stream: no
+ * further writes are allowed (see {@link #ensureWritable()}) and later
calls return the same
+ * instance, which callers such as {@code
MemoryRecordsBuilder#writeDefaultBatchHeader} rely on
+ * when they write the batch header directly into the returned buffer.
+ */
+ @Override
+ public ByteBuffer buffer() {
+ ensureNotDeallocated();
+ finalized = true;
+ if (flattenedBuffer == null)
+ flattenedBuffer = flatten();
+ return flattenedBuffer;
+ }
+
+ /**
+ * Flattens the written bytes across the data-bearing chunks into a single
new buffer (an extra
+ * copy). This will be removed once scatter-gather send (KAFKA-20580) is
implemented.
+ */
+ private ByteBuffer flatten() {
+ // Written bytes only live in chunks up to currentChunk, later chunks
are untouched.
+ int totalSize = 0;
+ for (int i = 0; i <= currentChunkIndex; i++) {
+ totalSize += chunks.get(i).position();
+ }
+ ByteBuffer flattened = ByteBuffer.allocate(totalSize);
+ for (int i = 0; i <= currentChunkIndex; i++) {
+ ByteBuffer chunk = chunks.get(i);
+ int chunkPos = chunk.position();
+ chunk.flip();
+ flattened.put(chunk);
+ chunk.limit(chunk.capacity());
+ chunk.position(chunkPos);
+ }
+ return flattened;
+ }
+
+ /**
+ * Releases the fully-unused chunks, given that the stream is closed for
appends.
+ */
+ @Override
+ public void close() {
+ releaseUnusedChunks();
+ }
+
+ /**
+ * Return the fully-unused chunks to the pool. The data-bearing chunks are
+ * kept until batch completion ({@link #deallocate()}), as they hold the
in-flight data.
+ */
+ private void releaseUnusedChunks() {
+ if (currentChunk == null) // already deallocated; nothing attached
+ return;
+ List<ByteBuffer> unused = chunks.subList(currentChunkIndex + 1,
chunks.size());
+ if (pool != null) {
+ for (ByteBuffer chunk : unused)
+ pool.deallocate(chunk);
+ }
+ // Remove the released chunks from `chunks`, so they are
+ // not deallocated again on batch completion.
+ unused.clear();
+ }
+
+ /**
+ * Total bytes written across all chunks.
+ */
+ @Override
+ public int position() {
+ ensureNotDeallocated();
+ // Written bytes only live in chunks up to currentChunk, later chunks
are untouched.
+ int total = 0;
+ for (int i = 0; i <= currentChunkIndex; i++) {
+ total += chunks.get(i).position();
+ }
+ return total;
+ }
+
+ /**
+ * Sets the write position, walking across pre-supplied chunks if the
requested position
+ * exceeds the first chunk's capacity. Only valid before any write.
+ */
+ @Override
+ public void position(int position) {
+ ensureNotDeallocated();
+ if (currentChunkIndex != 0 || currentChunk.position() != 0) {
+ throw new IllegalStateException("position() can only be called
before any writes");
+ }
+ int remaining = position;
+ int idx = 0;
+ while (remaining > 0 && idx < chunks.size()) {
+ ByteBuffer chunk = chunks.get(idx);
+ int take = Math.min(remaining, chunk.capacity());
+ chunk.position(take);
+ remaining -= take;
+ if (remaining > 0)
+ idx++;
+ }
+ if (remaining > 0) {
+ throw new IllegalArgumentException("position " + position
+ + " exceeds total pre-allocated capacity");
+ }
+ currentChunkIndex = idx;
+ currentChunk = chunks.get(idx);
+ }
+
+ /**
+ * Total capacity across all attached chunks (written + free).
+ * Every chunk has the same size, so this equals {@code position() +
remaining()} without walking the list.
+ */
+ int attachedCapacity() {
+ ensureNotDeallocated();
+ return chunks.size() * chunkSize;
+ }
+
+ /**
+ * Total bytes available across the current chunk and every queued
(not-yet-active) chunk.
+ */
+ @Override
+ public int remaining() {
+ ensureNotDeallocated();
+ int total = currentChunk.remaining();
+ for (int i = currentChunkIndex + 1; i < chunks.size(); i++)
+ total += chunks.get(i).remaining();
+ return total;
+ }
+
+ @Override
+ public int limit() {
+ return Integer.MAX_VALUE;
+ }
+
+ @Override
+ public int initialCapacity() {
+ return chunkSize;
+ }
+
+ @Override
+ public void ensureRemaining(int remainingBytesRequired) {
+ ensureNotDeallocated();
Review Comment:
Should we call this in limit() and initialCapacity() too?
##########
clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java:
##########
@@ -0,0 +1,338 @@
+/*
+ * 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.
+ deallocateExtensionChunks(extensionChunks);
+ 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()) {
+ // Attach the chunks we sized off-lock without
re-checking whether a
+ // concurrent appender already grew the batch
enough. Optimizes for the
+ // uncontended case; under concurrency two
appenders can each attach their
+ // own gap (temporary over-allocation, one could
have been enough for both).
+ // The unused chunks returns to the pool when the
batch closes for appends.
+ ((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.
+ deallocateExtensionChunks(extensionChunks);
+ extensionChunks = null;
+ continue;
+ }
+
+ // needsNewBatch path: extensionChunks == null here
implies needsNewBatch,
+ // so bufferStream was allocated (this iteration or
carried from a prior one).
+ if (bufferStream == null)
+ throw new IllegalStateException("needsNewBatch path
reached without an allocated buffer stream");
+ int firstRecordSize =
AbstractRecords.estimateSizeInBytesUpperBound(
Review Comment:
This code is different from the code when the new batch is allocated. Could
we remember the first record size when the new batch is allocated and reuse it
here?
```
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);
```
##########
clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedRecordAccumulator.java:
##########
@@ -0,0 +1,338 @@
+/*
+ * 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.
+ deallocateExtensionChunks(extensionChunks);
+ 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()) {
+ // Attach the chunks we sized off-lock without
re-checking whether a
+ // concurrent appender already grew the batch
enough. Optimizes for the
+ // uncontended case; under concurrency two
appenders can each attach their
+ // own gap (temporary over-allocation, one could
have been enough for both).
+ // The unused chunks returns to the pool when the
batch closes for appends.
+ ((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.
+ deallocateExtensionChunks(extensionChunks);
+ extensionChunks = null;
+ continue;
+ }
+
+ // needsNewBatch path: extensionChunks == null here
implies needsNewBatch,
+ // so bufferStream was allocated (this iteration or
carried from a prior one).
+ if (bufferStream == null)
+ throw new IllegalStateException("needsNewBatch path
reached without an allocated buffer stream");
+ int firstRecordSize =
AbstractRecords.estimateSizeInBytesUpperBound(
+ RecordBatch.CURRENT_MAGIC_VALUE,
compression.type(), key, value, headers);
+ 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.needsNewBatch())
+ throw new IllegalStateException("appendNewBatch must
not return a needsNewBatch result");
Review Comment:
Could we do this check right after we get appendResult back?
--
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]
