m1a2st commented on code in PR #22654: URL: https://github.com/apache/kafka/pull/22654#discussion_r3537069939
########## clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedBufferPool.java: ########## @@ -0,0 +1,200 @@ +/* + * 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; + boolean allocationCompleted = false; + Condition moreMemory = lock.newCondition(); + try { + long remainingTimeToBlockNs = TimeUnit.MILLISECONDS.toNanos(maxTimeToBlockMs); + waiters.addLast(moreMemory); + while ((long) pooled.size() * chunkSize + accumulated < memoryRequired) { + long startWaitNs = time.nanoseconds(); + long timeNs; + boolean waitingTimeElapsed; + try { + waitingTimeElapsed = !moreMemory.await(remainingTimeToBlockNs, TimeUnit.NANOSECONDS); + } finally { + long endWaitNs = time.nanoseconds(); + timeNs = Math.max(0L, endWaitNs - startWaitNs); + recordWaitTime(timeNs); + } + + if (this.closed) + throw new KafkaException("Producer closed while allocating memory"); + + if (waitingTimeElapsed) { Review Comment: `buffer-exhausted-records` is not updated on this path, so the incremental metrics never increment `buffer-exhausted-rate` or `buffer-exhausted-total` when a `send()` is dropped due to buffer pool exhaustion. ########## clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedProducerBatch.java: ########## @@ -0,0 +1,113 @@ +/* + * 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.Callback; +import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.header.Header; +import org.apache.kafka.common.record.internal.MemoryRecordsBuilder; +import org.apache.kafka.common.utils.ByteBufferOutputStream; + +import java.nio.ByteBuffer; +import java.util.List; + +/** + * A {@link ProducerBatch} for the incremental buffer.memory allocation strategy, backed + * by a {@link MemoryRecordsBuilder} whose stream is a {@link ChunkedByteBufferOutputStream}. + * It adds mid-batch chunk extension support ({@link #extensionBytesNeeded} / + * {@link #addBuffers}) and overrides the pool deallocation hooks so all chunks are returned to + * the pool rather than a single buffer. + * <p> + * This class is not thread safe and external synchronization must be used when modifying it. + */ +public class ChunkedProducerBatch extends ProducerBatch { + + // The parent's builder reference is private; keep our own to access stream capacity state. + private final MemoryRecordsBuilder recordsBuilder; + + public ChunkedProducerBatch(TopicPartition tp, MemoryRecordsBuilder recordsBuilder, long createdMs) { + super(tp, recordsBuilder, createdMs); + this.recordsBuilder = recordsBuilder; + } + + /** + * Bytes of chunk capacity this batch needs before {@code tryAppend} could accept the given + * record. Returns 0 when no extension is needed: the batch is at its batch-size limit, or the + * attached chunk capacity already has room (always the case for an empty batch, whose stream + * is pre-sized for the first record). Positive when the record is within the batch-size limit + * but the attached chunks lack capacity — the accumulator then allocates exactly the missing + * bytes (rounded up to whole chunks) and attaches them via {@link #addBuffers} before retrying. + */ + int extensionBytesNeeded(long timestamp, byte[] key, byte[] value, Header[] headers) { + if (!recordsBuilder.hasRoomFor(timestamp, key, value, headers)) + return 0; + // Size against the batch's projected total output after this record (header counted once, + // ratio-adjusted when compressed), not per-record. Per-record sizing would over-count the + // header and miss the compressor's flush-accumulation behavior. + int target = recordsBuilder.estimatedBytesWrittenAfter(key, value, headers); + ByteBufferOutputStream stream = recordsBuilder.bufferStream(); + int totalAttachedCapacity = stream.position() + stream.remaining(); Review Comment: Both `ChunkedByteBufferOutputStream#position()` and `#remaining()` walk the entire chunk list, making this line O(chunks). Since `extensionBytesNeeded()` is called on every `ChunkedRecordAccumulator#tryAppend()`, the cumulative cost becomes O(records × chunks) for a single batch. Because every chunk has a fixed size (chunkSize), `position() + remaining()` is always equal to `chunks.size() * chunkSize` and can be computed in O(1). -- 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]
