ezhou413 commented on code in PR #23483:
URL: https://github.com/apache/kafka/pull/23483#discussion_r4471659266
##########
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:
added
##########
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:
makes sense
--
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]
