[SPARK-59444][CORE] Re-register execution-memory tasks after concurrent release - #58747

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.

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.

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?

Thanks for the PR, @sunchao!

acquireMemory reads memoryForTask(taskAttemptId) with apply after every lock.wait(). A concurrent release of the task's last byte deregisters it, so the waiter throws NoSuchElementException on wake-up. Retaining the entry while an acquisition is waiting is the right shape for that. I reran your negative control and it reproduces exactly: 16 of the 18 new cases fail against 80479fa48a2, and the 2 that pass are the partial-release ones.

Two points already on the thread that I reached independently, so I am not opening threads for them. The TaskMemoryManager monitor argument at r3998051507 holds - acquireExecutionMemory keeps synchronized (this) across the blocking lock.wait(). So does r3998051511: memoryMode reaches nothing but poolName in this class.

Non-blocking

• 1. Cleanup half of the waiter bookkeeping is untested: the suite pins retaining the entry in depth and never checks that the waiter count is released. I leaked it on the success path only and all 18 cases still passed. inline: ExecutionMemoryPoolSuite.scala:152

Minor

• 2. Lifecycle note is on the wrappers, not on the method that enforces it: releaseMemory carries the new guard and its scaladoc is unchanged, while both delegating callers got the note. inline: ExecutionMemoryPool.scala:205

Review summary

The narrowed fix does what the PR describes. Each pass of ExecutionMemoryPool.acquireMemory now re-registers the task before it reads the task's entry. A waiting acquisition whose last bytes another consumer of the same task releases therefore re-enters normal fair-share allocation instead of throwing NoSuchElementException. Keeping notifyAll() with the insert is needed for the three-task wake-up case, because re-adding a task lowers the minimum share other waiters need. This review found no defects in the change or in the new ExecutionMemoryPoolSuite. The earlier inline requests are addressed at this revision: the waiter bookkeeping and the lifecycle-change tests are gone, an interruption case now checks that an existing reservation survives, pool-only cases run in one memory mode, and the two-consumer test checks that a completed request leaves no entry behind. One question for the author about task cleanup follows.

Findings

0 total: 0 P0, 0 P1, 0 P2, 0 P3.

No findings.

Decision challenges

Loop-head re-registration can revive a task after cleanUpAllAllocatedMemory releases it

This follows up on the inline discussion about whether task cleanup can race a waiting acquisition, and on the question of how callers coordinate cleanup with outstanding acquisitions.

Verified in the source: TaskMemoryManager.cleanUpAllAllocatedMemory holds synchronized (this) only through TaskMemoryManager.java:928. Its final memoryManager.releaseExecutionMemory(...) (931-932) and releaseAllExecutionMemoryForTask(taskAttemptId) (934) run after the monitor is released. So a thread of the same task that blocked on the monitor during cleanup, or that calls acquireExecutionMemory right after that section, can be waiting in ExecutionMemoryPool.acquireMemory below its 1/2N share when those calls remove the task's entry.

• On the base branch, that waiter throws NoSuchElementException when it wakes and leaves no state behind.
• With this PR, the loop-head block (ExecutionMemoryPool.scala:109-112) re-registers the finished task. A grant made then is charged to the task after cleanup has returned the freed total that the executor's leak check uses, so the check never sees it. If the waiter waits again and is interrupted, it leaves a zero-byte entry instead. Unless a later release for that task removes the entry, it lasts for the executor's lifetime: granted bytes stay unavailable, and the entry counts as one more active task in every other task's 1/N share.

This does not affect the concurrent-release case the PR fixes. Acquisitions that start after line 934 already behave this way on the base branch, so the PR widens an existing lifecycle gap rather than adding a new kind of leak.

On Spark's own threads: the pipelined Python writer is joined at task completion, and PipedRDD and the continuous-processing reader interrupt their threads, so a later lock.wait() throws. The script-transformation writers (SparkScriptTransformationExec, HiveScriptTransformationExec) and the R writer (BaseRRunner) also set the task context and pull rows through the child plan. Nothing joins or interrupts them at completion.

Open question: can a supported caller, Comet in particular, still have an acquisition waiting when cleanup reaches those release calls? And is memory granted at that point ever freed? If the schedule is unsupported, a short note saying why would settle this. If it is supported, two options:

• Have a waiter whose task was released through releaseAllMemoryForTask fail or return 0 instead of re-registering. An ordinary zero-byte release would keep this PR's behavior.
• Move the two release calls into cleanUpAllAllocatedMemory's synchronized block, so they serialize with same-task acquisitions.

Not verified by this review (environment unavailable). Repro sketch:

In ExecutionMemoryPoolSuite (package org.apache.spark.memory), using the conf of the existing two-consumer test: (1) Subclass UnifiedMemoryManager(conf, 1000L, 500L, 1) and override releaseExecutionMemory without synchronized: when the current thread is the cleanup thread, count down a 'paused' latch and await a 'resume' latch, then call super. Its first call from cleanUpAllAllocatedMemory is the one at TaskMemoryManager.java:931, after the monitor is released. (2) With ON_HEAP TestMemoryConsumers, the peer task (2) acquires 900 bytes and task 1's owner acquires 100. (3) Run task 1's cleanUpAllAllocatedMemory() on a thread named 'cleanup' and await 'paused'. (4) Start requester.acquireMemory(300L) on another thread and use awaitWaiting() to confirm that it waits in ExecutionMemoryPool.acquireMemory. (5) Count down 'resume', check that cleanup returns 100, then call peer.freeMemory(300L). Expected with this PR: the requester gets 300 bytes, and memory.getExecutionMemoryUsageForTask(1L) is still 300 after cleanup returned. Reference on the base branch: the requester fails with NoSuchElementException and task 1 holds 0 bytes. Control without the pause: start the requester first, then cleanup on its own thread, and wait until cleanup is BLOCKED on the TaskMemoryManager monitor. After peer.freeMemory(300L), the requester gets 300, cleanup returns 400, and task 1 holds 0 bytes on both branches.

This follows up on existing discussions 1, 2, 3.

Generated by Omnigent on Databricks