peterxcli commented on code in PR #5493:
URL: https://github.com/apache/datafusion-comet/pull/5493#discussion_r3870595305
##########
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:
Addressed in b8fadf587. The interrupt handler no longer rethrows the pending
`SparkOutOfMemoryError`; it restores the interrupt flag and throws a
`RuntimeException` wrapping the `InterruptedException`, which is NonFatal and
therefore matched by Spark's killed-task handler — an intentional
`killTaskAttempt` on a waiter is now classified as `TaskKilled` instead of
`ExceptionFailure`, so it does not consume the retry budget. Added a regression
test that interrupts a blocked waiter and asserts the surfaced error is not an
`OutOfMemoryError` and carries the `InterruptedException` cause.
--
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]
