test: cover native memory accounting boundaries

[rich7420]

Which issue does this PR close?

Part of #5212.

Rationale for this change

The existing CometTaskMemoryManagerSuite test uses one manager whose requests are fully granted. It cannot catch accounting requested bytes instead of granted bytes, or reporting the task total as a manager's own usage.

What changes are included in this PR?

Add tests for partial and zero grants, release and reacquisition, and separate balances for two managers sharing a task. Assert manager usage, Spark consumer usage and the task total throughout. Extract the existing task setup into a helper that clears the task context even if cleanup fails.

How are these changes tested?

All three suite tests pass through the root Maven reactor on Spark 4.1.3 / JDK 21. Two mutation checks confirm the new tests catch the accounting errors above while the original test passes.

Fork CI passes, including Linux Spark 3.4–4.2 and macOS Spark 4.0. These tests cover JVM accounting with Spark's TestMemoryManager; they do not exercise JNI or native pool rollback.

[sunchao]

Correctness

Reviewed 3377fb5c124a against 8320ae481b4e. This changes only the memory-manager test suite; production allocation behavior is unchanged. The prior test covered a single manager receiving its entire request. The added cases distinguish requested from granted bytes and each manager's balance from the task total. With a 1,024-byte fixture budget, requests of 768, 512 and 128 bytes yield 768, 256 and 0; release/reacquisition and interleaved releases from two managers check all three accounting views after each transition.

This matches the reviewed Spark 3.5/4.0 memory-manager semantics: a partial grant remains allocated until explicitly released, and a consumer reports its own usage. The cases are deterministic JVM accounting tests using Spark's TestMemoryManager. They do not exercise real off-heap allocation, native JNI rollback, parallel callbacks, overflow or invalid release arguments. The PR description correctly limits its claim. No verified P1/P2 finding.

Validation

• I compiled and ran the exact PR Java manager and ScalaTest suite with Spark 4.1.3, Scala 2.13.17 and JDK 17. All three tests passed twice in one JVM; task context was clear after each run.
• I independently ran two mutations. Counting requested bytes failed the partial-grant test (1280 != 1024); exposing task-total usage failed the separate-manager test (128 != 0). The original test passed under both mutations, and context cleanup still completed after the failed tests.
• This was focused JVM execution, not a root-reactor build or Spark/JNI integration run. Maintained Spark 3.4/4.1 source was unavailable, so those source comparisons remain qualified.
• At 08:32 UTC on September 11, CI had 20 successful checks, 24 running, 10 skipped and one failure. The macOS scans job hit a JVM SIGSEGV; its selected suites exclude this accounting suite. The accounting suite's execution jobs remain pending. The inspected CI checkout has the assigned base/head parents and the same tree as this head; CI is not fully green.

Performance

There is no production-path change or performance claim. The added work consists of small logical reservations, accounting assertions and a few reflective field reads in tests. The 1,024-byte limit controls the fixture ledger; these checks do not allocate that amount through Comet's native allocator or measure query performance.

Design

The shared task fixture makes each test's state transitions easy to follow and gives every invocation fresh Spark memory managers. Its nested finally removes the thread-local task context even if memory cleanup throws. The two-manager case checks release ownership explicitly by confirming that changing one balance leaves the other unchanged.

Abstraction & complexity

The private fixture helper removes duplicated context setup without introducing a new test framework. Small local checkUsage helpers express the relevant invariants together, while assertions on returned grants keep the tests sensitive to the allocator's actual response. The existing reflective consumer accessor remains narrowly scoped to verifying Spark's consumer-facing view.

[andygrove]

Test only. The partial and zero grant cases are the ones that matter here, since the original test could not distinguish requested bytes from granted bytes, and it is good that you confirmed that by mutation. The only red check is the macOS scans job, which does not run CometTaskMemoryManagerSuite, so it is unrelated. LGTM.