rust/lancedb
1
0
Fork 0
mirror of https://github.com/lancedb/lancedb.git synced 2026-11-13 22:26:53 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(python): close two use-after-free/hang windows in runtime shutdown

Browse Source 
The previous commit's shutdown() freed the shared Tokio runtime as soon
as it looked unused (an AtomicPtr swap + Box::from_raw), but callers
only ever held a bare reference into it, with no way to tell shutdown
they were still using it. A concurrent caller could dereference memory
shutdown had already freed.

Switching to an ArcSwapOption<Runtime> so every caller holds an owned
Arc clone for as long as it's using the runtime fixes the
use-after-free, but isn't sufficient on its own, for two reasons:

- future_into_py spawns a task that, once running, spawns a second one
  to do the real work and awaits its JoinHandle. If both nested calls
  independently re-resolved "the current runtime" via get_runtime(), a
  reclaim landing between them could bind the outer and inner task to
  two different runtime instances, permanently stalling the outer task
  if its own runtime was the one torn down while it awaited the inner
  one.

- Runtime::shutdown_timeout gives spawned (non-blocking) tasks no
  grace period at all: a task keeps running until it next yields, then
  is dropped. An outermost caller that only held its Arc for the
  instant it submitted a task let the runtime's reference count fall
  back to baseline immediately, long before the task itself finished,
  so shutdown could reclaim the runtime while a top-level task was
  still genuinely in flight, abandoning it silently.

spawn/spawn_blocking now use Handle::try_current() to pin any nested
call to whatever runtime is already executing it, and the outermost
call moves its Arc clone into the spawned task itself so the reference
count stays elevated for the task's whole lifetime, not just its
submission.

Verified with two new targeted unit tests plus the existing concurrent
stress test, and empirically against a concurrent reproducer built
from this PR's own review feedback: 100+ clean runs with no hangs or
crashes, versus a consistent ~35-55% failure rate beforehand.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
ChilePiquin
2026-09-15 05:41:43 -06:00
co-authored by Claude Sonnet 5
parent 2dec579641
commit 4206a727d2
3 changed files with 310 additions and 68 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/src/runtime.rs
+308 -68
View File
@@ -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<runtime::Runtime> = AtomicPtr::new(std::ptr::null_mut());
static RUNTIME_INSTALLING: AtomicBool = AtomicBool::new(false);
static RUNTIME: ArcSwapOption<runtime::Runtime> = 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<runtime::Runtime> {
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::<Arc<runtime::Runtime>>, 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<F: std::future::Future>(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(&current) <= 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<Box<dyn Future<Output = Result<(), Self::JoinError>> + 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<F>(fut: F) -> Self::JoinHandle
where
F: Future<Output = ()> + 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();
}
}
}
}