dwsmith1983 opened a new pull request, #59103: URL: https://github.com/apache/spark/pull/59103
### What changes were proposed in this pull request? `ExecutionMemoryPool.acquireMemory` registered the task's `memoryForTask` entry once, before its wait loop, and read it on every pass with `memoryForTask(taskAttemptId)`. `releaseMemory` removes the entry when the task's balance reaches zero and calls `notifyAll`. A caller parked in `lock.wait()` that woke after that removal threw `java.util.NoSuchElementException: key not found: <taskAttemptId>` instead of continuing to wait or being granted memory. This PR moves the registration check to the top of each loop pass. A task whose entry was removed while it waited is re-registered at zero, and the existing `notifyAll` on registration wakes the other waiters so they recount the active tasks. The first pass behaves exactly as before. The entry can be removed under a waiting caller because `TaskMemoryManager.releaseExecutionMemory` does not take the `TaskMemoryManager` monitor that `acquireExecutionMemory` holds while it waits. Any consumer of the same task that releases through it (`MemoryConsumer.freeMemory`, or a caller on another thread) can bring the balance to zero while another acquire of that task is parked below its share. Synchronizing the release on the `TaskMemoryManager` instead would deadlock with the parked acquirer, and keeping zero-balance entries in the pool would inflate the active task count after tasks end, so the loop is the right place for the fix. One related edge case changes shape: if `releaseAllMemoryForTask` runs while a thread of that task is still parked in `acquireMemory`, the waiter previously threw the same exception; now it re-registers the task and can be granted memory, the same as a thread of that task that calls `acquireMemory` after cleanup already could. ### Why are the changes needed? The exception is not a memory error, so callers such as `UnsafeExternalSorter` do not treat it as a signal to spill; the task fails. Seen in Apache DataFusion Comet, whose native memory consumer parks in this loop while other consumers of the same task release (apache/datafusion-comet#6224, apache/datafusion-comet#6304). Comet works around it by retrying the acquire. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? New test in `MemoryManagerSuite` (runs for both the static and unified memory managers): with two tasks and no free memory, a task's acquire blocks below its 1 / 2N share, another thread of the same task frees its whole balance, and the blocked acquire must keep waiting and later be granted its 1 / N cap once the other task releases. The test fails before the fix (the future completes with the exception) and passes after. Ran `core/testOnly *MemoryManagerSuite *UnifiedMemoryManagerSuite *TestMemoryManagerSuite` and `core/scalastyle`. ### Was this patch authored or co-authored using generative AI tooling? No -- 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]
