andygrove commented on code in PR #5494:
URL: https://github.com/apache/datafusion-comet/pull/5494#discussion_r3873207007
##########
native/core/src/execution/jni_api.rs:
##########
@@ -994,31 +986,23 @@ pub extern "system" fn
Java_org_apache_comet_Native_releasePlan(
exec_context: jlong,
) {
try_unwrap_or_throw(&e, |env| unsafe {
- let execution_context = get_execution_context(exec_context);
-
- // Move the guard out before the fallible metrics update. On error it
unregisters while
- // leaving the raw execution context alive for the JVM's existing
release retry.
- let memory_pool_registration =
execution_context.memory_pool_registration.take();
-
- // Update metrics
- update_metrics(env, execution_context)?;
-
- handle_task_shared_pool_release(
- execution_context.memory_pool_config.pool_type,
- execution_context.task_attempt_id,
- );
+ // Reclaim ownership of the context up front so that it is always
freed, even if updating
+ // metrics below fails. Dropping it releases the memory pool and every
JNI global ref the
+ // context holds.
+ let mut execution_context: Box<ExecutionContext> =
+ Box::from_raw(exec_context as *mut ExecutionContext);
Review Comment:
Reclaiming the `Box` up front is the right call, but it drops the null check
that `get_execution_context` was doing with its `.expect("Comet execution
context shouldn't be null!")`. A null or stale `jlong` coming from the JVM is
now undefined behaviour rather than a panic that `try_unwrap_or_throw` converts
into a Java exception.
I do not think it is reachable today, since `plan` is a `val` that only a
successful `createPlan` can set. But this is a JNI entry point taking a raw
pointer from Java, and the PR itself points out that a second `releasePlan` on
the same pointer is now a use-after-free rather than a leak. Could you keep an
assertion on the pointer before reclaiming it?
##########
spark/src/main/scala/org/apache/comet/CometExecIterator.scala:
##########
@@ -227,13 +227,39 @@ class CometExecIterator(
def close(): Unit = synchronized {
if (!closed) {
- if (currentBatch != null) {
- currentBatch.close()
- currentBatch = null
+ closed = true
+
+ // Attempt every resource's cleanup independently, so that one failure
does not skip the
+ // remaining resources: this close() is the only chance to release them,
since `closed` is
+ // already set and the task-completion retry is a no-op. The first
failure is rethrown with
+ // any later ones attached as suppressed exceptions.
+ var failure: Throwable = null
+ def attempt(cleanup: => Unit): Unit = {
+ try {
+ cleanup
+ } catch {
+ case t: Throwable =>
+ if (failure == null) failure = t else failure.addSuppressed(t)
+ }
+ }
+
+ attempt {
+ if (currentBatch != null) {
+ currentBatch.close()
+ currentBatch = null
+ }
+ }
+ attempt(nativeUtil.close())
+ shuffleBlockIterators.values.foreach(it => attempt(it.close()))
+
+ // Released last and exactly once, even if the teardown above failed:
dropping the native
+ // execution context frees this plan's task-shared memory pool reference
and several JNI
+ // global refs.
+ attempt(nativeLib.releasePlan(plan))
+
+ if (failure != null) {
Review Comment:
Throwing at this point means a teardown failure also loses the `memInUse !=
0` warning further down, and that is exactly the diagnostic that would tell you
whether the failed teardown stranded native memory.
Could the tracing call and the `memInUse` check go inside `attempt` blocks
as well, with the throw moved to the very end of the method? Position 8 stays
fixed either way, because `closed` is already set by then, so a throw out of
`traceMemoryUsage()` can no longer trigger a listener retry.
##########
native/core/src/execution/memory_pools/mod.rs:
##########
@@ -34,94 +34,67 @@ use unified_pool::CometUnifiedMemoryPool;
pub(crate) use config::*;
pub(crate) use task_shared::*;
+/// Creates the memory pool for a native plan.
+///
+/// Task-shared pools use their returned `Arc` as the RAII handle, so they
remain registered for as
+/// long as the plan or any of its reservations retain the pool.
pub(crate) fn create_memory_pool(
memory_pool_config: &MemoryPoolConfig,
comet_task_memory_manager: Arc<Global<JObject<'static>>>,
task_attempt_id: i64,
) -> Arc<dyn MemoryPool> {
const NUM_TRACKED_CONSUMERS: usize = 10;
- match memory_pool_config.pool_type {
- MemoryPoolType::GreedyUnified => {
- let mut memory_pool_map = TASK_SHARED_MEMORY_POOLS.lock().unwrap();
- let per_task_memory_pool =
- memory_pool_map.entry(task_attempt_id).or_insert_with(|| {
- let pool: Arc<dyn MemoryPool> =
Arc::new(TrackConsumersPool::new(
- CometUnifiedMemoryPool::new(
- Arc::clone(&comet_task_memory_manager),
- task_attempt_id,
- ),
- NonZeroUsize::new(NUM_TRACKED_CONSUMERS).unwrap(),
- ));
- PerTaskMemoryPool::new(pool)
- });
- per_task_memory_pool.num_plans += 1;
- Arc::clone(&per_task_memory_pool.memory_pool)
- }
- MemoryPoolType::FairUnified => {
- let mut memory_pool_map = TASK_SHARED_MEMORY_POOLS.lock().unwrap();
- let per_task_memory_pool =
- memory_pool_map.entry(task_attempt_id).or_insert_with(|| {
- let pool: Arc<dyn MemoryPool> =
Arc::new(TrackConsumersPool::new(
- CometFairMemoryPool::new(
- Arc::clone(&comet_task_memory_manager),
- memory_pool_config.pool_size,
- ),
- NonZeroUsize::new(NUM_TRACKED_CONSUMERS).unwrap(),
- ));
- PerTaskMemoryPool::new(pool)
- });
- per_task_memory_pool.num_plans += 1;
- Arc::clone(&per_task_memory_pool.memory_pool)
- }
- MemoryPoolType::Greedy => Arc::new(TrackConsumersPool::new(
- GreedyMemoryPool::new(memory_pool_config.pool_size),
- NonZeroUsize::new(NUM_TRACKED_CONSUMERS).unwrap(),
- )),
- MemoryPoolType::FairSpill => Arc::new(TrackConsumersPool::new(
- FairSpillPool::new(memory_pool_config.pool_size),
+
+ fn tracked(pool: impl MemoryPool + 'static) -> Arc<dyn MemoryPool> {
+ Arc::new(TrackConsumersPool::new(
+ pool,
NonZeroUsize::new(NUM_TRACKED_CONSUMERS).unwrap(),
- )),
+ ))
+ }
+
+ fn task_shared(
Review Comment:
This forwards its arguments unchanged, so it is only a second name for
`acquire_task_shared_pool`. Calling that directly in the four arms below would
make it obvious at a glance which arms touch the registry. `tracked` is pulling
real weight, so I would keep that one.
##########
native/core/src/execution/memory_pools/task_shared.rs:
##########
@@ -15,46 +15,174 @@
// specific language governing permissions and limitations
// under the License.
-use crate::execution::memory_pools::MemoryPoolType;
-use datafusion::execution::memory_pool::MemoryPool;
+use datafusion::execution::memory_pool::{
+ MemoryConsumer, MemoryLimit, MemoryPool, MemoryReservation,
+};
use once_cell::sync::Lazy;
+use parking_lot::Mutex;
+use std::collections::hash_map::Entry;
use std::collections::HashMap;
-use std::sync::{Arc, Mutex};
+use std::fmt;
+use std::sync::{Arc, Weak};
-/// The per-task memory pools keyed by task attempt id.
-pub(crate) static TASK_SHARED_MEMORY_POOLS: Lazy<Mutex<HashMap<i64,
PerTaskMemoryPool>>> =
+/// The memory pools for active task attempts. Weak references let the pool's
normal `Arc`
+/// ownership determine its lifetime, and each pool removes its entry when the
last reference drops.
+static TASK_SHARED_MEMORY_POOLS: Lazy<Mutex<HashMap<i64,
Weak<TaskSharedMemoryPool>>>> =
Lazy::new(|| Mutex::new(HashMap::new()));
-pub(crate) struct PerTaskMemoryPool {
- pub(crate) memory_pool: Arc<dyn MemoryPool>,
- pub(crate) num_plans: usize,
+/// A transparent `MemoryPool` wrapper whose lifetime also controls its
registry entry.
+#[derive(Debug)]
+struct TaskSharedMemoryPool {
+ task_attempt_id: i64,
+ inner: Arc<dyn MemoryPool>,
}
-impl PerTaskMemoryPool {
- pub(crate) fn new(memory_pool: Arc<dyn MemoryPool>) -> Self {
- Self {
- memory_pool,
- num_plans: 0,
+impl fmt::Display for TaskSharedMemoryPool {
+ fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+ fmt::Display::fmt(self.inner.as_ref(), f)
+ }
+}
+
+impl MemoryPool for TaskSharedMemoryPool {
+ fn name(&self) -> &str {
+ self.inner.name()
+ }
+
+ fn register(&self, consumer: &MemoryConsumer) {
+ self.inner.register(consumer)
+ }
+
+ fn unregister(&self, consumer: &MemoryConsumer) {
+ self.inner.unregister(consumer)
+ }
+
+ fn grow(&self, reservation: &MemoryReservation, additional: usize) {
+ self.inner.grow(reservation, additional)
+ }
+
+ fn shrink(&self, reservation: &MemoryReservation, shrink: usize) {
+ self.inner.shrink(reservation, shrink)
+ }
+
+ fn try_grow(
+ &self,
+ reservation: &MemoryReservation,
+ additional: usize,
+ ) -> datafusion::common::Result<()> {
+ self.inner.try_grow(reservation, additional)
+ }
+
+ fn reserved(&self) -> usize {
+ self.inner.reserved()
+ }
+
+ fn memory_limit(&self) -> MemoryLimit {
+ self.inner.memory_limit()
+ }
+}
+
+impl Drop for TaskSharedMemoryPool {
+ fn drop(&mut self) {
+ if let Entry::Occupied(entry) =
TASK_SHARED_MEMORY_POOLS.lock().entry(self.task_attempt_id)
+ {
+ // An acquire racing with this drop can replace our expired `Weak`
before we obtain the
+ // lock. Do not let the old pool remove that replacement's entry.
+ if std::ptr::eq(entry.get().as_ptr(), self) {
+ entry.remove();
+ }
}
}
}
-// This function reduces the refcount of a per-task memory pool when a native
plan is released.
-// If the refcount reaches zero, the memory pool is removed from the map and
dropped.
-pub(crate) fn handle_task_shared_pool_release(pool_type: MemoryPoolType,
task_attempt_id: i64) {
- if !pool_type.is_task_shared() {
- return;
- }
-
- // Decrement the number of native plans using the per-task shared memory
pool, and
- // remove the memory pool if the released native plan is the last native
plan using it.
- let mut memory_pool_map = TASK_SHARED_MEMORY_POOLS.lock().unwrap();
- if let Some(per_task_memory_pool) =
memory_pool_map.get_mut(&task_attempt_id) {
- per_task_memory_pool.num_plans -= 1;
- if per_task_memory_pool.num_plans == 0 {
- // Drop the memory pool from the per-task memory pool map if there
are no
- // more native plans using it.
- memory_pool_map.remove(&task_attempt_id);
+/// Returns the memory pool shared by every native plan in `task_attempt_id`,
creating it with
+/// `create` if no live pool exists for the task. The returned `Arc` is the
RAII handle: the pool
+/// stays registered until the last reference to it drops.
+pub(crate) fn acquire_task_shared_pool(
+ task_attempt_id: i64,
+ create: impl FnOnce() -> Arc<dyn MemoryPool>,
+) -> Arc<dyn MemoryPool> {
+ let mut memory_pool_map = TASK_SHARED_MEMORY_POOLS.lock();
+ if let Some(memory_pool) = memory_pool_map
+ .get(&task_attempt_id)
+ .and_then(Weak::upgrade)
+ {
+ return memory_pool;
+ }
+
+ let memory_pool = Arc::new(TaskSharedMemoryPool {
+ task_attempt_id,
+ inner: create(),
+ });
+ memory_pool_map.insert(task_attempt_id, Arc::downgrade(&memory_pool));
+ memory_pool
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use datafusion::execution::memory_pool::UnboundedMemoryPool;
+
+ /// Tests share the process-wide pool map, so each uses its own task
attempt id.
+ fn acquire(task_attempt_id: i64) -> Arc<dyn MemoryPool> {
+ acquire_task_shared_pool(task_attempt_id, ||
Arc::new(UnboundedMemoryPool::default()))
+ }
+
+ fn is_registered(task_attempt_id: i64) -> bool {
+ TASK_SHARED_MEMORY_POOLS
+ .lock()
+ .contains_key(&task_attempt_id)
+ }
+
+ #[test]
+ fn plans_in_the_same_task_share_one_pool() {
+ let first = acquire(-1001);
+ let second = acquire(-1001);
+ assert!(Arc::ptr_eq(&first, &second));
+ }
+
+ #[test]
+ fn plans_in_different_tasks_get_different_pools() {
+ let first = acquire(-1002);
+ let second = acquire(-1003);
+ assert!(!Arc::ptr_eq(&first, &second));
+ }
+
+ #[test]
+ fn pool_is_removed_only_after_the_last_reference_drops() {
+ let first = acquire(-1004);
+ let second = acquire(-1004);
+
+ drop(first);
+ assert!(
+ is_registered(-1004),
+ "pool must outlive the first reference to release it"
+ );
+
+ drop(second);
+ assert!(!is_registered(-1004));
+ }
+
+ #[test]
+ fn dropping_the_reference_releases_the_pool() {
+ // Stands in for `createPlan` failing after the pool was acquired. The
ordinary `Arc` drops
+ // on unwind, so no explicit release path is needed.
+ {
+ let _pool = acquire(-1005);
+ assert!(is_registered(-1005));
}
+ assert!(!is_registered(-1005));
+ }
+
+ #[test]
+ fn an_old_pool_does_not_remove_its_replacement() {
+ let old_pool = acquire(-1006);
+ TASK_SHARED_MEMORY_POOLS.lock().remove(&-1006);
Review Comment:
Removing the entry by hand does exercise the `ptr_eq` guard, but it does not
show that the real code can reach that state. If the guard ever regresses the
failure is silent and unpleasant: the old pool's `Drop` evicts a live
replacement, so later plans in the same task each build their own pool and the
per-task limit ends up enforced more than once.
Could you add a short concurrent churn test alongside this one? Something
like a handful of threads each acquiring and immediately dropping the same task
attempt id in a loop, then asserting the map is empty once they quiesce and
that a fresh acquire comes back registered.
##########
spark/src/test/scala/org/apache/spark/CometExecIteratorLifecycleSuite.scala:
##########
@@ -0,0 +1,213 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.spark
+
+import java.io.ByteArrayInputStream
+import java.lang.ref.WeakReference
+import java.util.Properties
+import java.util.concurrent.atomic.AtomicBoolean
+
+import org.apache.spark.executor.TaskMetrics
+import org.apache.spark.memory.{TaskMemoryManager, TestMemoryManager}
+import org.apache.spark.sql.CometTestBase
+import org.apache.spark.sql.catalyst.expressions.PrettyAttribute
+import org.apache.spark.sql.comet.{CometExec, CometExecUtils, CometMetricNode}
+import org.apache.spark.sql.comet.execution.arrow.CometArrowStream
+import org.apache.spark.sql.types.{LongType, StructField, StructType}
+
+import org.apache.comet.{CometExecIterator, CometShuffleBlockIterator, Native}
+import org.apache.comet.serde.Config.ConfigMap
+import org.apache.comet.serde.OperatorOuterClass
+
+/**
+ * Regression tests for the native plan lifecycle: every `createPlan` must be
balanced by exactly
+ * one release of the native execution context and of its task-shared memory
pool reference, even
+ * when a step of the lifecycle fails partway through. See issue #5212
(positions 2, 3 and 8).
+ */
+class CometExecIteratorLifecycleSuite extends CometTestBase {
+
+ private def withTaskContext[T](taskAttemptId: Long)(f: => T): T = {
+ val memoryManager = new TestMemoryManager(new SparkConf())
+ val taskMemoryManager = new TaskMemoryManager(memoryManager, taskAttemptId)
+ val taskContext = new TaskContextImpl(
+ stageId = 0,
+ stageAttemptNumber = 0,
+ partitionId = 0,
+ numPartitions = 1,
+ taskAttemptId = taskAttemptId,
+ attemptNumber = 0,
+ taskMemoryManager = taskMemoryManager,
+ localProperties = new Properties,
+ metricsSystem = null,
+ taskMetrics = TaskMetrics.empty,
+ cpus = 1,
+ resources = Map.empty)
+ TaskContext.setTaskContext(taskContext)
+ try {
+ f
+ } finally {
+ taskMemoryManager.cleanUpAllAllocatedMemory()
+ TaskContext.unset()
+ }
+ }
+
+ /** Retries GC until every weak reference clears or the deadline passes;
returns survivors. */
+ private def survivorsAfterGc(refs: Seq[WeakReference[_]]): Int = {
+ val deadline = System.nanoTime() + 30L * 1000 * 1000 * 1000
+ while (refs.exists(_.get() != null) && System.nanoTime() < deadline) {
+ System.gc()
+ Thread.sleep(50)
+ }
+ refs.count(_.get() != null)
+ }
+
+ test("createPlan failure releases the task-shared memory pool reference") {
+ val nativeLib = new Native()
+ val emptyPlan =
OperatorOuterClass.Operator.newBuilder().build().toByteArray
+ // An unknown DataFusion config makes createPlan fail while building the
session context,
+ // which happens after the task-shared memory pool has been registered for
the task.
+ val badConfigs = ConfigMap
+ .newBuilder()
+ .putEntries("spark.comet.datafusion.no_such_namespace.option", "1")
+ .build()
+ .toByteArray
+
+ val managerRefs = (0 until 10).map { i =>
+ // Unique synthetic task attempt ids keep each iteration's pool entry
independent.
+ val taskAttemptId = 4200000L + i
+ withTaskContext(taskAttemptId) {
+ val manager = new CometTaskMemoryManager(i, taskAttemptId)
+ val thrown = intercept[Throwable] {
+ nativeLib.createPlan(
+ i,
+ Array.empty[Object],
+ emptyPlan,
+ badConfigs,
+ 1,
+ CometMetricNode(Map.empty),
+ 0L,
+ manager,
+ Array(System.getProperty("java.io.tmpdir")),
+ 8192,
+ true,
+ "fair_unified",
+ 64L << 20,
+ 64L << 20,
+ taskAttemptId,
+ 1L,
+ null,
+ null,
+ null)
+ }
+ // Guard against a vacuous pass: the failure must be the injected
config error thrown
+ // inside createPlan, not e.g. an UnsatisfiedLinkError from a missing
native library.
+ assert(
+ thrown.getMessage != null &&
thrown.getMessage.contains("no_such_namespace"),
+ s"expected the injected DataFusion config failure, got: $thrown")
+ new WeakReference(manager)
+ }
+ }
+
+ // A stranded TASK_SHARED_MEMORY_POOLS entry holds a JNI global ref to the
+ // CometTaskMemoryManager, so the manager staying reachable means the pool
leaked.
+ val survivors = survivorsAfterGc(managerRefs)
+ assert(
+ survivors == 0,
+ s"$survivors of ${managerRefs.size} CometTaskMemoryManagers stayed
reachable: " +
+ "createPlan failure leaked their task-shared memory pool references")
+ }
+
+ test("close() is idempotent and still releases the plan when teardown
throws") {
+ withTaskContext(4300000L) {
+ val boom = new java.io.IOException("injected shuffle block close
failure")
+ val throwingBlockIter =
+ new CometShuffleBlockIterator(new
ByteArrayInputStream(Array.emptyByteArray)) {
+ override def close(): Unit = throw boom
+ }
+ @volatile var laterInputClosed = false
+ val trackingBlockIter =
+ new CometShuffleBlockIterator(new
ByteArrayInputStream(Array.emptyByteArray)) {
+ override def close(): Unit = {
+ laterInputClosed = true
+ super.close()
+ }
+ }
+ val limitOp =
+ CometExecUtils.getLimitNativePlan(Seq(PrettyAttribute("test",
LongType)), 100).get
+ val iter = new CometExecIterator(
+ id = 1L,
+ inputObjects = Array.empty[Object],
+ numOutputCols = 1,
+ protobufQueryPlan = limitOp.toByteArray,
+ nativeMetrics = CometMetricNode(Map.empty),
+ numParts = 1,
+ partitionIndex = 0,
+ shuffleBlockIterators = Map(0 -> throwingBlockIter, 1 ->
trackingBlockIter))
+
+ val thrown = intercept[java.io.IOException](iter.close())
+ assert(thrown eq boom)
+ // One input's close failure must not skip the remaining resources: this
close() is the only
+ // chance to release them, since the task-completion retry is a no-op
once `closed` is set.
+ assert(laterInputClosed, "a later shuffle input was not closed after an
earlier one threw")
+ // The first close() must have marked the iterator closed and released
the plan despite the
+ // teardown failure: a second close() re-running releasePlan would free
the native
+ // execution context twice, and skipping the release would strand it.
+ iter.close()
+ }
+ }
+
+ test("releasePlan frees the native context even when the final metrics
update fails") {
+ withTaskContext(4400000L) {
+ val failMetrics = new AtomicBoolean(false)
+ class ThrowingMetricNode extends CometMetricNode(Map.empty, Nil) {
+ override def set_all_from_bytes(bytes: Array[Byte]): Unit = {
+ if (failMetrics.get()) {
+ throw new IllegalStateException("injected metrics update failure")
+ }
+ }
+ }
+ val schema = StructType(Seq(StructField("test", LongType, nullable =
false)))
+ val stream = CometArrowStream.fromColumnarBatchIter(
+ Iterator.empty,
+ schema,
+ CometArrowStream.NATIVE_TIMEZONE,
+ "lifecycle-test")
+ val limitOp =
+ CometExecUtils.getLimitNativePlan(Seq(PrettyAttribute("test",
LongType)), 100).get
+ val iter = CometExec.getCometIterator(
+ Array(stream.asInstanceOf[Object]),
+ 1,
+ limitOp,
+ new ThrowingMetricNode,
+ 1,
+ 0,
+ None,
+ Seq.empty)
+
+ failMetrics.set(true)
+ // Exhausting the iterator closes it, and the close propagates the
metrics failure thrown
+ // by the native releasePlan call.
+ intercept[Throwable](iter.hasNext)
Review Comment:
The `createPlan` test is careful to assert the injected message so an
unrelated failure cannot make it pass vacuously, but this one accepts any
`Throwable`. That matters more here than usual, because `executePlan` also
calls `update_metrics` on the periodic interval, so the throw could come from
mid-execution rather than from `releasePlan` and the test would still pass
without covering the path it is named after.
Could you assert the message contains `injected metrics update failure`, and
ideally arm `failMetrics` only once the plan has produced its final batch, so
the failure is pinned to the release path?
--
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]
