peterxcli commented on code in PR #5493:
URL: https://github.com/apache/datafusion-comet/pull/5493#discussion_r3879867847
##########
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 the [cross-pool
cycle](https://github.com/apache/datafusion-comet/pull/5493#pullrequestreview-5050181096)
in 3ca39cf33, taking your first option: the unsafe writer no longer retains
pre-write Comet allocations across blocking input construction. `open()` moved
from the constructor into `write()`, so the sorter's pointer array is only
allocated after Spark has finished evaluating the input iterator — a task
blocked in Spark's execution-memory pool during an eager input sort now holds
zero Comet pool bytes, and in your scenario B's 999448-byte request fits beside
its own array (1015808 free) and proceeds, breaking the cycle. Everything that
touches the sorter already runs inside `write()`, and `stop()`, peak-memory
accounting, and the task-completion listener all null-guard it; the listener
stays as a backstop for fatal errors during `write()` itself. I preferred this
over cross-pool cycle detection since the Comet allocator cannot observe
waiters in Spark's pool. Rework
ed the regression test accordingly: it now asserts a full-pool allocation
succeeds right after writer construction (nothing retained pre-write; this
fails on the previous head), then drives `write()` into a fatal mid-iteration
error and asserts the pool is fully reclaimed again afterwards.
--
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]
