peterxcli opened a new pull request, #5494: URL: https://github.com/apache/datafusion-comet/pull/5494
## Which issue does this PR close? Part of #5212 (positions 2, 3 and 8). Supersedes the closed draft #5217, replacing its explicit refcount guard with plain `Arc` lifetime. ## 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. For the unified pool types each entry transitively holds a JNI global ref to the task's `CometTaskMemoryManager`, which pins its `TaskMemoryManager` and `TaskContext` for the lifetime of the executor. There were two leak paths, plus an unsafe close path: 1. **`createPlan` failing after the pool was registered.** The pool is registered early in `createPlan`, but later steps can still fail — building the DataFusion session context (e.g. an invalid `spark.comet.datafusion.*` config), the key-unwrapper global ref, and the `TaskContext`/`ClassLoader` global refs. On the JVM side `plan` is a field initializer evaluated *before* the task-completion listener is registered, so a failed `createPlan` means `releasePlan` is never called. 2. **`releasePlan` failing before the release.** It ran `update_metrics(...)?` *before* the pool release and the `Box::from_raw`, so a metrics failure stranded both the registry entry and the native `ExecutionContext`. 3. **`CometExecIterator.close()` was not idempotent and could skip or repeat `releasePlan`.** `closed = true` was set last, so a throw from `currentBatch.close()`, `nativeUtil.close()` or a shuffle block iterator skipped `releasePlan` entirely (stranding the context), and the task-completion listener's retry re-entered the whole teardown. I verified all three paths empirically against unmodified `main` with an instrumented build and deterministic failure injections (a bogus `spark.comet.datafusion.*` key to fail `createPlan` after registration, a metrics node that throws to fail `releasePlan`, and a shuffle block iterator whose `close()` throws): | Injected failure (on unmodified main) | Observed | |---|---| | 10 × `createPlan` failure | registry grew 1→10, zero releases; 10/10 `CometTaskMemoryManager`s still pinned after 30 s of GC | | teardown throw in `close()` | `releasePlan` never called; native context and pool entry stranded | | metrics failure in `releasePlan` | `releasePlan` called **twice** (listener retry), context freed **zero** times | With this change the same three injections leave the registry empty, free every native context exactly once, and the pinned managers become collectible. ## What changes are included in this PR? **Tie the registry entry to the pool's own lifetime.** `TASK_SHARED_MEMORY_POOLS` now maps task attempt id → `Weak<TaskSharedMemoryPool>`, where `TaskSharedMemoryPool` is a transparent `MemoryPool` wrapper whose `Drop` removes its own entry. The `Arc` returned by `create_memory_pool` is the RAII handle: the pool stays registered exactly as long as the plan's session context (or any reservation) holds it, and both failure paths above are covered by construction — `createPlan` unwinding drops the `Arc`, and `releasePlan` dropping the `ExecutionContext` drops it too. A `ptr_eq` check in `Drop` prevents an old pool's drop from removing a racing replacement's entry. Using `Weak` in the map means the map itself never keeps a pool alive. **`releasePlan` reclaims the `Box` up front and flushes metrics last**, so the context (and with it the pool reference and every JNI global ref) is freed even when the metrics update fails; the metrics error is still thrown. **`CometExecIterator.close()` sets `closed` first and always calls `releasePlan`.** Teardown runs in a `try`/`catch`; the release runs unconditionally afterwards, with the teardown exception propagated and any release failure attached as a suppressed exception. This has to land together with the `releasePlan` change: now that the `Box` is always freed, a second `releasePlan` on the same pointer would be a use-after-free rather than a leak. **Deletions that fall out of the above:** `MemoryPoolType::is_task_shared`, `handle_task_shared_pool_release`, the per-plan refcount (`PerTaskMemoryPool.num_plans`), and the `ExecutionContext::task_attempt_id` / `memory_pool_config` fields, all of which existed only to feed the old explicit release call. The duplicated task-shared arms of `create_memory_pool` collapse behind an `acquire_task_shared_pool` helper. The registry moves to `parking_lot::Mutex`, matching `fair_pool.rs` in the same module. ## How are these changes tested? **New `CometExecIteratorLifecycleSuite`** with one deterministic reproducer per failure path, exercising the real JNI lifecycle: - `createPlan` failure: 10 plans created with an unknown `spark.comet.datafusion.*` key (fails inside `createPlan` after pool registration) against `fair_unified` pools; a `WeakReference` per `CometTaskMemoryManager` plus a GC loop asserts none stay pinned. The test asserts the injected error message so a missing native library cannot make it pass vacuously. - teardown throw: a `CometShuffleBlockIterator` stub whose `close()` throws; asserts the exception propagates and a second `close()` is a no-op. - metrics failure: a `CometMetricNode` subclass that throws in `set_all_from_bytes`, armed after a real plan executed; asserts the failure propagates and a second `close()` is a no-op. On unmodified `main` all three fail deterministically (first with "10 of 10 CometTaskMemoryManagers stayed reachable", the other two with a second `close()` re-throwing); with this change all three pass. **Unit tests** in `task_shared.rs` cover the registry: plans in one task share a pool, different tasks do not, the entry survives until the last reference drops, dropping the `Arc` alone releases it (the `createPlan`-unwind case), and an old pool's drop does not remove a racing replacement's entry. **Regression suites:** `CometAggregateSuite`, `CometNativeShuffleSuite`, `CometNativeSuite`, `CometTaskMemoryManagerSuite` and the new suite — 125 tests, all passing. Native: full `datafusion-comet` crate suite (167 passed), `cargo clippy --all-targets -- -D warnings` clean, `cargo fmt --check` clean, `spotless:check` and `scalastyle:check` clean. ## Are there any user-facing changes? No API or configuration changes. Long-running executors no longer accumulate memory pool registry entries and pinned JVM task objects when `createPlan`, the final metrics update, or iterator teardown fails, and a teardown failure no longer risks a double release of the native plan. One behavioral nuance: the registry entry now lives until the last `Arc` to the pool drops rather than until the last plan's `releasePlan` call, so a reservation that outlives its plan keeps the pool registered instead of the entry being removed while the pool is still in use. 🤖 Generated with [Claude Code](https://claude.com/claude-code) -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
