fix: retry a native acquire that Spark failed after dropping the task's entry

[dwsmith1983]

Which issue does this PR close?

Closes #6224.

Rationale for this change

Spark's ExecutionMemoryPool.acquireMemory registers the task's memoryForTask entry once, before its wait loop, and reads it with memoryForTask(taskAttemptId) on every pass. releaseMemory removes the entry when the task's balance reaches zero and wakes the waiters. A caller that wakes after the removal throws NoSuchElementException: key not found: <taskAttemptId>. The code is the same in Spark 3.4 through 4.1 and on master.

TaskMemoryManager.acquireExecutionMemory holds the task memory manager's monitor while it is parked, but releaseExecutionMemory takes no monitor. So a release can empty the task's balance while an acquire of the same task is parked: under greedy_unified, one native thread's release fails another native thread's parked acquire; under fair_unified, a JVM consumer of the task freeing its last bytes with freeMemory does the same. The exception reaches the native side as a plain error rather than ResourcesExhausted, so the operator cannot spill and the task fails.

What changes are included in this PR?

CometTaskMemoryManager.acquireMemory calls Spark through a helper that catches this one exception, identified by its message (key not found: plus this task's attempt id) and a frame in ExecutionMemoryPool, and calls acquireExecutionMemory again. The retry re-registers the entry, re-checks the task's share and parks again if it has to, so the caller ends up with what Spark would have granted had the entry stayed. After three attempts it returns 0, which the native side treats as a refusal and spills on. Any other exception is rethrown. A retry logs at info level, and giving up logs a warning.

The retry is safe because Spark only removes the entry at a zero balance, which cannot hold while the failed call has a partial grant. So the failed call has nothing to reconcile, and used and Spark's counters have not moved. The helper holds no lock of its own across attempts.

Retrying from the acquire side reaches the same end state as fixing it in Spark (SPARK-59444, apache/spark#58747, which re-registers a waiting task in ExecutionMemoryPool), without coordinating the release side, which does not know whether an acquire is parked. The Spark fix would also cover Spark's own operators in the same task, such as its sorters and aggregates, which this does not. CometUnifiedShuffleMemoryAllocator has its own retry from #6403.

The contributor guide's memory management page describes the new failure and why the retry is safe.

How are these changes tested?

Four tests in CometTaskMemoryManagerSuite on a 100-byte off-heap UnifiedMemoryManager:

• two threads share one CometTaskMemoryManager; the task holds 10 bytes and another task holds 90; an acquire of 20 parks below the minimum share; releasing the 10 bytes empties the balance. On main the parked acquire throws key not found: 0. With this change it completes with 20 once the other task frees its memory, and getUsed and the task's consumption both read 20.
• the same, with the 10 bytes held by a sibling MemoryConsumer of the task and released with freeMemory.
• a NoSuchElementException with another task's id, or without an ExecutionMemoryPool frame, is rethrown unchanged after one call.
• when every attempt throws the matching exception, the helper returns 0 after the cap with accounting unchanged.

CometTaskMemoryManagerSuite, CometUnboundedShuffleMemoryAllocatorSuite and CometExecIteratorLifecycleSuite pass on Spark 3.4, 3.5, 4.0 and 4.1, the spill tests in CometExecSuite pass on 3.5, and the native memory pool tests pass. With the retry disabled, the two race tests and the cap test fail again.

[sunchao]

Summary

• Prior state and problem: Releasing a task’s last reserved bytes could remove Spark’s bookkeeping entry while a native acquisition waited, causing NoSuchElementException and task failure.
• Design approach: Retry the specifically identified missing-entry exception, with three attempts before returning a zero grant.
• Correctness / compatibility analysis: Checked Spark’s memory-pool and task-manager sources for 3.4.3, 3.5.9, 4.0.4, 4.1.3 and 4.2.0. The zero-balance removal rule supports retrying without losing a partial grant. Unrelated exceptions still propagate.
• Key design decisions: Matching both the task-specific message and an ExecutionMemoryPool frame limits recovery to the intended failure. Successful acquisitions still make one Spark call. Exception inspection and retry logging occur only on failure, and no new locks are introduced.
• Implementation sketch: Two private helpers wrap acquisition. Four regression tests extend the existing suite. Reviewed all five files in the supplied base-to-head diff, including both window-documentation edits.
• Behavioral changes worth calling out: A missing entry becomes recoverable. Repeated failures return zero through the existing native refusal/overcommit handling. Direct JVM page allocations remain outside this fix, as documented in #6304.
• Suggested improvements: None meeting the P1/P2 reporting threshold. No introduced P1/P2 issues found within this review.

Reviewed full SHA 6bc7575621819b7dd5ff21cab5c07c84e58756fc against base 76493618769ff85e7a43e6f6c79f8b58d9e0b1ad. The PR remains non-draft. Routed skills: review-comet-pr and review-comet-memory-pr. The snapshot and live discussion checks contained no existing review concerns.

Validation: Freshly compiled the changed Java class and Scala suite against cached Spark 4.1.3 dependencies. All seven tests passed. Running the same suite against the base implementation reproduced failures in both race tests and the retry-cap test. This was isolated component validation. Full Maven/native builds, end-to-end spill tests, other runtime profiles and performance benchmarks were not run.

Exact-head CI: The label job passed. Comet CI, CodeQL and Check PR Title report action_required, so no build/test CI verdict is available.

[dwsmith1983]

@sunchao could you approve a CI run on 2cb2c34ce? CI has not run on this PR yet.

[sunchao]

Summary

• Prior state and problem: Releasing a task’s last reserved bytes could remove Spark’s bookkeeping entry while a native acquisition waited, causing NoSuchElementException and task failure.
• Design approach: Retry the specifically identified missing-entry exception, with three attempts before returning a zero grant.
• Correctness / compatibility analysis: Checked ExecutionMemoryPool and TaskMemoryManager sources for Spark 3.4.3, 3.5.9, 4.0.4, 4.1.3 and 4.2.0. Their zero-balance removal rule supports retrying without stranding a partial grant. Unrelated exceptions still propagate, and successful grants retain the existing accounting.
• Key design decisions: Matching the task-specific message and an ExecutionMemoryPool frame narrowly identifies the recoverable failure. Successful acquisition still makes one Spark call. Stack inspection and retry logging stay on the exception path. No new locks or shared state are introduced.
• Implementation sketch: Two private helpers wrap acquisition, four regression tests extend the existing suite, and the memory-management guide documents the behavior. Reviewed all three files in the full base-relative diff.
• Behavioral changes worth calling out: Compared with branch-1.1, this intentionally converts the missing-entry failure into recovery or the existing zero-grant refusal/overcommit handling. Direct JVM page allocations remain outside this fix, as tracked in #6304.
• Suggested improvements: None meeting the reporting threshold. No introduced P1/P2 issues found within this review.

Reviewed full SHA 790d6fcaee799ebd162d83472b522b470bf0f92a against base fef94f6cd78b18151dff57b7a936798385356de5. The PR remains non-draft. Routed skills: review-comet-pr and review-comet-memory-pr. Read the snapshot and live discussions. No unresolved review blockers were identified.

Validation: Freshly compiled the current Java class and Scala suite against cached Spark 4.1.3 dependencies using JDK 17. All nine tests passed. Running the same suite against the supplied base implementation reproduced failures in both race tests and the retry-cap test, with the other six passing. This was isolated component validation. Full Maven/native builds, end-to-end spilling, other runtime profiles and performance benchmarks were not run.

Exact-head CI: The label job passed. Comet CI and CodeQL report action_required, so no build/test CI verdict is available.

[dwsmith1983]

@andygrove this is current with main and ready for review. The CI workflow for the head commit still needs an approval to start.

[dwsmith1983]

@andygrove @viirya CI passed on a9db5f52e and it still merges cleanly with main. #6813 rebases onto this once it lands, so could one of you merge it?

[viirya]

Thanks @dwsmith1983, and sorry this sat for a while. I went through the Spark side again on 3.4 through master, and the core argument holds. Any earlier grant inside the same acquireExecutionMemory call keeps the task's balance above zero, so the call that fails with the missing entry really was granted nothing, and retrying is safe. Matching the exact task id plus the ExecutionMemoryPool frame is also tighter than what #6403 does, which is good. The race tests fail on the old code and pass here, so they cover the change itself.

I'd like to land this as it is. There are a few things I'll pick up in #6813, since that PR already moves both call sites onto one helper:

• Plan teardown. releasePlan expects that a Tokio worker may be parked in Spark's acquireMemory while it drops the stream. If that worker holds nothing itself, dropping the stream's reservations can take the task's balance to zero. Before this change the worker failed fast. Now it registers the task again and can park a second time while holding the TaskMemoryManager monitor, so cleanUpAllAllocatedMemory waits for other tasks to free memory. I'll first try to reproduce this with a test where the retry parks again after the task is marked completed. If it does, the shared helper will stop retrying once the task is completed or interrupted.
• The doc paragraph. A zero grant is a refusal on the try_grow path but overcommit on grow. The "#6304" sentence and the shuffle paragraph's "Only the shuffle allocator's callers are guarded" also go stale once both paths retry. #6813 will fold them into one paragraph.

One thing that has to happen before merge: the PR description becomes the squash commit message, and SPARK-59827 was closed as a duplicate of SPARK-59444. Could you switch it to SPARK-59444 (and apache/spark#58747), to match what the shuffle paragraph from #6403 uses? Whoever merges can also fix it in the squash message.

[dwsmith1983]

Could you switch it to SPARK-59444 (and apache/spark#58747), to match what the shuffle paragraph from #6403 uses?

Switched. The description now cites SPARK-59444 and apache/spark#58747, and it no longer lists the shuffle allocator as uncovered, since #6403 added its retry.