sunchao commented on code in PR #5493:
URL: https://github.com/apache/datafusion-comet/pull/5493#discussion_r3870526306
##########
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:
[P2] Check that allocation can make progress before waiting
In opt-in on-heap JVM shuffle, this retry also runs from the unsafe sorter,
whose empty `spill()` retains its initial pointer array. I reproduced two real
Spark 3.5.9 tasks with a 1 MiB pool, 256 KiB pages,
`spark.shuffle.sort.bypassMergeThreshold=0`, pointer capacity 4096, and a first
binary row of 976 KiB each: both request 999448 bytes, but their arrays retain
65536 bytes, leaving only 983040 available. Both executor tasks remain here
after the input gate opens, and neither can reach cleanup to free its array.
Each task alone and both with a 2 MiB pool complete correctly; the exact base
Java overlay raises an error promptly. A sole oversized row also waits with an
empty pool. Please handle retained allocations and unsatisfiable requests
before entering an unbounded wait.
--
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]
