fix: attach tokio runtime threads to the JVM as daemon threads

.github/workflows/pr_build_linux.yml

@@ -409,6 +409,7 @@ jobs:
               org.apache.comet.CometSetOpWithGroupBySuite
               org.apache.comet.CometSparkSessionExtensionsSuite
               org.apache.spark.CometPluginsSuite
+              org.apache.spark.CometRuntimeShutdownSuite
               org.apache.spark.CometTaskMemoryManagerSuite
               org.apache.spark.CometExecIteratorLifecycleSuite
               org.apache.spark.CometPluginsDefaultSuite

.github/workflows/pr_build_macos.yml

@@ -182,6 +182,7 @@ jobs:
               org.apache.comet.CometSetOpWithGroupBySuite
               org.apache.comet.CometSparkSessionExtensionsSuite
               org.apache.spark.CometPluginsSuite
+              org.apache.spark.CometRuntimeShutdownSuite
               org.apache.spark.CometTaskMemoryManagerSuite
               org.apache.spark.CometExecIteratorLifecycleSuite
               org.apache.spark.CometPluginsDefaultSuite

docs/source/contributor-guide/development.md

@@ -59,9 +59,11 @@ in the operator struct or be shared via `Arc`.
 **JNI calls work from any thread, but have overhead.** `JVMClasses::get_env()` calls
 `AttachCurrentThread`, which acquires JVM internal locks. The attachment is cached in
 thread-local storage and is only released when the worker thread itself exits, not when the
-`AttachGuard` is dropped. This is why the tokio runtime has to be shut down (releasing its
-worker threads) for the JVM to be able to exit. Acquiring the JVM locks on each `get_env()`
-call adds overhead, so avoid calling into the JVM on hot paths during stream execution.
+`AttachGuard` is dropped. Tokio runtime threads are attached as JVM daemon threads when they
+start (see `build_runtime`), so they do not prevent the JVM from exiting when an application
+returns from `main` without calling `SparkContext.stop()`. Acquiring the JVM locks on each
+`get_env()` call adds overhead, so avoid calling into the JVM on hot paths during stream
+execution.
 
 **Do not call `TaskContext.get()` from JVM callbacks during execution.** Spark's `TaskContext` is
 a `ThreadLocal` on the executor task thread. JVM methods invoked from tokio worker threads will

native/core/src/execution/jni_api.rs

@@ -116,7 +116,7 @@ use crate::execution::spark_config::{
 };
 use crate::parquet::encryption_support::{CometEncryptionFactory, ENCRYPTION_FACTORY_ID};
 use datafusion_comet_proto::spark_operator::operator::OpStruct;
-use log::info;
+use log::{info, warn};
 use std::sync::OnceLock;
 #[cfg(feature = "jemalloc")]
 use tikv_jemalloc_ctl::{epoch, stats};
@@ -238,10 +238,49 @@ fn build_runtime(default_worker_threads: Option<usize>) -> Runtime {
     }
     builder
         .enable_all()
+        .on_thread_start(attach_thread_as_daemon)
+        .on_thread_stop(detach_thread)
         .build()
         .expect("Failed to create Tokio runtime")
 }
 
