fix: share the missing task entry retry and stop retrying it after the task has ended

[viirya]

Which issue does this PR close?

None. Follow-up to #6310 and #6403, which added the two retries this PR merges.

Rationale for this change

#6403 made CometUnifiedShuffleMemoryAllocator retry an allocation that Spark fails with NoSuchElementException: key not found: <taskAttemptId> after the task's entry was dropped from ExecutionMemoryPool (SPARK-59444, fix pending in apache/spark#58747). #6310 added its own copy to CometTaskMemoryManager.acquireMemory. As #6403's description said, the two call sites should share one helper.

The shared retry also has a gap at task teardown. When a task ends, releasePlan drops a plan's stream while a Tokio worker may still be waiting in Spark's acquireMemory for the memory that stream holds. If the release empties the task's balance, the worker loses the task's entry. Before #6310 it failed at once. With the retry, it registers the ended task again and, below its minimum share, waits for other tasks to release memory. While it waits it holds the TaskMemoryManager monitor that the executor's cleanUpAllAllocatedMemory needs.

What changes are included in this PR?

• MissingTaskEntryRetry in org.apache.comet holds the retry and the check for Spark's missing-entry exception, moved out of CometUnifiedShuffleMemoryAllocator (fix: retry a shuffle page allocation that Spark failed after dropping the task's entry #6403) and CometTaskMemoryManager (fix: retry a native acquire that Spark failed after dropping the task's entry #6310).
  • CometTaskMemoryManager passes its task id and keeps fix: retry a native acquire that Spark failed after dropping the task's entry #6310's exact key not found: <id> match. The shuffle allocator keeps fix: retry a shuffle page allocation that Spark failed after dropping the task's entry #6403's prefix match.
  • When the retries run out, the native pool still gets a zero grant and the shuffle allocator still throws SparkOutOfMemoryError. The warning now carries the last exception for both.
• The helper no longer retries once the task has completed or been killed (TaskContext.isCompleted() / isInterrupted()): it refuses at once, with a zero grant or SparkOutOfMemoryError as above. Both callers pass the TaskContext they were created in, because native threads have none of their own. Spark runs completion listeners without holding the context's monitor, so a worker calling isCompleted() does not wait on releasePlan.
• The fixtures and the waiting-request harness of both suites move to a shared TaskMemoryTestUtils trait.
• memory_management.md describes the retry once for both call sites. It says a zero grant is a refusal for try_grow (the operator spills) and overcommit for grow, covers the teardown case, and drops the stale A JVM consumer's parked page allocation fails when another consumer of the task empties its balance #6304 reference.

How are these changes tested?

• The race tests from fix: retry a shuffle page allocation that Spark failed after dropping the task's entry #6403 and fix: retry a native acquire that Spark failed after dropping the task's entry #6310 now run through the shared harness against a real UnifiedMemoryManager. CometTaskMemoryManagerSuite keeps one race test, the refusal test (which now also checks that the warning carries the exception), and fix: retry a native acquire that Spark failed after dropping the task's entry #6310's rethrow test for another task's id.
• New tests mark the task completed or killed while a native acquire waits in Spark, then release the task's last bytes. Without the guard, the acquire registers the task again and is still waiting in Spark 10 seconds later. With it, the acquire returns 0 and cleanUpAllAllocatedMemory completes.
• A shuffle allocator test covers the same refusal after the task has ended.
• CometTaskMemoryManagerSuite, CometUnifiedShuffleMemoryAllocatorSuite, and CometUnboundedShuffleMemoryAllocatorSuite pass locally on Spark 4.1. The run-all-spark-profiles label covers the other profiles.

[andygrove]

This looks like the same fix as #6310 from @dwsmith1983. That PR retries the missing-entry exception in CometTaskMemoryManager.acquireMemory and closes #6224, which is the same bug as #6812. It's approved and its CI is green. When #6403 merged yesterday, its description said #6310 would move both call sites onto one shared helper, and this PR does that move. MissingTaskEntryRetry and the fixtures in TaskMemoryTestUtils are mostly #6403's code moved out of the allocator and its suite.

Could you and @dwsmith1983 agree on which PR carries the fix? If it's this one, could it also close #6224, with #6812 closed as a duplicate, and could the description credit #6310 and #6403?

