zhangfengcdt opened a new pull request, #5748:
URL: https://github.com/apache/datafusion-comet/pull/5748
## 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.
--
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]
