fix: [branch-1.1] retry a native acquire that Spark failed after dropping the task's entry (#6310)

[andygrove]

Which issue does this PR close?

Closes #6224.

Backport of #6310 to branch-1.1.

Rationale for this change

Spark can remove a task's execution-memory entry while a native Comet acquire is waiting, causing the task to fail with NoSuchElementException instead of allowing the operator to spill. The fix should be available in the 1.1 release line.

What changes are included in this PR?

Cherry-picks #6310 with -x. The native memory manager retries only Spark's specific missing-task-entry failure, caps retries, and adds focused regression coverage and documentation.

The only branch-specific adaptation resolves a comment conflict around the existing 1.1 native-plan identifier; the 1.1 wording is retained and the retry constant is added alongside it.

How are these changes tested?

Ran:

./mvnw test -Dtest=none -Dsuites='org.apache.spark.CometTaskMemoryManagerSuite' -Dscalastyle.skip=true

Result: 1 suite completed, 9 tests passed.

[sunchao]

Reviewed the full three-file diff at 128ec245fce09db58c1f0c7bc8b479f4aa7a6c43 against e9efd9f764ee0a59b7898ff028d6985d4a7a28e1. The PR is non-draft. The snapshot and live discussion checks contained no existing reviews or comments.

Routed skills: review-comet-pr and review-comet-memory-pr.

No introduced P1/P2 issues found within this review.

Summary

• Prior state and problem: A native acquire waiting in Spark could fail with NoSuchElementException when another consumer released the task’s last bytes, preventing normal spill handling. This backports the fix to branch-1.1.
• Design approach: Retry Spark’s specific missing-task-entry failure through acquireFromSpark, with three total attempts before returning a zero grant.
• Correctness: Spark registers the task before its wait loop and grants memory only when leaving that loop. A missing entry after waking therefore leaves no grant from that call to reconcile. Accounting advances only after a successful return. Both real race regressions pass at head and fail against the base.
• 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 wait/remove behavior is consistent. No SQL semantics, configuration, or JNI signatures change.
• Key design decisions: Recognition requires both the exact task-specific exception message and an ExecutionMemoryPool stack frame. Unrelated exceptions propagate. No additional lock is held across retries.
• Implementation sketch: One private acquisition helper and one exception predicate wrap the existing Spark call. Four added tests cover native and sibling-consumer releases, exception filtering, and retry exhaustion. The contributor guide explains the recovery boundary.
• Performance: Successful acquisitions retain one Spark call. Stack inspection and retry logging occur only on the targeted exception, and attempts are bounded. No evidence-backed performance regression was identified.
• Design: The recovery stays at the shared native-to-Spark acquisition boundary and preserves Spark’s allocation policy. This is a focused design for the backport.
• Abstraction & complexity: The two private helpers separate retry control from exception recognition without adding a framework, public API, or configuration surface. No concrete simplification met the reporting bar.
• Behavioral changes worth calling out: The targeted race now retries or returns zero. Native try_grow treats refusal as memory pressure, allowing spill-capable operators to spill. Infallible grow retains its existing overcommit handling. JVM callers that bypass this bridge remain outside the fix.
• Suggested improvements: No additional change meeting the P1/P2 reporting bar was identified.

Exact-head CI: At inspection, seven checks had succeeded, eighteen were running, and four were skipped, with no recorded failures. Linux/macOS lint and preflight checks passed. Builds, Rust tests, Spark SQL 3.5/4.0/4.1, and Iceberg validation remained underway in Comet CI. Spark SQL 3.4 was skipped. There was no final CI verdict yet.

Validation: Freshly compiled the exact-head Java manager and Scala suite against cached Spark 4.1.3 dependencies on JDK 17. All nine tests passed. Running the two real waiting-acquire tests against freshly compiled base code reproduced NoSuchElementException: key not found: 0 in both. git diff --check passed and the checkout remained unchanged. This was isolated JVM validation, not a full Maven/native build or end-to-end SQL run. Other supported versions received source comparison, not local runtime testing.