From a8656af287952234b880e128fa9368e744e9086a Mon Sep 17 00:00:00 2001 From: zhangfengcdt Date: Sun, 6 Sep 2026 11:22:56 -0700 Subject: [PATCH 1/2] fix: attach tokio runtime threads to the JVM as daemon threads jni-rs attaches runtime threads lazily as non-daemon JVM threads, and DestroyJavaVM waits for those before running the shutdown hook that would release the runtime. An application that returns from main without SparkContext.stop() therefore never exits. Attach each runtime thread with AttachCurrentThreadAsDaemon on start and detach it on stop, so the shutdown hook can run and release the runtime. Add a forked-JVM regression test and update the threading section of the contributor guide. Closes #5664 --- docs/source/contributor-guide/development.md | 8 +- native/core/src/execution/jni_api.rs | 41 ++++++- .../spark/CometRuntimeShutdownSuite.scala | 113 ++++++++++++++++++ 3 files changed, 158 insertions(+), 4 deletions(-) create mode 100644 spark/src/test/scala/org/apache/spark/CometRuntimeShutdownSuite.scala diff --git a/docs/source/contributor-guide/development.md b/docs/source/contributor-guide/development.md index 26fdbbb4043..04952d5da31 100644 --- a/docs/source/contributor-guide/development.md +++ b/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 diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 84abd98f2b6..1c470e39296 100644 --- a/native/core/src/execution/jni_api.rs +++ b/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) -> 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) { diff --git a/spark/src/test/scala/org/apache/spark/CometRuntimeShutdownSuite.scala b/spark/src/test/scala/org/apache/spark/CometRuntimeShutdownSuite.scala new file mode 100644 index 00000000000..69425c7eae4 --- /dev/null +++ b/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() + } + } +} From 7f127afc2b7d94bc9a4425135b0f7f7650009aa2 Mon Sep 17 00:00:00 2001 From: zhangfengcdt Date: Sun, 6 Sep 2026 11:41:17 -0700 Subject: [PATCH 2/2] ci: register CometRuntimeShutdownSuite in PR workflows --- .github/workflows/pr_build_linux.yml | 1 + .github/workflows/pr_build_macos.yml | 1 + 2 files changed, 2 insertions(+) diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 79fea7fb881..c62efa70c9b 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.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 diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index 5473d3988a8..b40de4b5e68 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.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