Skip to content

fix: attach tokio runtime threads to the JVM as daemon threads - #5748

Open
zhangfengcdt wants to merge 2 commits into
apache:mainfrom
zhangfengcdt:fix/daemon-runtime-threads
Open

fix: attach tokio runtime threads to the JVM as daemon threads#5748
zhangfengcdt wants to merge 2 commits into
apache:mainfrom
zhangfengcdt:fix/daemon-runtime-threads

Conversation

@zhangfengcdt

Copy link
Copy Markdown
Member

Which issue does this PR close?

Closes #5664.

Rationale for this change

#4734 releases the tokio runtime from the driver/executor plugin shutdown(), which fixes the
JVM exit hang whenever SparkContext.stop() is called. It does not help when an application
simply returns from main (or a PySpark driver exits) without calling spark.stop():

  1. Runtime threads that call back into the JVM (memory pool acquireMemory, scans fed from JVM
    iterators, scalar subqueries, ...) are attached lazily by jni-rs with AttachCurrentThread,
    which makes them non-daemon JVM threads.
  2. DestroyJavaVM waits until all non-daemon threads have exited before running shutdown hooks.
  3. The shutdown hook is what would call SparkContext.stop() -> plugin shutdown() ->
    release_runtime(), so the threads can only be released by a hook that can only run after the
    threads have exited. The JVM hangs forever.

Spark's own thread pools are daemon threads for exactly this reason.

What changes are included in this PR?

  • build_runtime registers tokio on_thread_start / on_thread_stop hooks that attach each
    runtime thread with AttachCurrentThreadAsDaemon and detach it before the thread exits. jni-rs
    reuses an existing attachment, so its lazy non-daemon attach is never triggered on these
    threads. jni-rs intentionally offers no daemon attach API, hence the raw invoke-interface
    calls. The hooks are no-ops when no JVM is registered (Rust unit tests and benchmarks).
  • The runtime is still released on plugin shutdown as before, so workers exit cleanly when
    SparkContext.stop() is called.
  • Updated the threading section of the contributor guide.

How are these changes tested?

New CometRuntimeShutdownSuite forks a child JVM that runs a native Comet query with the unified
memory pool (which calls into the JVM from a tokio worker) and returns from main without
SparkContext.stop(). The parent asserts the child exits with status 0 within a deadline. A
child JVM is required because the failure mode is "the JVM cannot exit", which cannot be observed
from inside that JVM.

Without the fix the child prints its marker and never exits (killed at the deadline); with the
fix it exits in a few seconds. Also verified CometPluginsSuite (the SparkContext.stop() path)
and the native planner unit tests, which build the runtime without a JVM.

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 apache#5664
@zhangfengcdt
zhangfengcdt force-pushed the fix/daemon-runtime-threads branch from 7c0e46d to 7f127af Compare September 6, 2026 18:41
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

JVM exit hangs when the application returns from main without SparkContext.stop()

1 participant