sunchao commented on code in PR #5493:
URL: https://github.com/apache/datafusion-comet/pull/5493#discussion_r3873174944
##########
spark/src/main/java/org/apache/spark/shuffle/comet/CometBoundedShuffleMemoryAllocator.java:
##########
@@ -112,6 +128,81 @@ 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();
+ boolean logged = false;
+ 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()) {
+ throw e;
+ }
+ if (!logged) {
+ logger.warn(
+ "Waiting for other tasks to free up {} bytes of Comet shuffle
pool memory", size);
+ logged = true;
+ }
+ try {
+ wait();
+ } catch (InterruptedException ie) {
+ Thread.currentThread().interrupt();
+ // Not an allocation failure: stay non-fatal so that an
intentional task kill is
+ // classified as TaskKilled rather than ExceptionFailure.
+ throw new RuntimeException(
+ "Interrupted while waiting for Comet shuffle pool memory", ie);
+ }
+ }
+ }
+ } finally {
+ waitingThreads.remove(self);
Review Comment:
[P2] Wake remaining waiters when the only fitting request is cancelled
`anyWaiterCanProceed()` can keep other threads waiting on an entry that this
`finally` removes without notification. In a 1-MiB pool, let two large waiters
each retain 32,768 bytes and request 999,448 bytes, while a zero-retained
waiter requests 262,144 bytes. If the small waiter is interrupted, a healthy
holder frees its last page, and the large waiters reacquire first, both wait
again because the still-registered small request fits the 983,040 free bytes.
The interrupted waiter then removes itself here. Neither remaining request
fits, and there is no notification to re-evaluate the deadlock check, so they
can remain blocked absent another pool event.
With the freshly compiled current allocator, a controlled legal monitor
schedule reproduced this state in 2/20 trials (both survivors still waiting
after 500 ms); notification alone restored termination. The freshly compiled
`b8fadf58` control terminated 20/20. Source tracing maps the zero-retained
request to a bypass writer waiting for its first page: its empty failure
cleanup calls no allocator `free()`. Cancelling only that writer's independent
job with `interruptOnCancel=true` need not retry it or cancel the unsafe
waiters retaining pointer arrays; Spark task/memory-manager cleanup notifies
different monitors. This mixed-job impact is source-inferred—the dedicated
Spark cancellation control could not run within the disk reserve. Please
notify/recheck surviving waiters when removing this entry.
--
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]
