sunchao opened a new pull request, #58747: URL: https://github.com/apache/spark/pull/58747
### Why are the changes needed? An execution-memory task can have no bytes reserved while still waiting to acquire memory. `ExecutionMemoryPool` currently treats releasing the last byte as the end of the task's participation. If another acquisition for that task is waiting, it wakes up and indexes a task entry that has already been removed, throwing `NoSuchElementException`. For example, consider a 1,000-byte pool: | Step | Task A | Task B | What happens | | --- | ---: | ---: | --- | | Initial allocations | 100 | 900 | The pool is full. | | Another A consumer requests 300 | 100 | 900 | A waits for its minimum fair share. | | A's first consumer releases 100 | 0 | 900 | Today, A's accounting entry is removed despite the waiting acquisition. | | A wakes up | — | 900 | Looking up the removed entry throws instead of continuing to wait. | This interleaving is reproducible through Spark's `MemoryConsumer` and `TaskMemoryManager` APIs in both memory modes; it does not require an external execution engine. It requires concurrent memory operations within the same task. The tests reproduce the allocator/API schedule, not an end-to-end SQL or Python query failure. Tracks [SPARK-59444](https://issues.apache.org/jira/browse/SPARK-59444). ### What changes were proposed in this pull request? Keep a task registered until its waiting acquisitions have finished, even if it temporarily owns zero bytes. After B releases 300 bytes in the example, A can acquire its requested 300 bytes normally. The pool records a waiter only when an acquisition first needs to wait. All updates use the existing memory-manager monitor. Multiple waiters still count as one task for fairness, and a `finally` block removes each waiter on success, interruption, or an exception. When the last waiter leaves, an empty task entry is removed and other contenders are notified; a nonempty reservation remains charged. The waiter map is allocated lazily and discarded when empty. Non-waiting acquisitions do not access it. The fair-share limits, storage-reclamation callbacks, and lock ordering remain unchanged. The release-all documentation now distinguishes freeing currently reserved bytes from canceling outstanding acquisitions. ### How was this patch tested? Built and tested against Apache Spark master `80479fa48a25a28519456e5707b818f66c5d3f02`, using JDK 21: ```text build/sbt 'core/testOnly org.apache.spark.memory.ExecutionMemoryPoolSuite org.apache.spark.memory.UnifiedMemoryManagerSuite' ``` **42 tests passed:** all 18 new regression cases and 24 existing unified-memory tests. One existing unified-memory test was canceled by its Apple Silicon platform guard. The new suite covers last-byte release, release-all, multiple waiters, partial release, interruption, callback failure after waiting, and two public memory consumers sharing one task manager. Every case runs in both on-heap and off-heap mode. Worker waits and cleanup are bounded. As a negative control, I compiled the unmodified master `ExecutionMemoryPool` separately and ran the new suite against it, using the same freshly built master dependencies. **16 cases failed and the two partial-release controls passed.** The failures detect the removed task entry or stale zero-byte fairness participation. The fixed allocator passed all 18 cases. This control did not change the reviewed checkout. ### Does this PR introduce _any_ user-facing change? Yes: the concurrent release/wait schedule described above no longer fails the memory allocation. There is no API or configuration change. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: OpenAI Codex (tool version not recorded). Codex assisted with implementation, tests, review, and this description. -- 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]
