peterxcli commented on code in PR #5493:
URL: https://github.com/apache/datafusion-comet/pull/5493#discussion_r3874372158
##########
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:
Addressed in a21f3f19e by reclaiming the failed-constructor allocations —
your first option, since it restores the invariant the progress condition
relies on (every allocated byte has an owner with a cleanup path) instead of
trying to enumerate departed owners, which is unreliable anyway with pooled
executor threads that outlive their tasks. `SpillSorter`'s constructor now
wraps the pointer-array allocation and, on any failure, frees both the array
(if allocated but not yet adopted) and the sorter's initial one-entry array
before rethrowing. I audited the other consumers of the bounded allocator for
the same shape: `CometDiskBlockWriter` allocates only lazily after
construction, `CometShuffleExternalSorter` allocates nothing before creating
its `SpillSorter`, and the mid-task paths (`reset`, `growPointerArray`) belong
to writers Spark already cleans up via `stop(false)` — so this constructor was
the only orphan window. Added a regression test with your sizes: a holder
retainin
g 1,032,192 bytes of the 1-MiB pool, a `SpillSorter` constructor that fails on
its 32-KiB array after the 8-byte one succeeds, and after the holder frees, a
full-pool 1,048,576-byte allocation must succeed — it fails forever if the 8
bytes leak.
--
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]
