peterxcli commented on code in PR #5493:
URL: https://github.com/apache/datafusion-comet/pull/5493#discussion_r3872193083
##########
spark/src/main/java/org/apache/spark/shuffle/comet/CometBoundedShuffleMemoryAllocator.java:
##########
@@ -112,6 +118,36 @@ 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, so a
waiting task holds no
+ * pool memory itself and the tasks still holding memory can always progress
and eventually free
+ * it. Interrupting the task (e.g. task kill) aborts the wait.
+ */
+ @Override
+ public synchronized MemoryBlock allocateBlocking(long required) {
+ long size = Math.max(pageSize, required);
+ boolean logged = false;
+ while (true) {
+ try {
+ return allocateMemoryBlock(size);
+ } catch (SparkOutOfMemoryError e) {
+ if (!logged) {
+ logger.warn(
+ "Waiting for other tasks to free up {} bytes of Comet shuffle
pool memory", size);
+ logged = true;
+ }
+ try {
+ wait();
Review Comment:
Follow-up addressed in de5d00a9b. You're right — the deadlock check treated
notified-but-not-yet-resumed waiters as permanently blocked. `waitingThreads`
now maps each waiter to its pending request size, and deadlock is only declared
when, in addition to all allocated memory being retained by blocked waiters,
none of their requests fits in the free pool. In your schedule, the large
waiter that resumes first sees the small waiter's 262144-byte request fits the
983040 free bytes, so it keeps waiting; the small waiter proceeds, finishes,
and releases its array, after which the large request fits (999448 ≤ 1048576 −
32768). The genuine mutual-block case is still detected, since there no
waiter's request fits the free pool. Added a regression test with exactly your
sizes (900-KiB holder, two waiters retaining 32768-byte arrays,
999448/262144-byte requests) asserting both waiters complete after the holder
frees, regardless of wake order.
--
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]
