fix: RAII for memory pool registration

[peterxcli]

Which issue does this PR close?

Part of #5212 (finding #20 and the correctness portion of finding #19).

Rationale for this change

Tracing registers each execution context's memory pool in THREAD_MEMORY_POOLS, but registration currently happens without a lifetime guard. A failure during later plan creation or during the metrics update in releasePlan skips unregistration and retains the pool indefinitely.

Debug tracing also registers each plan's distinct LoggingMemoryPool wrapper. Two contexts sharing one 4,096-byte task pool therefore report 8,192 bytes instead of the 4,096-byte ground truth.

On unmodified upstream main, a deterministic failure reproducer left 100 registry entries after plan-creation failures and 200 after an additional 100 metrics failures. Those entries reported 819,200 bytes and retained 200 extra references to the shared base pool.

What changes are included in this PR?

• Register the task-shared base pool before constructing a per-plan logging wrapper.
• Store registration in a drop guard so every post-registration creation error removes its entry.
• Move the guard out before the fallible release metrics update, preserving the raw execution context for the existing JVM retry while cleaning the tracing registry on error.
• Add a focused regression covering two debug wrappers over one base pool, exact reserved-byte deduplication, repeated creation/metrics failures, and retained-reference cleanup.

The existing registry locking and scan algorithm are unchanged.

How are these changes tested?

• cargo test -p datafusion-comet --no-default-features execution::jni_api::tests::thread_memory_pool_registration_is_scoped_and_deduplicates_base_pool --profile ci -- --exact --nocapture --test-threads=1
  • 1 passed.
• cargo test -p datafusion-comet --no-default-features --lib --profile ci -- --test-threads=1 --skip parquet::objectstore::s3::tests::
  • 126 passed.
• cargo clippy -p datafusion-comet --no-default-features --lib --tests --profile ci -- -D warnings
• cargo fmt --check --all
• git diff --check

An unfiltered library run passed 155 of 164 tests. Eight S3 credential tests failed because this host has no usable native root certificates; a representative failure reproduced with the unmodified-main binary. The remaining order-sensitive cache-refresh failure passed independently on both versions.

A warmed optimized A-B-B-A comparison used 31 interleaved samples per arm:

Workload	Upstream main	This PR	Change
Tiny, one context/pool	95.4 ns/op	98.9 ns/op	+3.6%
Two contexts sharing one pool	110.5 ns/op	109.0 ns/op	-1.3%
Registration lifecycle	193.4 ns/op	187.9 ns/op	-2.9%

Per-arm IQRs overlapped for all three workloads.

[sunchao]

Summary

Reviewed the complete one-file diff at head a34fc27b3d0a9ba622a07faa621325718acc9496 against base bd7dc601aa9a99289530598d4778c40e6ded7772. The change scopes tracing-registry ownership to a native execution context and makes shared-pool accounting independent of per-plan debug wrappers.

Five independent passes covered correctness and edge cases; compatibility and semantics; tests and CI; performance, concurrency, and resource safety; and adversarial cross-checking and discussion deduplication. After reconciliation and source validation, I found no outstanding P1/P2 issue in this scoped change and am approving it.

Prior state and problem

THREAD_MEMORY_POOLS owns an Arc for each registered context, but previously registration and unregistration were separate calls without a lifetime guard. A failure after registration during context construction, or an error from the metrics update at the start of releasePlan, could bypass tracing-registry cleanup and retain that reference.

The registry already deduplicates pools by their allocation address. When memory debugging is enabled, however, registering each context's distinct LoggingMemoryPool wrapper gives the same task-shared base pool multiple identities. Its reservation can consequently be counted once per wrapper instead of once per base allocation.

Design approach

ThreadMemoryPoolRegistration records the registering thread and context IDs and unregisters on drop. createPlan constructs this optional guard immediately after obtaining the base memory pool, before creating a logging wrapper, and transfers it into ExecutionContext only when construction succeeds.

releasePlan takes the guard out before the fallible metrics update. An error then drops the local guard, while the raw context remains alive for the JVM's existing opportunity to retry close. On a normal release, the consuming unregister method returns the remaining total for the existing tracing emission and disables a second unregister from Drop.

Correctness / compatibility analysis

The runtime still receives the logging wrapper when configured, so consumer registration, allocation limits, growth, shrinking, and debug logging retain their existing behavior. Only the identity stored in the tracing registry changes. Task-shared and global pool factories return cloned base Arcs, while per-plan pools remain distinct, matching the registry's pointer-based deduplication.

The context and DataFusion reservations retain their own pool references, so removing a tracing reference does not free a pool still used by execution. The guard preserves the original thread ID for cleanup, and taking its Option prevents a later release retry from unregistering the same registration again. Tracing-disabled execution does not create a registration.

This change is scoped to the tracing registry. It does not repair the separate pre-existing task-shared num_plans rollback problem or guarantee disposal of a raw context after a persistent metrics failure; those should not be confused with this guard's cleanup guarantee.

Key design decisions

• Register the shared base allocation, while retaining per-plan wrappers for actual execution and logging.
• Let Rust ownership cover fallible construction paths instead of adding cleanup beside each ? return.
• Move registration ownership before metrics publication without prematurely freeing the raw context used by the existing JVM close path.
• Preserve the existing registry mutex, per-thread buckets, and deduplicating scan; this is not a redesign of the registry's locking or complexity.

Implementation sketch

The patch adds the guard and one optional field to ExecutionContext, moves registration before wrapper creation, and replaces the manual release-time unregister with the guard's consuming method. The guard's ordinary destructor handles early returns, while the successful release path explicitly obtains the remaining reservation total.

The new regression test constructs two wrappers over one reserved base pool, checks registry entry counts and deduplicated bytes, repeats error-scope cleanup, and verifies that registry references do not keep the pool alive. These are direct guard tests rather than JNI fault-injection tests.

Behavioral changes worth calling out

With tracing and memory-debug logging enabled, contexts sharing one base pool now contribute that pool's reservation once per thread. Failed post-registration construction and failed release metrics updates stop retaining a tracing-registry entry. Normal successful release continues to emit the remaining registered total.

I inspected the native CI job: 901 tests passed and 4 were skipped, including an explicit pass for the new regression. Its tested merge commit has the exact reviewed base and head as parents. Five additional local source-extracted lifecycle checks passed, covering explicit unregister, shared/distinct pools, error and unwind cleanup, release retry ownership, and cross-thread destruction; those checks use pool/mutex stand-ins and are not native/JNI integration runs.

CI is not entirely green: the Spark 4.1 SQL shard failed before SQL tests started because downloading Maven 3.9.6 returned HTTP 429. This is an infrastructure failure, not a test failure attributable to this diff. No workflow was rerun as part of this review.

Suggested improvements

No P1/P2 code changes are requested for this patch. The remaining CI action is to rerun the Spark 4.1 shard that failed downloading Maven; until that succeeds, the current checks should not be described as entirely green.

[sunchao]

Merged, thanks @peterxcli !