Nikita-Shupletsov commented on code in PR #23483:
URL: https://github.com/apache/kafka/pull/23483#discussion_r4383641014
##########
clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedByteBufferOutputStream.java:
##########
@@ -316,7 +395,7 @@ public int initialCapacity() {
}
/**
- * Throws if {@code requiredBytes} cannot be written to this stream.
+ * Test-only. Throws if {@code requiredBytes} cannot be written to this
stream.
Review Comment:
if it's test only, maybe it makes sense to have this logic in the test?
##########
clients/src/main/java/org/apache/kafka/clients/producer/internals/BufferPool.java:
##########
@@ -370,8 +353,148 @@ public List<ByteBuffer> allocateChunks(int totalSize,
long maxTimeToBlockMs) thr
}
}
- // Allocate raw chunks for the reserved non-pooled portion outside the
lock. On error,
- // refund all reserved bytes (memoryRequired) and let the next waiter
try.
+ return materializeChunks(pooled, numChunks, chunkSize, memoryRequired);
+ }
+
+ /**
+ * Non-blocking variant of {@link #allocateChunks} for callers that can
neither wait nor fail, such as a
+ * chunked stream growing in the middle of a record. Returns the chunks if
the memory is available right
+ * away, and an empty list otherwise, without reserving anything, joining
{@link #waiters} or recording a
+ * wait time. It never waits, so it ignores the thread's interrupt status
and leaves it untouched. Also
+ * returns an empty list once the pool is closed: the producer closes the
pool before it drains or aborts the
+ * last batches, and closing those batches still needs to grow their
streams. A stream growing
+ * mid-write, which can take a partial grant, uses {@link
#tryAllocateChunksUpTo} instead. Also returns an
+ * empty list while any thread is queued for memory, so non-blocking
callers never take memory ahead of a waiter.
+ *
+ * @param totalSize minimum total bytes of capacity required across the
returned chunks
+ * @return list of {@code ceil(totalSize / poolableSize())} {@code
ByteBuffer}s, each of capacity
+ * {@code poolableSize()}, or an empty list if the memory is not
available right now or the pool is closed
+ * @throws IllegalArgumentException if {@code totalSize <= 0}, or if the
request rounded up to
+ * whole chunks exceeds {@code totalMemory()}
+ */
+ public List<ByteBuffer> tryAllocateChunks(int totalSize) {
+ int chunkSize = poolableSize();
+ int numChunks = numChunksFor(totalSize, chunkSize);
+ long memoryRequired = (long) numChunks * chunkSize;
+ List<ByteBuffer> pooled;
+
+ lock.lock();
+ try {
+ if (this.closed)
+ return List.of();
+ // Memory free while a thread is queued is the window between a
deallocation signalling it and it re-taking
+ // the lock; taking it here would starve that FIFO waiter, so
refuse and let the caller fall back.
+ if (!this.waiters.isEmpty())
+ return List.of();
+ pooled = reserveChunksIfAvailable(numChunks, chunkSize,
memoryRequired);
+ if (pooled == null)
+ return List.of();
+ } finally {
+ lock.unlock();
+ }
+ return materializeChunks(pooled, numChunks, chunkSize, memoryRequired);
+ }
+
+ /**
+ * Best-effort variant of {@link #tryAllocateChunks} for a chunked stream
growing mid-write: returns as many
+ * of the {@code ceil(totalSize / poolableSize())} chunks as are available
right now, possibly none, and the
+ * caller covers the rest from the heap. A request larger than the pool is
clamped to the pool's size rather
+ * than rejected. Never waits; returns an empty list, never {@code null},
when nothing is free or the pool
+ * is closed. Also returns an empty list while any thread is queued for
memory, so non-blocking callers never
+ * take memory ahead of a waiter.
+ *
+ * @throws IllegalArgumentException if {@code totalSize <= 0}
+ */
+ public List<ByteBuffer> tryAllocateChunksUpTo(int totalSize) {
+ if (allocationMode != AllocationMode.INCREMENTAL)
+ throw new IllegalStateException("allocateChunks() and
tryAllocateChunks() are not supported in "
+ + allocationMode + " allocation mode; use allocate()");
+ if (totalSize <= 0)
+ throw new IllegalArgumentException("totalSize must be positive: "
+ totalSize);
+ int chunkSize = poolableSize();
+ int numChunks = (int) Math.min(((long) totalSize + chunkSize - 1L) /
chunkSize, totalMemory / chunkSize);
+ List<ByteBuffer> pooled = new ArrayList<>();
+ int granted;
+
+ lock.lock();
+ try {
+ if (this.closed)
+ return List.of();
+ // Memory free while a thread is queued is the window between a
deallocation signalling it and it re-taking
+ // the lock; taking it here would starve that FIFO waiter, so
refuse and let the caller fall back.
+ if (!this.waiters.isEmpty())
+ return List.of();
+ granted = reserveChunksUpTo(numChunks, chunkSize, pooled);
+ } finally {
+ lock.unlock();
+ }
+ if (granted == 0)
+ return List.of();
+ return materializeChunks(pooled, granted, chunkSize, (long) granted *
chunkSize);
+ }
+
+ /**
+ * Validates a chunk request and returns the number of whole chunks it
needs. Shared by
+ * {@link #allocateChunks} and {@link #tryAllocateChunks}.
+ */
+ private int numChunksFor(int totalSize, int chunkSize) {
+ if (allocationMode != AllocationMode.INCREMENTAL)
+ throw new IllegalStateException("allocateChunks() and
tryAllocateChunks() are not supported in "
+ + allocationMode + " allocation mode; use allocate()");
+ if (totalSize <= 0)
+ throw new IllegalArgumentException("totalSize must be positive: "
+ totalSize);
+ int numChunks = (int) (((long) totalSize + chunkSize - 1L) /
chunkSize);
+ throwIfChunksNeededExceedsPool(totalSize, numChunks, chunkSize, (long)
numChunks * chunkSize);
Review Comment:
nit: I think if we call throwIfChunksNeededExceedsPool in tryAllocateChunks
and allocateChunks, we would be able to reuse this method in
tryAllocateChunksUpTo as well.
##########
clients/src/main/java/org/apache/kafka/clients/producer/internals/ChunkedProducerBatch.java:
##########
@@ -64,27 +71,50 @@ int extensionBytesNeeded(long timestamp, byte[] key, byte[]
value, Header[] head
}
/**
- * Appends the record. This batch never acquires memory itself: the
capacity is arranged before
- * the append, by the accumulator attaching chunks (see {@link
#extensionBytesNeeded} and
- * {@link #addBuffers}).
+ * Appends the record. Its capacity is pre-sized by the accumulator before
the append: the batch is
+ * created with chunks for its first record (a blocking acquire), and each
later record gets extension
+ * chunks for the shortfall reported by {@link #extensionBytesNeeded} (see
{@link #addBuffers}) through a
+ * non-blocking acquire; if the pool cannot supply them, the accumulator
closes this batch and retries the
+ * record on a new one. Pre-sizing is an exact upper bound for an
uncompressed batch. For a compressed batch
+ * it scales with the topic's compression ratio estimate, which is only a
heuristic and can change between
+ * pre-sizing and batch creation, so the backing {@link
ChunkedByteBufferOutputStream} also grows during the
+ * write, attaching all the chunks the write still needs: from the pool
without blocking, the rest from the
+ * heap. The check here therefore covers only an uncompressed batch's
first record, the one append no caller
+ * checks ({@code RecordAccumulator.appendNewBatch} appends straight into
the new batch); every later append
+ * comes through {@code ChunkedRecordAccumulator.tryAppend}, which has
already evaluated
+ * {@link #extensionBytesNeeded}.
* <p>
- * The capacity is verified here only for the batch's first record, which
is the one append no
- * caller checks: {@code RecordAccumulator.appendNewBatch} creates the
batch and appends straight
- * into it, relying on the stream having been pre-sized. Every later
append arrives through
- * {@code ChunkedRecordAccumulator.tryAppend}, which evaluates {@link
#extensionBytesNeeded}
- * itself and attaches chunks before retrying, so repeating the check for
those would size the
- * record a second time on every append for no added safety.
+ * If the growth fell back to the heap, the record is still appended but
the batch is then closed for
+ * record appends, so it is sent early rather than growing further on the
heap past buffer.memory.
*
* @return the record's future, or null if the batch is at its batch-size
limit
- * @throws IllegalStateException if the stream was not pre-sized to hold
the batch's first record
+ * @throws IllegalStateException if the uncompressed stream was not
pre-sized to hold the batch's first record
*/
@Override
public FutureRecordMetadata tryAppend(long timestamp, byte[] key, byte[]
value, Header[] headers, Callback callback, long now) {
- if (recordCount == 0 && extensionBytesNeeded(timestamp, key, value,
headers) != 0)
+ if (recordCount == 0 && !isCompressed() &&
extensionBytesNeeded(timestamp, key, value, headers) != 0)
throw new IllegalStateException(
"Unexpected append to a chunked batch whose chunks lack
capacity for the record; " +
"the stream should have been pre-sized for the
batch's first record");
- return super.tryAppend(timestamp, key, value, headers, callback, now);
+ FutureRecordMetadata future = super.tryAppend(timestamp, key, value,
headers, callback, now);
+ if (future != null && stream().fallbackAllocations() > 0) {
+ // The stream grew on the heap mid-record: the pool is exhausted,
and the builder's fullness checks
+ // trust a compression ratio estimate that has proven too low.
Close for appends so later records
+ // go to a new batch (which blocks for pool memory) instead of
growing this one further on the heap.
+ // The caller's append result reads isFull(), now true, so the
sender is woken to send it early, and
+ // the drain's ProducerBatch.close() then feeds the observed ratio
back to the estimator sooner. The
+ // compressor flush in closeForRecordAppends may fall back to the
heap once more, which is bounded since
+ // nothing is appended after it. Closing is idempotent, and
matches how the extension path closes a
+ // batch on pool exhaustion.
+ log.debug("Closing batch for {} to record appends so it is sent
early: the buffer pool was "
Review Comment:
if looks like a pretty heavy log. maybe it's worth wrapping it in
`if(log.isDebugEnabled())`?
--
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]
