peterxcli commented on code in PR #5494:
URL: https://github.com/apache/datafusion-comet/pull/5494#discussion_r3873491669
##########
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:
Restored in 9ea52ffa6 as an `assert_ne!(exec_context, 0, "Comet execution
context shouldn't be null!")` before the `Box::from_raw`, with the same message
`get_execution_context` uses — `try_unwrap_or_throw` converts the panic into a
Java exception as before.
##########
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:
Done in 9ea52ffa6 — the tracing call and the `memInUse` check now run inside
`attempt` blocks and the throw moved to the very end of `close()`, so a
teardown failure still emits the non-zero-memory warning (and any diagnostic
failure is attached as suppressed).
##########
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:
Done in 9ea52ffa6 — the test now runs under
`spark.comet.metrics.updateInterval=0`, which disables the periodic updates
inside `executePlan`, so the releasePlan update is the only place the stub can
fire; and it asserts the thrown message contains `injected metrics update
failure` so an unrelated throw can't make it pass vacuously.
##########
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:
Done in 9ea52ffa6 — dropped the pass-through and the four arms call
`acquire_task_shared_pool` directly; `tracked` stays.
##########
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:
Added in 9ea52ffa6:
`concurrent_acquire_and_drop_leaves_a_consistent_registry` — 8 threads each
acquire-and-drop the same task attempt id 1000 times, then the test asserts the
registry entry is gone once they quiesce and that a fresh acquire comes back
registered. The hand-removal test stays as the deterministic unit-level check
of the `ptr_eq` guard.
--
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]