[andygrove]

CI is red on every build job, even though all the suites pass. The jobs fail afterwards in check-jar-contents, because org/apache/spark/MissingTaskEntryRetry.class isn't on the allowlist in dev/ensure-jars-have-correct-contents.sh. The helper doesn't use anything private to Spark, so could it move under org.apache.comet? If it needs to stay in org.apache.spark, the allowlist needs an entry and a reason.

[dwsmith1983]

Could you and @dwsmith1983 agree on which PR carries the fix?

@viirya #6310 has been open since September 28 with CI green on its current head, and #5613 builds on it. I'd land #6310 as it is and have #6813 follow as the change that moves both call sites onto one helper, which #6403 said the second PR would do and #6310 didn't. #6812 would then close as a duplicate of #6224. The other order works too: #6310 takes the move, and #6813 closes.

[comphead]

The retry and the zero grant look safe for both native pools from reading the code. The inline comments cover the duplicated race tests, the harness that copies the shuffle allocator suite, and a doc paragraph that a merge with #6310 would silently duplicate.

[comphead]

These two race tests reach Spark through the same call. From reading Spark 3.4.3 and 4.1.3, CometTaskMemoryManager.releaseMemory and MemoryConsumer.freeMemory both end in TaskMemoryManager.releaseExecutionMemory against the task id, so I expect the releaser's type not to change what Spark or acquireMemory does. Could we keep one of the two race tests and the refusal test?

[comphead]

checkWaitingAcquireSurvives is close to a copy of checkWaitingAllocationSurvives in CometUnifiedShuffleMemoryAllocatorSuite. The pool setup, the other task's 90 bytes, the waiting thread, awaitWaitingInSpark, the cleanup in finally and the single INFO event check are the same. Only the holder, the request and the final used checks differ. Since this PR adds TaskMemoryTestUtils to share these fixtures, could the harness move there too, taking the holder's acquire and release and the request as functions?

[comphead]

A note for whichever of this PR and #6310 lands second. git merge-tree reports conflicts only in CometTaskMemoryManager.java and CometTaskMemoryManagerSuite.scala. memory_management.md merges cleanly and keeps both #6310's "A parked acquire can wake up to a missing task entry" paragraph and this bullet, which describe the same retry. That paragraph also says the allocatePage gap is tracked in #6304, which #6403 has since closed. Could the rebase keep only one of the two?

[viirya]

I'd land #6310 as it is and have #6813 follow as the change that moves both call sites onto one helper, which #6403 said the second PR would do and #6310 didn't.

Sounds good to me. Thanks @dwsmith1983

[sunchao]

Reviewed the entire seven-file diff from 00a4b422f0ed6560ef76dce004b94c3697613108 to 55a6f70fa69f0a00638d68b013d4247e8e516aa7. The PR is open and not a draft. Read the existing reviews, conversation comments, inline comments and threads. Routed skills: review-comet-pr, review-comet-memory-pr and review-comet-shuffle-pr.

No additional introduced P1/P2 issues found within this review. The previously reported P2 packaging blocker remains unresolved.

Summary

