From 2dec579641faa7d2c72c5b8447f09a727fe5644f Mon Sep 17 00:00:00 2001 From: ChilePiquin Date: Mon, 14 Sep 2026 15:38:45 -0700 Subject: [PATCH] fix(python): shut down the shared Tokio runtime on interpreter exit The Python bindings route every async call through a single, lazily created Tokio runtime, stored behind a raw pointer so it can survive being reset (not dropped) across fork() -- dropping it there would hang joining now-dead worker threads. That leak-on-fork path is deliberate and documented, but normal process exit had no equivalent handling at all: nothing ever tells the runtime to shut down, so its worker threads keep running, completely uncoordinated with the interpreter, until the process ends. If a worker thread is mid-task exactly as Py_Finalize begins tearing down interpreter state, it can panic on state that's already gone. That panic happens on a background thread with no PyO3-wrapped call frame to catch it -- and if it lands while another unwind is already in progress, Rust aborts the whole process outright rather than merely failing that one call. This reproduces as an intermittent SIGABRT on otherwise-successful short-lived processes that open a connection, do a handful of operations, and exit normally. Give the runtime a coordinated, bounded shutdown at the one point it was missing one: register a Python atexit callback that calls runtime::shutdown(), which takes the runtime through the same atomic slot the fork handler uses and calls shutdown_timeout() rather than a bare drop -- bounded, so a wedged background task can never hang process exit, but real in-flight work gets a genuine chance to finish first. Runs with the GIL released (Python::detach), since the wait would otherwise be able to deadlock against any in-flight task that itself needs the GIL to complete. Co-Authored-By: Claude Sonnet 5 --- python/src/lib.rs | 26 +++++++++++++-- python/src/runtime.rs | 78 +++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 101 insertions(+), 3 deletions(-) diff --git a/python/src/lib.rs b/python/src/lib.rs index 06b31d033..2b3868f8a 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -8,8 +8,8 @@ use expr::{PyExpr, expr_col, expr_func, expr_lit}; use index::IndexConfig; use permutation::{PyAsyncPermutationBuilder, PyPermutationReader}; use pyo3::{ - Bound, PyResult, Python, pymodule, - types::{PyModule, PyModuleMethods}, + Bound, PyResult, Python, pyfunction, pymodule, + types::{PyAnyMethods, PyModule, PyModuleMethods}, wrap_pyfunction, }; use query::{FTSQuery, HybridQuery, Query, VectorQuery}; @@ -38,8 +38,24 @@ pub mod sql; pub mod table; pub mod util; +/// Shut down the shared Tokio runtime (see `runtime::shutdown`). +/// +/// Registered below as a Python `atexit` callback rather than called +/// directly: `atexit` runs while the interpreter is still fully valid, +/// which is the coordinated, bounded exit the runtime otherwise never gets. +/// +/// Runs the actual wait with the GIL released (`Python::detach`): shutdown +/// blocks the calling thread waiting on the runtime's own worker threads, +/// and if any in-flight task needs the GIL to finish (e.g. one that calls +/// back into Python), holding it here while waiting on that same task would +/// deadlock rather than time out. +#[pyfunction] +fn shutdown_runtime(py: Python<'_>) { + py.detach(|| runtime::shutdown(std::time::Duration::from_secs(5))); +} + #[pymodule] -pub fn _lancedb(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> { +pub fn _lancedb(py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> { let env = Env::new() .filter_or("LANCEDB_LOG", "warn") .write_style("LANCEDB_LOG_STYLE"); @@ -95,5 +111,9 @@ pub fn _lancedb(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(expr_lit, m)?)?; m.add_function(wrap_pyfunction!(expr_func, m)?)?; m.add("__version__", env!("CARGO_PKG_VERSION"))?; + // Give the shared runtime a coordinated, bounded shutdown at normal + // process exit -- see `shutdown_runtime` and `runtime::shutdown` for why. + py.import("atexit")? + .call_method1("register", (wrap_pyfunction!(shutdown_runtime, m)?,))?; Ok(()) } diff --git a/python/src/runtime.rs b/python/src/runtime.rs index 13951a7aa..e8660b38b 100644 --- a/python/src/runtime.rs +++ b/python/src/runtime.rs @@ -13,10 +13,23 @@ //! by an [`AtomicPtr`] to a tokio runtime that we own. A `pthread_atfork` //! child handler nulls the pointer; the next `spawn` rebuilds the runtime in //! the child. This mirrors the pattern used in the Lance Python bindings. +//! +//! Normal (non-fork) process exit has no equivalent handling: nothing ever +//! tells this runtime to shut down, so its worker threads are simply still +//! running, uncoordinated with the interpreter, right up until the process +//! ends. If one of them is mid-task exactly as `Py_Finalize` starts tearing +//! down interpreter state, it can panic on state that's already gone -- and +//! since that happens on a background thread with no PyO3-wrapped call frame +//! to catch it (or, worse, while another unwind is already in progress), +//! Rust aborts the whole process rather than failing that one call. See +//! [`shutdown`], registered as a Python `atexit` callback, which closes that +//! gap by giving the runtime a coordinated, bounded exit while the +//! interpreter is still fully valid. use std::future::Future; use std::pin::Pin; use std::sync::atomic::{AtomicBool, AtomicPtr, Ordering}; +use std::time::Duration; use pyo3::{Bound, PyAny, PyResult, Python, conversion::IntoPyObject}; use pyo3_async_runtimes::{ @@ -65,6 +78,37 @@ pub fn block_on(fut: F) -> F::Output { get_runtime().block_on(fut) } +/// Gracefully quiesce the owned runtime, meant to run at normal process exit. +/// +/// Takes the runtime through the same atomic slot `get_runtime()` uses, so a +/// caller that races with this either gets the runtime before it is swapped +/// out, or transparently builds a fresh one after (safe, if pointless, this +/// late in the program's life) -- never a dangling reference. +/// +/// `shutdown_timeout` rather than a bare `drop`: dropping a tokio `Runtime` +/// waits (in the worst case indefinitely) for its worker threads to join, +/// which is exactly what the fork handler above avoids by leaking instead. +/// A bounded timeout gives genuinely in-flight work -- a connection pool's +/// keep-alive, a graceful close -- a real chance to finish, then forcibly +/// abandons whatever has not, so this can never hang process exit. +/// +/// Must reset `RUNTIME_INSTALLING` alongside `RUNTIME`, exactly as +/// `atfork_child` does: that flag only ever transitions false -> true, for +/// whichever thread wins the race to build the current runtime, and nothing +/// else ever clears it back to false again. `get_runtime()` relies on it +/// being false precisely when `RUNTIME` is null; leaving it true after +/// nulling `RUNTIME` here would strand every later call in its wait loop +/// forever, since no thread could ever again win that false -> true edge to +/// build a replacement. +pub fn shutdown(timeout: Duration) { + let ptr = RUNTIME.swap(std::ptr::null_mut(), Ordering::SeqCst); + if !ptr.is_null() { + let runtime = unsafe { Box::from_raw(ptr) }; + runtime.shutdown_timeout(timeout); + } + RUNTIME_INSTALLING.store(false, Ordering::SeqCst); +} + /// Runs in async-signal context after `fork()` in the child. We can only /// touch atomics here; we deliberately leak the previous runtime because /// dropping a tokio `Runtime` would try to join its (now-dead) worker @@ -149,3 +193,37 @@ where { pyo3_async_runtimes::generic::future_into_py::(py, fut) } + +#[cfg(test)] +mod tests { + use super::*; + + // One test, not several: RUNTIME/RUNTIME_INSTALLING are process-wide + // statics, and Rust runs tests in parallel by default, so separate test + // functions touching them would race each other. Exercising the whole + // create -> shutdown -> lazily-rebuild -> shutdown-again sequence in one + // function keeps it self-contained instead of pulling in a + // test-serialization dependency for a single file. + #[test] + fn test_shutdown_stops_and_the_runtime_rebuilds_lazily_after() { + // No runtime created yet in this process: shutdown must be a no-op, + // not a null-pointer dereference. + shutdown(Duration::from_secs(1)); + + // Force the runtime into existence, then shut it down. Bounded by + // the timeout, so a hang here means shutdown itself is broken, not + // that the test is slow. + assert_eq!(block_on(async { 1 + 1 }), 2); + shutdown(Duration::from_secs(5)); + + // A caller after shutdown -- e.g. a stray call racing with the + // atexit callback -- must get a fresh, working runtime rather than + // a dangling reference into the one just torn down. + assert_eq!(block_on(async { 2 + 2 }), 4); + + // Shutting down twice in a row (e.g. atexit firing more than once) + // must not panic or double-free. + shutdown(Duration::from_secs(5)); + shutdown(Duration::from_secs(5)); + } +}