sunchao commented on code in PR #5493:
URL: https://github.com/apache/datafusion-comet/pull/5493#discussion_r3870526317
##########
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();
+ } catch (InterruptedException ie) {
+ Thread.currentThread().interrupt();
+ throw e;
Review Comment:
[P2] Preserve task-kill semantics when the wait is interrupted
Throwing the preceding `SparkOutOfMemoryError` converts an intentional task
kill into a counted failure: it extends `OutOfMemoryError`, so Spark's
killed-task handler for `InterruptedException | NonFatal(_)` does not match it
and reports `ExceptionFailure`. In a real Spark 3.5.9 Comet JVM shuffle,
`killTaskAttempt(id, true)` on a small-row allocation waiter produced
`ExceptionFailure` and aborted the job with `local[2,1]`. The same-head control
killed during an input wait reported `TaskKilled`; its retry returned all 752
correct rows. Both used a healthy holder with an independent finite release.
Please preserve interruption/task-kill classification here so intentional kills
do not consume the retry budget or abort the job when that budget is exhausted.
--
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]