+/// Attaches a runtime thread to the JVM as a daemon thread.
+///
+/// jni-rs attaches threads lazily with `AttachCurrentThread`, which makes them non-daemon JVM
+/// threads. `DestroyJavaVM` waits for all non-daemon threads to exit before it runs shutdown
+/// hooks, but runtime threads only exit once the shutdown hook has called `SparkContext.stop()`
+/// and [`release_runtime`]. An application that returns from `main` without calling
+/// `SparkContext.stop()` would therefore never exit. Daemon threads are not waited for, and
+/// jni-rs reuses an existing attachment rather than attaching again.
+fn attach_thread_as_daemon() {
+    let Some(vm) = crate::JAVA_VM.get() else {
+        return;
+    };
+    let vm = vm.get_raw();
+    let mut env: *mut std::ffi::c_void = std::ptr::null_mut();
+    // SAFETY: `vm` is the JavaVM stored by `NativeBase.init` and outlives the runtime. Null
+    // attach args select the default JNI version, thread name and thread group.
+    let rc =
+        unsafe { ((**vm).v1_4.AttachCurrentThreadAsDaemon)(vm, &mut env, std::ptr::null_mut()) };
+    if rc != jni::sys::JNI_OK {
+        warn!("Failed to attach tokio runtime thread to the JVM as a daemon thread: {rc}");
+    }
+}
+
+/// Detaches a thread attached by [`attach_thread_as_daemon`] before it exits. jni-rs only
+/// detaches threads it attached itself.
+fn detach_thread() {
+    let Some(vm) = crate::JAVA_VM.get() else {
+        return;
+    };
+    let vm = vm.get_raw();
+    // SAFETY: see `attach_thread_as_daemon`. Detaching an unattached thread is a JNI error,
+    // not undefined behavior.
+    unsafe {
+        ((**vm).v1_1.DetachCurrentThread)(vm);
+    }
+}
+
 /// Initialize the global Tokio runtime with the given default worker thread count.
 /// If the runtime is already initialized, this is a no-op.
 pub fn init_runtime(default_worker_threads: usize) {

spark/src/test/scala/org/apache/spark/CometRuntimeShutdownSuite.scala

@@ -0,0 +1,113 @@
+/*
+ * 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.File
+import java.lang.management.ManagementFactory
+import java.nio.charset.StandardCharsets
+import java.nio.file.{Files, Paths}
+import java.util.concurrent.TimeUnit
+
+import scala.jdk.CollectionConverters._
+
+import org.scalatest.funsuite.AnyFunSuite
+
+import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.functions.{expr, sum}
+import org.apache.spark.util.Utils
+
+/**
+ * Entry point run in a child JVM by [[CometRuntimeShutdownSuite]]. Executes a native Comet query
+ * and then returns from `main` without calling `SparkContext.stop()`.
+ */
+object CometRuntimeShutdownMain {
+  val Marker = "comet native query executed"
+
+  def main(args: Array[String]): Unit = {
+    val spark = SparkSession
+      .builder()
+      .master("local[2]")
+      .config("spark.plugins", "org.apache.spark.CometPlugin")
+      .config(
+        "spark.shuffle.manager",
+        "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager")
+      .config("spark.sql.extensions", "org.apache.comet.CometSparkSessionExtensions")
+      .config("spark.comet.enabled", "true")
+      .config("spark.comet.exec.enabled", "true")
+      .config("spark.memory.offHeap.enabled", "true")
+      .config("spark.memory.offHeap.size", "256m")
+      .config("spark.ui.enabled", "false")
+      .getOrCreate()
+    val dir = Files.createTempDirectory("comet-runtime-shutdown").toFile
+    try {
+      val path = dir.getAbsolutePath
+      spark.range(100000).write.mode("overwrite").parquet(path)
+      // The aggregate reserves memory through the unified pool, which calls back into the JVM
+      // from a tokio worker thread.
+      val df = spark.read.parquet(path).groupBy(expr("id % 10")).agg(sum("id"))
+      df.collect()
+      val plan = df.queryExecution.executedPlan.toString
+      require(plan.contains("CometHashAggregate"), s"Expected a native plan but got:\n$plan")
+    } finally {
+      Utils.deleteRecursively(dir)
+    }
+    // scalastyle:off println
+    println(Marker)
+    // scalastyle:on println
+  }
+}
+
+/**
+ * Runs Comet in a child JVM because the failure mode under test, the JVM being unable to exit,
+ * cannot be observed from inside the affected JVM.
+ */
+class CometRuntimeShutdownSuite extends AnyFunSuite {
+
+  test("JVM exits when main returns without SparkContext.stop()") {
+    val javaBin = new File(System.getProperty("java.home"), "bin/java").getAbsolutePath
+    val jvmArgs = ManagementFactory.getRuntimeMXBean.getInputArguments.asScala
+      .filter(a => a.startsWith("--add-opens") || a.startsWith("--add-exports"))
+    val cmd = Seq(javaBin, "-Xmx1g") ++ jvmArgs ++ Seq(
+      "-cp",
+      System.getProperty("java.class.path"),
+      CometRuntimeShutdownMain.getClass.getName.stripSuffix("$"))
+
+    // java.io.tmpdir is pinned to target/tmp by the pom and may not exist yet
+    val tmpDir = Files.createDirectories(Paths.get(System.getProperty("java.io.tmpdir")))
+    val output = Files.createTempFile(tmpDir, "comet-runtime-shutdown", ".log").toFile
+    try {
+      val process =
+        new ProcessBuilder(cmd.asJava).redirectErrorStream(true).redirectOutput(output).start()
+      val exited = process.waitFor(2, TimeUnit.MINUTES)
+      if (!exited) {
+        process.destroyForcibly().waitFor()
+      }
+      val log = new String(Files.readAllBytes(output.toPath), StandardCharsets.UTF_8)
+      val tail = log.linesIterator.toSeq.takeRight(30).mkString("\n")
+      assert(
+        log.contains(CometRuntimeShutdownMain.Marker),
+        s"Child JVM did not run query:\n$tail")
+      assert(exited, s"Child JVM did not exit after main returned:\n$tail")
+      assert(process.exitValue() == 0, s"Child JVM exited with ${process.exitValue()}:\n$tail")
+    } finally {
+      output.delete()
+    }
+  }
+}