andygrove opened a new pull request, #5217: URL: https://github.com/apache/datafusion-comet/pull/5217
## Which issue does this PR close? Part of https://github.com/apache/datafusion-comet/issues/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`. That ordering matters: the pool's own reservations must be released before the last reference removes it from the map, otherwise `CometUnifiedMemoryPool::drop` can fire its "dropped with N bytes still reserved" warning spuriously. **`releasePlan` reclaims the `Box` up front and flushes metrics last**, so the context is freed even when the metrics update fails. **`CometExecIterator.close` now sets `closed` before tearing down.** 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. `close()` previously set `closed = true` only after `releasePlan` and `traceMemoryUsage()`, so a throw from either left the iterator eligible for a second `close()` from the task-completion listener. **Smaller items in the same code:** - The three duplicated `or_insert_with` blocks in `create_memory_pool` collapse into one `acquire_task_shared_pool` helper. This is what makes the refcounting unit testable — the helper takes the pool factory as a closure, so tests can pass an `UnboundedMemoryPool` and exercise the map without a JVM. - The refcount uses `saturating_sub` and the map lock recovers from poisoning instead of propagating it (position 3 in the epic). Both are in the function being rewritten: an underflow panic while holding that lock would poison it, and since every `createPlan` and release goes through the map, that would fail all subsequent native execution in the executor. - `MemoryPoolType` derives `Debug` so the "no pool was registered" warning can name the pool type. - `ExecutionContext::task_attempt_id` and `memory_pool_config` are removed: both existed only to drive the release call, and CI runs `clippy -- -D warnings`, which rejects the now-dead fields. ## 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. ``` test execution::memory_pools::task_shared::tests::plans_in_the_same_task_share_one_pool ... ok test execution::memory_pools::task_shared::tests::plans_in_different_tasks_get_different_pools ... ok test execution::memory_pools::task_shared::tests::pool_is_removed_only_once_the_last_plan_releases_it ... ok test execution::memory_pools::task_shared::tests::dropping_the_reference_releases_the_pool ... ok test execution::memory_pools::task_shared::tests::releasing_an_unregistered_pool_does_not_panic ... ok ``` Full native crate suite passes (140 passed, 4 ignored). `cargo clippy --all-targets -- -D warnings` is 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`. The guard makes the release unconditional at the type level, which is the property I would want a test to establish, but I want to be clear 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. -- 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]
