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)); + } +}