diff --git a/Cargo.lock b/Cargo.lock index b3ff0e299..2e2c0c296 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5571,6 +5571,7 @@ dependencies = [ name = "lancedb-python" version = "0.39.0-beta.8" dependencies = [ + "arc-swap", "arrow", "async-trait", "bytes", diff --git a/python/Cargo.toml b/python/Cargo.toml index d7ec8cd61..314664e14 100644 --- a/python/Cargo.toml +++ b/python/Cargo.toml @@ -15,6 +15,7 @@ name = "_lancedb" crate-type = ["cdylib"] [dependencies] +arc-swap = "1.9" arrow = { workspace = true, features = ["pyarrow"] } async-trait.workspace = true bytes.workspace = true diff --git a/python/src/runtime.rs b/python/src/runtime.rs index e8660b38b..970595505 100644 --- a/python/src/runtime.rs +++ b/python/src/runtime.rs @@ -10,27 +10,65 @@ //! //! We sidestep the global by routing every future through our own //! [`LanceRuntime`] (a [`pyo3_async_runtimes::generic::Runtime`] impl) backed -//! 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 +//! by an [`ArcSwapOption`] holding the tokio runtime we own. A +//! `pthread_atfork` child handler clears the slot without dropping what it +//! finds there (see `atfork_child`); the next call 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. +//! Normal (non-fork) process exit has its own gap: nothing tells the runtime +//! to shut down, so its worker threads keep 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, Rust +//! aborts the whole process rather than failing that one call. [`shutdown`], +//! registered as a Python `atexit` callback, closes that gap by giving the +//! runtime a coordinated, bounded exit while the interpreter is still fully +//! valid. +//! +//! Every reader holds its own `Arc` clone for as long as it is using the +//! runtime (see `get_runtime`), rather than a bare reference into the slot. +//! `shutdown` only ever reclaims (drops) the runtime once it can prove it +//! holds the only remaining reference; if some other clone is still +//! outstanding when its bound elapses, it abandons the runtime instead -- +//! the same trade `atfork_child` already makes -- rather than ever risk +//! freeing memory a live caller might still be using. +//! +//! That alone is not enough, for two compounding reasons. +//! +//! First, `future_into_py` does not make one `get_runtime()`-mediated call +//! per logical operation -- it spawns a task that, once running, spawns a +//! second one to do the real work and awaits its `JoinHandle`. If the two +//! nested calls independently re-resolved "the current runtime", a reclaim +//! landing between them could bind the outer and inner task to two +//! *different* instances; tearing down the outer task's runtime while it +//! awaits the inner one then leaves it parked forever. `spawn`/ +//! `spawn_blocking` below close that gap with `Handle::try_current`: a call +//! already running on one of our runtime's worker threads is pinned to that +//! same instance regardless of what the global slot holds, so only the +//! first, outermost call of a chain (from a thread outside any runtime) +//! ever consults it. +//! +//! Second, and more fundamentally: `Runtime::shutdown_timeout` does not +//! give spawned (non-blocking) tasks a bounded grace period at all -- per +//! its own docs, a task "keeps running until it yields, then is dropped". +//! An outermost call that only held its `Arc` for the instant it submitted +//! the task (as `get_runtime()`'s doc once assumed of every caller) let +//! `strong_count` fall back to baseline immediately, long before the task +//! itself finished -- so `shutdown` could, and empirically did, reclaim the +//! runtime while a top-level task was still in flight, silently abandoning +//! it and leaving whatever Python future it was going to resolve unresolved +//! forever. `spawn`/`spawn_blocking`'s outermost branch now moves an `Arc` +//! clone *into* the task itself, so `strong_count` stays elevated for the +//! task's entire lifetime, not just the submission call. use std::future::Future; use std::pin::Pin; -use std::sync::atomic::{AtomicBool, AtomicPtr, Ordering}; -use std::time::Duration; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::time::{Duration, Instant}; +use arc_swap::ArcSwapOption; use pyo3::{Bound, PyAny, PyResult, Python, conversion::IntoPyObject}; use pyo3_async_runtimes::{ TaskLocals, @@ -38,8 +76,7 @@ use pyo3_async_runtimes::{ }; use tokio::{runtime, task}; -static RUNTIME: AtomicPtr = AtomicPtr::new(std::ptr::null_mut()); -static RUNTIME_INSTALLING: AtomicBool = AtomicBool::new(false); +static RUNTIME: ArcSwapOption = ArcSwapOption::const_empty(); static ATFORK_INSTALLED: AtomicBool = AtomicBool::new(false); fn create_runtime() -> runtime::Runtime { @@ -50,23 +87,35 @@ fn create_runtime() -> runtime::Runtime { .expect("Failed to build tokio runtime") } -fn get_runtime() -> &'static runtime::Runtime { +/// Get a live, owned handle to the shared runtime. +/// +/// Returns an owned `Arc` rather than a bare reference so the runtime cannot +/// be freed out from under a caller still using it: as long as any clone is +/// held, `shutdown` will not reclaim the runtime it points to. Callers below +/// rely on Rust's own temporary-drop timing to hold the returned `Arc` for +/// the duration of one `block_on`, `spawn`, or `spawn_blocking` call -- +/// exactly as long as the runtime is actually in use, not just to obtain it. +fn get_runtime() -> Arc { loop { - let ptr = RUNTIME.load(Ordering::SeqCst); - if !ptr.is_null() { - return unsafe { &*ptr }; + if let Some(existing) = RUNTIME.load_full() { + return existing; } - if !RUNTIME_INSTALLING.fetch_or(true, Ordering::SeqCst) { - break; + if !ATFORK_INSTALLED.fetch_or(true, Ordering::SeqCst) { + install_atfork(); } - std::thread::yield_now(); + // Built optimistically, outside any lock: on the rare race where two + // threads both find the slot empty, one candidate wins the + // compare-and-swap below and the other is simply dropped here, + // tearing down its own (never shared, never used) worker pool the + // ordinary way. + let candidate = Arc::new(create_runtime()); + let previous = + RUNTIME.compare_and_swap(&None::>, Some(Arc::clone(&candidate))); + if previous.is_none() { + return candidate; + } + // Someone else's candidate won; go around and load it. } - if !ATFORK_INSTALLED.fetch_or(true, Ordering::SeqCst) { - install_atfork(); - } - let new_ptr = Box::into_raw(Box::new(create_runtime())); - RUNTIME.store(new_ptr, Ordering::SeqCst); - unsafe { &*new_ptr } } /// Block the current thread on a future using the shared runtime. @@ -78,44 +127,70 @@ 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. +/// Gracefully quiesce the shared 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. +/// Polls (bounded by `timeout`) for every other outstanding `Arc` clone -- +/// each held by a caller still actually using the runtime, see +/// `get_runtime` -- to be dropped, and only once none remain does it remove +/// the runtime from the slot and reclaim it, so reclaiming can never free +/// memory anyone else might still touch. If the bound elapses first, it +/// leaves the runtime exactly where it is rather than forcing the issue: +/// the process is exiting either way, so an unreclaimed runtime here costs +/// nothing a crash would cost more. /// -/// `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. +/// Deliberately does not empty the slot up front and reclaim once ownership +/// clears, the way a first version of this function did: `get_runtime()` +/// treats an empty slot as "nothing built yet" and responds by building a +/// whole new multi-threaded runtime, worker pool included. Emptying the slot +/// before we are actually ready to consume what was in it means every other +/// concurrent caller sees that empty slot too, and a caller looping tightly +/// (a busy poller, say) would rebuild a brand new runtime on every single +/// call for as long as the slot stayed empty -- not a crash, but a +/// self-inflicted thundering herd that starves real work for the same +/// bounded window this function is supposed to just wait out quietly. /// -/// 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. +/// `shutdown_timeout` rather than a bare `drop`, once exclusive ownership is +/// proven: dropping a tokio `Runtime` waits (in the worst case indefinitely) +/// for its worker threads to join, whereas `shutdown_timeout` gives real +/// in-flight work -- a connection pool's keep-alive, a graceful close -- a +/// bounded chance to finish first, then forcibly ends whatever has not. +/// +/// Safe to call even while operations are still in flight on the runtime +/// being reclaimed: `spawn`/`spawn_blocking` below pin every nested call +/// spawned from an already-running task to that same task's runtime (via +/// `Handle::try_current`), so reclaiming this instance out from under a +/// task that outlives the exclusivity check can, at worst, make +/// `shutdown_timeout` leak that task's worker threads to finish in the +/// background -- it can never split one logical operation across two +/// different runtime instances. 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); + let Some(current) = RUNTIME.load_full() else { + return; + }; + let deadline = Instant::now() + timeout; + loop { + if Arc::strong_count(¤t) <= 2 { + RUNTIME.store(None); + if let Ok(runtime) = Arc::try_unwrap(current) { + runtime.shutdown_timeout(deadline.saturating_duration_since(Instant::now())); + } + return; + } + if Instant::now() >= deadline { + return; + } + std::thread::sleep(Duration::from_millis(1)); } - 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 -/// threads and hang. +/// Runs in async-signal context after `fork()` in the child. We can only +/// touch the atomic slot here; we deliberately do not drop whatever runtime +/// we find, because dropping a tokio `Runtime` would try to join its +/// (now-dead) worker threads and hang. extern "C" fn atfork_child() { - RUNTIME.store(std::ptr::null_mut(), Ordering::SeqCst); - RUNTIME_INSTALLING.store(false, Ordering::SeqCst); + if let Some(orphaned) = RUNTIME.swap(None) { + std::mem::forget(orphaned); + } } #[cfg(not(windows))] @@ -146,11 +221,46 @@ impl Runtime for LanceRuntime { type JoinError = LanceJoinError; type JoinHandle = Pin> + Send>>; + /// `pyo3_async_runtimes::generic::future_into_py` spawns a task that, + /// once it starts running, spawns a second one to do the real work (and + /// spawn_blocking calls beyond that to hand the result back to Python). + /// If each of those calls re-resolved "the current runtime" independently + /// via `get_runtime()`, a reclaim landing between them could bind the + /// outer and inner tasks to two different runtime instances -- and if + /// the outer task's own runtime is the one torn down while it awaits the + /// inner task's `JoinHandle`, it never resumes. `Handle::try_current` + /// sidesteps the global slot entirely whenever we're already running + /// inside a runtime, pinning nested spawns to that same instance; only a + /// call from a genuine outside thread (the first, outermost spawn of a + /// call chain) falls through to `get_runtime()`. + /// + /// That call also has to hold its `Arc` for as long as the task it + /// starts is actually running, not just for the moment it submits it: + /// `Runtime::shutdown_timeout` does not wait for spawned (non-blocking) + /// tasks to finish, it lets each run until it next yields and then + /// *drops* it -- so `shutdown` reclaiming the runtime while a top-level + /// task is still in flight would silently abandon it mid-await, and + /// whatever Python future it was going to resolve would never resolve. + /// Moving a clone into the task itself keeps `strong_count` elevated for + /// the task's whole lifetime, so `shutdown`'s exclusivity check (see + /// `shutdown`) does not consider the runtime free until every top-level + /// task genuinely has completed or been abandoned by some earlier, + /// already-accounted-for reclaim. fn spawn(fut: F) -> Self::JoinHandle where F: Future + Send + 'static, { - let handle = get_runtime().spawn(fut); + let handle = match tokio::runtime::Handle::try_current() { + Ok(handle) => handle.spawn(fut), + Err(_) => { + let rt = get_runtime(); + let rt_for_task = Arc::clone(&rt); + rt.spawn(async move { + let _rt = rt_for_task; + fut.await; + }) + } + }; Box::pin(async move { handle.await.map_err(LanceJoinError) }) } @@ -158,7 +268,17 @@ impl Runtime for LanceRuntime { where F: FnOnce() + Send + 'static, { - let handle = get_runtime().spawn_blocking(f); + let handle = match tokio::runtime::Handle::try_current() { + Ok(handle) => handle.spawn_blocking(f), + Err(_) => { + let rt = get_runtime(); + let rt_for_task = Arc::clone(&rt); + rt.spawn_blocking(move || { + let _rt = rt_for_task; + f(); + }) + } + }; Box::pin(async move { handle.await.map_err(LanceJoinError) }) } } @@ -198,12 +318,38 @@ where 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. + // Before an outermost `spawn`/`spawn_blocking` call held its `Arc` for + // the task's whole lifetime, `strong_count` fell back to baseline as + // soon as the task was merely *submitted* -- long before it actually + // finished running. `shutdown_timeout` gives spawned tasks no grace + // period of its own (it lets one run until it next yields, then drops + // it), so a concurrent `shutdown()` could, and empirically did, reclaim + // the runtime out from under a task still in flight, abandoning it + // silently. Yielding a few times keeps this task genuinely in flight + // (not yet started, and not yet finished) when `shutdown` runs its + // exclusivity check immediately after. + #[test] + #[allow(unused_must_use)] // fire-and-forget spawn, same as future_into_py itself + fn test_top_level_task_survives_concurrent_shutdown_reclaim() { + use std::sync::mpsc; + + for _ in 0..50 { + let (tx, rx) = mpsc::channel::<()>(); + + LanceRuntime::spawn(async move { + for _ in 0..5 { + task::yield_now().await; + } + let _ = tx.send(()); + }); + + shutdown(Duration::from_secs(2)); + + rx.recv_timeout(Duration::from_secs(5)) + .expect("top-level task was abandoned by a concurrent shutdown reclaim"); + } + } + #[test] fn test_shutdown_stops_and_the_runtime_rebuilds_lazily_after() { // No runtime created yet in this process: shutdown must be a no-op, @@ -226,4 +372,98 @@ mod tests { shutdown(Duration::from_secs(5)); shutdown(Duration::from_secs(5)); } + + // Adapted from the concurrent reproducer that found the original bug: + // many threads hammering the runtime while shutdown races them, not + // just the sequential rebuild path above. Before every reader held its + // own `Arc`, this could dereference a runtime `shutdown` had already + // freed, or hang forever. Repeated, since a race is not guaranteed to + // show up on any single attempt. + #[test] + fn test_shutdown_is_safe_concurrently_with_live_callers() { + use std::sync::Barrier; + use std::sync::atomic::AtomicBool as StopFlag; + + for _ in 0..50 { + let barrier = Arc::new(Barrier::new(9)); + let stop = Arc::new(StopFlag::new(false)); + let workers: Vec<_> = (0..8) + .map(|_| { + let barrier = Arc::clone(&barrier); + let stop = Arc::clone(&stop); + std::thread::spawn(move || { + barrier.wait(); + while !stop.load(Ordering::Relaxed) { + assert_eq!(block_on(async { 1 + 1 }), 2); + } + }) + }) + .collect(); + + barrier.wait(); + shutdown(Duration::from_millis(50)); + stop.store(true, Ordering::Relaxed); + for worker in workers { + worker.join().unwrap(); + } + } + } + + // Reproduces the actual bug mechanism, not just concurrent `block_on` + // traffic: `future_into_py` spawns an outer task that, once running, + // spawns a second (inner) one for the real work and awaits its + // `JoinHandle` (see `pyo3_async_runtimes::generic::future_into_py_with_locals`). + // Before `spawn`/`spawn_blocking` pinned nested calls to whatever + // runtime is already executing them, a reclaim landing between the + // outer and inner spawn could bind them to two different runtime + // instances -- and if the outer task's own runtime was the one torn + // down while it awaited the inner task, it never resumed. Every wait + // here is bounded so a reintroduced bug fails this test instead of + // hanging the suite. + #[test] + #[allow(unused_must_use)] // fire-and-forget outer spawn, same as future_into_py itself + fn test_nested_spawn_survives_concurrent_shutdown() { + use std::sync::mpsc; + + for _ in 0..50 { + let stop = Arc::new(AtomicBool::new(false)); + let (done_tx, done_rx) = mpsc::channel::<()>(); + + let workers: Vec<_> = (0..8) + .map(|_| { + let stop = Arc::clone(&stop); + let done_tx = done_tx.clone(); + std::thread::spawn(move || { + while !stop.load(Ordering::Relaxed) { + let (tx, rx) = mpsc::sync_channel::<()>(1); + LanceRuntime::spawn(async move { + let inner = LanceRuntime::spawn(async { + let _ = 1 + 1; + }); + let _ = inner.await; + let _ = tx.send(()); + }); + // A hang here across a concurrent shutdown is + // exactly the bug this test exists to catch. + let _ = rx.recv_timeout(Duration::from_secs(2)); + } + let _ = done_tx.send(()); + }) + }) + .collect(); + + std::thread::sleep(Duration::from_millis(10)); + shutdown(Duration::from_millis(50)); + stop.store(true, Ordering::Relaxed); + + for _ in 0..8 { + done_rx + .recv_timeout(Duration::from_secs(10)) + .expect("worker hung after shutdown raced a nested spawn"); + } + for worker in workers { + worker.join().unwrap(); + } + } + } }