lianetm commented on code in PR #22654: URL: https://github.com/apache/kafka/pull/22654#discussion_r3590446044
########## clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java: ########## @@ -0,0 +1,204 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.kafka.clients.producer.internals; + +import org.apache.kafka.clients.producer.BufferExhaustedException; +import org.apache.kafka.common.KafkaException; +import org.apache.kafka.common.metrics.Metrics; +import org.apache.kafka.common.utils.Time; + +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.locks.Condition; + +/** + * A {@link BufferPool} dedicated to chunk-sized buffer reuse (chunk size = {@link #poolableSize()}). + * <p> + * Adds {@link #allocateChunks(int, long)} to acquire multiple chunks atomically. + */ +public class ChunkedBufferPool extends BufferPool { + + public ChunkedBufferPool(long memory, int chunkSize, Metrics metrics, Time time, String metricGrpName) { + super(memory, chunkSize, metrics, time, metricGrpName); + } + + /** + * Allocate {@code ceil(totalSize / chunkSize)} chunk-sized buffers atomically, mirroring + * {@link BufferPool#allocate}: satisfied immediately if memory is available, else blocks up to + * {@code maxTimeToBlockMs} for the whole request (FIFO on {@link #waiters}). + * The reservation is tracked as bytes against {@link #nonPooledAvailableMemory} plus chunks polled + * from {@link #free}. Any failure refunds the whole reservation and signals the next waiter + * before the exception propagates, so a failed request leaves nothing reserved. + * + * @param totalSize minimum total bytes of capacity required across the returned chunks + * @param maxTimeToBlockMs maximum time in milliseconds to block waiting for memory + * @return list of {@code ceil(totalSize / chunkSize)} {@code ByteBuffer}s, each of capacity + * {@code chunkSize} + * @throws InterruptedException if interrupted while waiting + * @throws IllegalArgumentException if {@code totalSize <= 0}, or if the request rounded up to + * whole chunks exceeds {@code totalMemory()} + * @throws BufferExhaustedException if the request can't be satisfied within {@code maxTimeToBlockMs} + * @throws KafkaException if the pool is closed during the wait + */ + public List<ByteBuffer> allocateChunks(int totalSize, long maxTimeToBlockMs) throws InterruptedException { + if (totalSize <= 0) + throw new IllegalArgumentException("totalSize must be positive: " + totalSize); + throwIfChunksNeededExceedsPool(totalSize); + + int chunkSize = poolableSize(); + int numChunks = (int) (((long) totalSize + chunkSize - 1L) / chunkSize); + long memoryRequired = (long) numChunks * chunkSize; + + // Chunks taken from the free list. The remaining bytes are reserved against + // nonPooledAvailableMemory and materialized as raw allocations after the lock is released. + List<ByteBuffer> pooled = new ArrayList<>(numChunks); + + lock.lock(); + if (this.closed) { + lock.unlock(); + throw new KafkaException("Producer closed while allocating memory"); + } + try { + long freeListBytes = (long) free.size() * chunkSize; + if (this.nonPooledAvailableMemory + freeListBytes >= memoryRequired) { + // Enough memory available to allocate the chunks needed + while (pooled.size() < numChunks && !free.isEmpty()) + pooled.add(free.pollFirst()); + long remainingBytes = memoryRequired - (long) pooled.size() * chunkSize; + if (remainingBytes > 0) { + // remainingBytes > 0 means the free list was fully drained into `pooled`, so the + // remainder comes entirely from non-pooled memory (sufficient per the check above). + this.nonPooledAvailableMemory -= remainingBytes; + } + } else { + // Not enough memory available to allocate the chunks needed, so we need to wait for memory. + // Same as in BufferPool.allocate, but wait to acquire the memory needed for all the chunks. + // A single Condition is added to the waiter's list to ensure FIFO fairness at the request level. + // + // `accumulated` tracks bytes drawn from nonPooledAvailableMemory only (pool chunks + // already taken live in `pooled`), and is always a whole-chunk multiple. If the wait + // does not complete (timeout / close / interrupt), the finally refunds the whole + // reservation: `accumulated` back to non-pooled memory, `pooled` back to the free chunks list. + long accumulated = 0; Review Comment: done ########## clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java: ########## @@ -0,0 +1,291 @@ +/* + * 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) { Review Comment: agreed, done. Added a "finalized" flag to gate the writes, and extracted the flatten logic better. -- 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]
