[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.

Codex assisted with implementation, tests, review, and this description.

[sunchao]

cc @viirya @cloud-fan @dongjoon-hyun @peter-toth could you take a look? thanks!

[viirya]

The fix looks reasonable for the concurrent acquire/release interleaving covered by the tests. Retaining the task entry while an acquisition is waiting, and removing an empty entry when the last waiter exits, appears consistent with the existing accounting and fairness rules.

My main question is about the practical use case. In the usual single-threaded task execution, once a consumer blocks acquiring memory, that thread cannot proceed to release memory through another consumer. The new test explicitly introduces another thread to produce this interleaving.

Could you point to an existing Spark execution path or a real-world integration that performs these concurrent memory operations for the same task? The API-level reproducer demonstrates the failure, but a concrete use case would help establish whether this fixes a supported usage pattern and justify the additional bookkeeping. If this is intended to support a new concurrency pattern, it would also be useful to clarify its scope and how callers coordinate task cleanup with outstanding acquisitions.

[viirya]

Could we also cover interrupting the waiter while the task still holds its original 100-byte reservation? Both variants here have zero reserved bytes by the time the waiter is interrupted. The additional case would verify that exiting the last waiter preserves an existing reservation until it is explicitly released.

[viirya]

Could we make the lifecycle conditions more explicit here? A successful acquisition can leave the task active because it now holds memory. Perhaps state that this releases currently reserved memory without canceling pending acquisitions, and that the task entry is removed only when it has neither reserved memory nor waiting acquisitions. The corresponding comment in MemoryManager should use the same wording.

[dongjoon-hyun]

Thank you for the PR, @sunchao. The bug is real. To answer @viirya's question about a concrete path: the Python UDF writer thread. PythonRunner's writer thread sets the same TaskContext (TaskContext.setTaskContext(context) in PythonRunner.scala) and drives the upstream iterator, which can spill and free execution memory, while the main task thread runs the downstream operators after the UDF and requests memory. Both share one TaskMemoryManager, and TaskMemoryManager.releaseExecutionMemory does not take the TMM monitor, so the release/wait interleaving described here can happen without any external engine.

That said, I think the fix is broader than what is reachable through the public API. See the inline comments.

[dongjoon-hyun]

Through TaskMemoryManager, at most one acquisition per task can be waiting in this pool at any time, because TaskMemoryManager.acquireExecutionMemory is synchronized (this) and holds the TMM monitor across lock.wait(). For the same reason cleanUpAllAllocatedMemory (also synchronized (this)) cannot run while a waiter exists, so releaseAllMemoryForTask never races with a waiter either. So the per-task counter, the lazy allocation / null reset, the multiple-waiter cases and the releaseAll=true variant only cover schedules that require calling this private[memory] pool directly.

A much smaller fix would be to re-register the entry at the top of the loop, e.g.

// The entry may have been removed by a concurrent release of this task's last byte
// while we were waiting. Re-register so the accounting below stays consistent.
val curMem = memoryForTask.getOrElseUpdate(taskAttemptId, 0L)
val numActiveTasks = memoryForTask.keys.size

The trade-off is a transient fairness gap: between the removal and the wake-up another task may compute its share with N-1 tasks. That is not a correctness issue, and it keeps the existing lifecycle semantics unchanged. If you prefer to keep the current structure, a mutable.HashSet[Long] allocated once would already be enough; the counter and the null handling do not buy anything through the public API.

[dongjoon-hyun]

This case (previous=false) and the pool-growth-callback case below assert a behavior change rather than the regression. On master, a task that registers with 0 bytes and is then interrupted stays active until releaseAllMemoryForTask is called; that is by design, and it is still what happens for a non-waiting acquisition that is granted 0 bytes. With the finally block, the same task is removed only if it happened to wait first, so the lifecycle now depends on whether the acquisition waited.

This also affects the "16 cases failed" negative control in the description: some of those failures are these new-behavior assertions, not the NoSuchElementException. If we go with the smaller fix, I would drop these two cases; otherwise the description should distinguish them.

[dongjoon-hyun]

Nit: ExecutionMemoryPool does not branch on MemoryMode at all (it only affects the pool name), so running every pool-level case in both modes doubles the runtime without adding coverage. Only the UnifiedMemoryManager case above benefits from the mode loop.

[viirya]

Thanks, @dongjoon-hyun. The pipelined Python UDF writer path helps explain why concurrent acquire/release operations within one task are relevant to Spark itself, beyond the API-level reproducer.

On the fix scope, I think we should distinguish preventing the missing-entry exception from changing which tasks participate in memory arbitration. Removing an entry when its allocation reaches zero is already the current policy, even if the task is still running and may request execution memory again shortly afterward. The transient reduction in the task count with the re-registration approach therefore follows the existing policy; it is not entirely new behavior introduced by that fix.

Retaining the entry while an acquisition is waiting also seems reasonable: in that case, we know there is an outstanding demand. But that is a more specific lifecycle guarantee, and the additional cleanup on interruption or callback failure should be considered separately from fixing the exception.

Another possible policy would be to retain entries until task completion, since zero allocation does not imply that a task has finished using execution memory. That would have a broader trade-off, though: tasks with no current demand would continue to limit other tasks' allocations. I would not expand this PR into that policy change.

For this PR, I lean toward the smaller re-registration fix unless there is a concrete reason that a waiting task must remain continuously counted. If we retain the waiter-based approach, could we document that requirement and keep the bookkeeping and lifecycle changes limited to what is needed for it?