fix: release task-shared memory pool references via RAII guard

[andygrove]

Which issue does this PR close?

Part of #5212 (positions 2, 3 and 8).

Rationale for this change

Entries in TASK_SHARED_MEMORY_POOLS leak whenever a plan's lifecycle does not run to completion, and a stranded entry is never reclaimed: keys are unique task attempt ids and nothing prunes the map. Each entry holds an Arc<Global<JObject>> for the task's CometTaskMemoryManager, which transitively pins its TaskMemoryManager and TaskContext, so this accumulates JVM objects for the lifetime of the executor.

There were two leak paths.

createPlan failing after the pool was acquired. create_memory_pool increments the refcount, but createPlan can still fail afterwards — local_dirs decoding, prepare_datafusion_session_context, and the key-unwrapper global ref all use ?. On the JVM side plan is a field initializer (CometExecIterator.scala:87) evaluated before the task-completion listener is registered (:139), so when createPlan throws there is no releasePlan and no listener — the refcount is never decremented.

releasePlan failing before the release. It ran update_metrics(env, execution_context)? before handle_task_shared_pool_release, so a metrics failure stranded the entry and skipped the Box::from_raw that frees the context.

What changes are included in this PR?

Tie the release to a guard rather than an explicit call. create_memory_pool now returns a TaskSharedPoolRef alongside the pool, stored on the ExecutionContext. Dropping the context releases the reference, so both paths above are covered by construction: createPlan unwinding drops the guard, and releasePlan dropping the box drops the guard.

The guard is declared as the last field on ExecutionContext so it drops after session_ctx and the root_op/stream that hold reservations. If it dropped first, the pool would leave the map while still in use and the next plan in the same task attempt would build a second full-budget pool, letting the task exceed its per-task limit. The field comment says so.

releasePlan reclaims the Box up front and flushes metrics last, so the context is freed even when the metrics update fails.

CometExecIterator.close sets closed first and releases the native plan unconditionally. This is position 8 in the epic, and it has to land with this change rather than separately: now that the Box is always freed, a second releasePlan on the same pointer is a use-after-free rather than a leak, and close() previously set closed = true only after releasePlan and traceMemoryUsage().

Setting closed first alone would have traded that double-release for a guaranteed leak — a throw from currentBatch.close(), nativeUtil.close() or a shuffle block iterator would skip releasePlan entirely with no retry, stranding the very context this PR frees. So the release now runs even when teardown fails, with the teardown exception propagated and any release failure attached as a suppressed exception.

Deletions that fall out of the above:

• MemoryPoolType::is_task_shared and the debug_assert that was its only remaining caller. Task-sharedness is now encoded solely by which create_memory_pool arms return a guard, so the predicate had become a second source of truth reconciled only by an assertion that compiles out in release builds. config.rs is now a net deletion.
• ExecutionContext::task_attempt_id and memory_pool_config — both existed only to feed the old explicit release call.

Smaller items: the duplicated task-shared arms of create_memory_pool collapse behind an acquire_task_shared_pool helper (which is also what makes the refcounting unit testable, since it takes the pool factory as a closure); TASK_SHARED_MEMORY_POOLS moves to parking_lot::Mutex, matching fair_pool.rs in the same module and the sibling THREAD_MEMORY_POOLS registry, which removes the need for poison handling entirely; the refcount uses saturating_sub; and the release path uses the Entry API so it hashes the key once.

How are these changes tested?

Five new unit tests in task_shared.rs covering the refcount lifecycle: plans in the same task share one pool, plans in different tasks do not, the pool survives until the last plan releases it, dropping the guard alone releases the pool (the createPlan-unwind case), and releasing an already-removed entry does not panic.

Full native crate suite passes (140 passed, 4 ignored). cargo clippy --all-targets -- -D warnings clean, as are spotless:check and scalastyle:check.

JVM suites exercising plan create/release against the default on-heap greedy_task_shared pool:

• CometAggregateSuite — 88 succeeded, 0 failed
• CometNativeShuffleSuite — 27 succeeded, 0 failed

What is not directly tested: the leak paths themselves. Both require injecting a failure into createPlan or update_metrics, which there is no hook for, so the unit test for the unwind case exercises the guard in isolation rather than through createPlan. Likewise the new close() behaviour is not covered by a test that makes nativeUtil.close() throw. The guard makes the release unconditional at the type level, which is the property I would want a test to establish, but it is verified by construction and inspection rather than by a test that reproduces the leak.

I also have not measured the leak's magnitude on a real workload — it requires plan-creation or metrics failures, which are not the normal path. The severity comes from entries never being reclaimed once stranded, not from a high rate of stranding.

Are there any user-facing changes?

No API or configuration changes. Long-running executors that hit createPlan or metrics failures will no longer accumulate memory pool entries and pinned JVM task objects.

Deliberately left for follow-up

THREAD_MEMORY_POOLS has the same leak and is not converted here. register_memory_pool (jni_api.rs:494) registers a (thread_id, context_id) entry holding an Arc<dyn MemoryPool>, and the jni_new_global_ref! for task_context a few lines later can still fail before the ExecutionContext exists — stranding an entry in a map nothing prunes, exactly as described above. It only affects runs with tracing enabled, and converting it is a second guard rather than a change to this one, so I have kept it out to keep this diff reviewable and will add it to #5212 as its own item. Flagging it explicitly rather than leaving the inconsistency silent: after this PR, one of the two per-plan global registries is RAII-managed and the other is not.

Sibling close paths now differ in discipline. ArrowReaderIterator and NativeBatchDecoderIterator set their closed flag last, and NativeBase uses compareAndSet. Only CometExecIterator owns a native context whose double-free matters, so I have not touched the others, but converging on one idiom would be worthwhile.