peterxcli commented on code in PR #5493:
URL: https://github.com/apache/datafusion-comet/pull/5493#discussion_r3870594868
##########
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:
Addressed in b8fadf587. The allocator now tracks the pool memory retained
per thread (owner recorded per page) and the set of threads blocked in
`allocateBlocking`, and gives up before waiting when the wait can never
succeed: (1) when the request does not fit next to the requester's own retained
memory — this covers your sorter scenario (each task requesting 999448 bytes
while retaining its 65536-byte pointer array fails immediately, since 999448 >
1048576 − 65536 even if everything else were freed) and the sole oversized row
on an empty pool; and (2) when every allocated byte is retained by threads that
are themselves blocked waiting, so nothing can be freed anymore (the
mutual-block case where each request would individually fit). A new waiter
wakes existing waiters so they re-evaluate (2) against the enlarged waiting
set. Both cases have regression tests that fail fast (<25 ms) where they
previously hung; the healthy-holder wait scenario is unaffected since a
non-waiting ho
lder keeps the wait alive.
--
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]
