sunchao commented on code in PR #5493:
URL: https://github.com/apache/datafusion-comet/pull/5493#discussion_r3873969645
##########
spark/src/main/java/org/apache/spark/shuffle/comet/CometBoundedShuffleMemoryAllocator.java:
##########
@@ -112,6 +131,96 @@ public synchronized MemoryBlock allocate(long required) {
return allocateMemoryBlock(size);
}
+ /**
+ * Like {@link #allocate(long)}, but waits for other tasks of this shared
pool to free memory,
+ * mirroring how Spark's unified memory manager blocks a task until memory
becomes available.
+ * Callers must only use this after spilling their own buffered data. The
wait fails fast when it
+ * can never succeed: when the request does not fit next to the memory this
thread itself still
+ * retains (e.g. the sorter's pointer array), or when all allocated memory
is retained by threads
+ * that are themselves blocked here and none of their requests fits in the
free pool. Interrupting
+ * the task (e.g. task kill) aborts the wait.
+ */
+ @Override
+ public synchronized MemoryBlock allocateBlocking(long required) {
+ long size = Math.max(pageSize, required);
+ Thread self = Thread.currentThread();
+ long waitStart = 0;
+ long lastLog = 0;
+ try {
+ while (true) {
+ try {
+ return allocateMemoryBlock(size);
+ } catch (SparkOutOfMemoryError e) {
+ if (waitingThreads.put(self, size) == null) {
+ // Wake existing waiters so they re-evaluate the deadlock check
against the enlarged
+ // waiting set.
+ notifyAll();
+ }
+ // This thread cannot free what it retains while it waits, so a
request that does not
+ // fit next to its own retained memory can never be satisfied.
+ if (size > totalMemory - retainedMemory.getOrDefault(self, 0L)) {
+ throw e;
+ }
+ // The allocation just failed, so the request does not fit in the
unallocated pool.
+ // Waiting can only succeed while some thread can still free memory:
either a thread
+ // outside the waiting set retains pool memory, or another waiter's
request fits in the
+ // free pool, in which case that waiter can proceed and eventually
free what it retains.
+ if (allocatedMemory <= retainedByWaitingThreads() &&
!anyWaiterCanProceed()) {
Review Comment:
[P2] Avoid unbounded waits on orphaned constructor allocations
An owner outside `waitingThreads` is not necessarily able to free its
allocation. `SpillSorter` first constructs `ShuffleInMemorySorter(allocator, 1,
true)` (8 bytes), then allocates its default 4,096-long array (32 KiB). If the
second allocation fails, expansion never frees the first array. This happens
inside the unsafe-writer constructor, before Spark receives a writer to clean
up; Spark's task-memory cleanup does not own Comet's private pages, and
`SparkOutOfMemoryError` does not terminate the executor.
A source-derived on-heap JVM-shuffle case uses a 1-MiB pool, 256-KiB pages
and three task slots: a healthy bypass task retains 1,032,192 bytes; an unsafe
task then allocates 8 bytes and fails its 32-KiB allocation. After that
independent job fails and the healthy holder finishes, 8 orphaned bytes remain.
An already-running bypass task on a different thread requesting its first
1,048,576-byte page has zero retained bytes, but only 1,048,568 bytes are free.
Every timeout sees `allocatedMemory=8 > retainedByWaitingThreads()=0`, so this
branch waits again without a remaining release path. The payload sizes are
reachable with binary rows; mixed bypass/unsafe routing can use 2 versus 3
partitions and bypass threshold 2.
The tiny constructor leak predates this PR; changing the subsequent prompt
allocation failure into an unbounded wait is the regression. Please reclaim
failed-constructor allocations or otherwise account for departed owners before
relying on this progress condition. This is static source/cleanup-path
verification, not an executed reproduction.
--
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]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