• Prior state and problem: JVM shuffle allocations already retried Spark’s missing-task-entry race, but CometTaskMemoryManager.acquireMemory did not. A waiting native reservation could fail when another consumer released the task’s last bytes.
• Design approach: Extract the existing retry policy into MissingTaskEntryRetry and reuse it for native reservations. Keep the caller-specific refusal behavior: zero bytes for native memory and SparkOutOfMemoryError for shuffle allocations.
• Correctness: Checked Spark’s registration, wait, grant and release paths. The matching exception occurs before that pool request grants memory, and a legitimate outstanding grant prevents the task balance from reaching zero. Retrying does not double-count a grant. Comet updates used only after a successful return. Both native pools interpret zero through their existing refusal and overcommit paths. Unrelated exceptions still propagate. No additional P1/P2 correctness issue identified.
• Compatibility analysis: Compared ExecutionMemoryPool, TaskMemoryManager and MemoryConsumer sources for Spark 3.4.3, 3.5.9, 4.0.4, 4.1.3 and 4.2.0. The relevant race and retry assumptions hold across these versions. Spark 4.2’s atomic consumer accounting does not change this conclusion. No SQL expression semantics, configuration defaults or shuffle formats change.
• Key design decisions: Limit requests to three attempts, identify the exception using both its message prefix and an ExecutionMemoryPool stack frame, preserve the last failure for diagnostics, and leave the refusal policy with each caller. These choices keep the workaround narrowly scoped.
• Implementation sketch: Add the shared Java helper, wrap the native acquire, delegate the existing shuffle retry to it, extract TaskMemoryTestUtils, add five native tests, and update memory_management.md. The existing shuffle tests retain their behavioral checks.
• Performance: Successful requests still call Spark once. Exception-stack inspection and retry logging occur only after the matching failure. The generic callback and boxed return introduce some implementation overhead, but no measured P1/P2 performance regression was established. No benchmark was run.
• Design: The fix belongs at the shared Spark-memory boundary and preserves Spark’s arbitration. It changes reservation failure handling without claiming to solve container RSS limits or native reclamation. The different native and shuffle refusal contracts remain explicit.
• Abstraction & complexity: One helper appropriately shares the exception classifier, retry bound and logging across two callers returning different types. The shared test fixture centralizes task setup, failure injection and logger cleanup. No additional abstraction concern meets the P1/P2 bar.
• Behavioral changes worth calling out: Compared the touched paths with branch-1.1 at e9efd9f764ee0a59b7898ff028d6985d4a7a28e1. Native requests now retry this race and can return a normal refusal instead of propagating it. Spillable operators can consequently spill, while grow retains its overcommit policy. Shuffle retry behavior differs from that release because of the earlier #6403 change already present in this PR’s base. This PR preserves that behavior.
• Suggested improvements: Resolve the already-reported packaging blocker before merge. No additional P1/P2 changes are requested.

Exact-head CI: 24 checks succeeded, 24 failed and 37 were skipped. Failures include the Spark-profile build shards, both TPC build jobs and their aggregate gates. Inspected execution logs for all five Spark profiles show all ten CometTaskMemoryManagerSuite tests passing, followed by check-jar-contents rejecting org/apache/spark/MissingTaskEntryRetry.class. The default shuffle log likewise shows its five allocator tests passing before the same rejection. This substantiates the existing unresolved packaging comment, so it is not repeated as a new finding.

Validation: Freshly compiled the relevant head Java classes and both Scala suites in a disposable directory, using Spark 4.1.3 and JDK 21. All 15 tests passed. Replacing only CometTaskMemoryManager with the base version made both new race tests fail with NoSuchElementException: key not found: 0. The repository’s JAR-content checker also rejected a temporary JAR containing the freshly compiled helper, reproducing the existing blocker.

Validation limits: Local testing was isolated JVM validation, not a full Maven/native build or an end-to-end native spill/replay run. Spark SQL, Iceberg and macOS checks were skipped in exact-head CI. Project source was left unchanged, and nothing was published.

[viirya]

Thanks both. Per the plan with @dwsmith1983, I'll rebase this onto #6310 once it lands, turning it into the follow-up that moves both call sites onto one helper. In that rebase I'll move the helper under org.apache.comet for the jar check, keep one race test plus the refusal test, move the waiting-request harness into TaskMemoryTestUtils, and keep a single doc paragraph (dropping the stale #6304 reference).

[viirya]

Rebased onto #6310 and addressed the review:

• MissingTaskEntryRetry now lives under org.apache.comet, which fixes the check-jar-contents failure. It takes an optional task id, so CometTaskMemoryManager keeps fix: retry a native acquire that Spark failed after dropping the task's entry #6310's exact match.
• CometTaskMemoryManagerSuite keeps one race test plus the refusal and rethrow tests. The waiting-request harness now lives in TaskMemoryTestUtils, shared with the shuffle suite.
• memory_management.md has a single paragraph for both call sites. The stale A JVM consumer's parked page allocation fails when another consumer of the task empties its balance #6304 sentence and "Only the shuffle allocator's callers are guarded" are gone.

One behavior change was added, deferred from the #6310 review: the retry no longer runs once the task has completed or been killed. A test reproduces the issue. Releasing the task's last bytes after the task ended made the retried acquire register the task again and wait in Spark while holding the TaskMemoryManager monitor. It is in a separate commit, and the title and description now cover it.