[SPARK-59444][CORE] Retain execution-memory task registration while allocations wait

[sunchao]

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.

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:

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.