rust(bug): fix python bindings segfault

rust/CHANGELOG.md

@@ -16,6 +16,20 @@ the word "link".
 
 ### Bug Fixes
 
+#### Runtime threads no longer crash the Python process at interpreter exit
+
+##### `sift-stream-bindings` 0.5.2
+
+The bindings ran every async call on a process-global tokio runtime that was never shut down.
+A worker that finished work while CPython was finalizing re-entered Python and killed the
+process with SIGSEGV or SIGABRT after the program had already completed. The crate now owns its
+runtime and registers a new `shutdown(timeout=5.0)` function with `atexit`, so the runtime is
+stopped and its threads joined while the interpreter is still intact. `shutdown()` can also be
+called directly, for example before `os._exit()`, and `is_shut_down()` reports whether it has
+run. Calls that need the runtime raise `RuntimeError` afterwards. If threads are still alive when
+the timeout expires, `shutdown()` logs a warning and issues a `RuntimeWarning`, since those threads
+can still crash the process. Runtime threads are now named `sift-stream-worker`.
+
 #### Asset-name check is now case-insensitive
 
 `SiftStream` no longer rejects an ingestion config whose asset name differs from the one in Sift

rust/crates/sift_stream_bindings/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "sift-stream-bindings"
-version = "0.5.1"
+version = "0.5.2"
 # Released to PyPI as a wheel, never to crates.io.
 publish = false
 edition = { workspace = true }

rust/crates/sift_stream_bindings/src/lib.rs

@@ -1,11 +1,15 @@
+use pyo3::exceptions::{PyRuntimeWarning, PyValueError};
 use pyo3::prelude::*;
 use pyo3_stub_gen::define_stub_info_gatherer;
+use std::ffi::CString;
 use std::sync::{Mutex, Once};
+use std::time::Duration;
 use tracing::Level;
 use tracing_subscriber::{Layer, filter, layer::SubscriberExt};
 
 mod error;
 mod metrics;
+mod runtime;
 mod sift;
 mod stream;
 
@@ -135,6 +139,54 @@ fn is_tracing_initialized() -> bool {
     INIT_TRACING.is_completed()
 }
 
+/// Stop the background runtime that runs this library's async work and wait up to
+/// `timeout` seconds for its threads to finish.
+///
+/// The module registers this with `atexit` when it is imported, so a normal interpreter
+/// shutdown calls it for you. Without it, runtime threads that complete work while the
+/// interpreter is finalizing re-enter Python and crash the process (SIGSEGV or SIGABRT).
+/// `atexit` handlers registered after this module was imported run before it, so they can
+/// still finish a stream. Call it yourself to stop the runtime earlier, for example right
+/// before `os._exit()`.
+///
+/// Finish your streams first. In-flight work is cancelled, awaitables that have not
+/// completed never resolve, and any later call that needs the runtime raises `RuntimeError`.
+/// Calling this more than once is harmless.
+///
+/// If runtime threads are still alive when the timeout expires, this logs a warning through
+/// `tracing` and issues a Python `RuntimeWarning`, because those threads can still re-enter
+/// Python during finalization. Raise the timeout or finish streams earlier if you see it.
+///
+/// Args:
+///     timeout: Seconds to wait for in-flight work before giving up (default: 5.0)
+#[pyfunction]
+#[pyo3(signature = (timeout = runtime::DEFAULT_SHUTDOWN_TIMEOUT_SECS))]
+fn shutdown(py: Python<'_>, timeout: f64) -> PyResult<()> {
+    if !timeout.is_finite() || timeout < 0.0 {
+        return Err(PyValueError::new_err(
+            "timeout must be a non-negative number of seconds",
+        ));
+    }
+    let leftover = runtime::shutdown(py, Duration::from_secs_f64(timeout));
+    if leftover > 0 {
+        let message = format!(
+            "sift_stream_bindings.shutdown() timed out after {timeout}s with {leftover} runtime \
+             thread(s) still running; they may re-enter Python during interpreter finalization \
+             and crash the process. Finish streams before exit or raise the timeout."
+        );
+        tracing::warn!("{message}");
+        let message = CString::new(message)?;
+        PyErr::warn(py, py.get_type::<PyRuntimeWarning>().as_any(), &message, 1)?;
+    }
+    Ok(())
+}
+
+/// Whether `shutdown()` has run. Calls that need the runtime raise `RuntimeError` after it.
+#[pyfunction]
+fn is_shut_down() -> bool {
+    runtime::is_shut_down()
+}
+
 // Cannot organize into submodules right now
 // See below for issues with submodules using pyo3
 // https://github.com/PyO3/pyo3/issues/759
@@ -180,5 +232,14 @@ fn sift_stream_bindings(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(init_tracing, m)?)?;
     m.add_function(wrap_pyfunction!(init_tracing_with_file, m)?)?;
     m.add_function(wrap_pyfunction!(is_tracing_initialized, m)?)?;
+    m.add_function(wrap_pyfunction!(is_shut_down, m)?)?;
+
+    // Stop the runtime while the interpreter is still intact. `atexit` runs before CPython
+    // marks itself finalizing, so runtime threads never re-enter Python after that point.
+    let shutdown_fn = wrap_pyfunction!(shutdown, m)?;
+    m.add_function(shutdown_fn.clone())?;
+    m.py()
+        .import("atexit")?
+        .call_method1("register", (shutdown_fn,))?;
     Ok(())
 }

rust/crates/sift_stream_bindings/src/runtime.rs

@@ -0,0 +1,185 @@
+//! The tokio runtime that backs every awaitable this crate hands to Python.
+//!
+//! `pyo3_async_runtimes::tokio` keeps its runtime in a `OnceLock` with no way to
+//! shut it down, so its worker threads stay alive through interpreter
+//! finalization. A worker that completes a future after CPython has torn down
+//! its thread-state bookkeeping re-enters Python to resolve the asyncio future
+//! and crashes the process: SIGSEGV once `PyGILState_Ensure` dereferences the
+//! freed interpreter state, or SIGABRT when CPython calls `pthread_exit` on a
+//! thread that has Rust frames on its stack.
+//!
+//! This module owns the runtime instead. [`shutdown`] stops it and joins its
+//! threads while the interpreter is still intact, and the module init registers
+//! it with `atexit` so normal interpreter shutdown always runs it. Everything
+//! else (task locals, resolving the asyncio future, cancellation) still comes
+//! from `pyo3_async_runtimes::generic`, driven through [`SiftRuntime`].
+
+use pyo3::exceptions::PyRuntimeError;
+use pyo3::prelude::*;
+use pyo3_async_runtimes::TaskLocals;
+use pyo3_async_runtimes::generic::{self, ContextExt, Runtime as GenericRuntime};
+use std::cell::OnceCell;
+use std::future::Future;
+use std::pin::Pin;
+use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
+use std::sync::{Mutex, OnceLock, PoisonError};
+use std::time::Duration;
+use tokio::runtime::{Builder, Handle, Runtime};
+
+/// How long [`shutdown`] waits for in-flight work when the caller gives no timeout.
+pub const DEFAULT_SHUTDOWN_TIMEOUT_SECS: f64 = 5.0;
+
+/// The runtime itself. `None` before first use and again after [`shutdown`].
+static RUNTIME: Mutex<Option<Runtime>> = Mutex::new(None);
+/// Handle used to spawn work. Set once, when the runtime is first built.
+static HANDLE: OnceLock<Handle> = OnceLock::new();
+/// Set by [`shutdown`] before it drains `RUNTIME`; never cleared.
+static SHUT_DOWN: AtomicBool = AtomicBool::new(false);
+/// Runtime threads (workers and blocking pool) currently alive, kept by tokio's thread hooks.
+static LIVE_THREADS: AtomicUsize = AtomicUsize::new(0);
+
+tokio::task_local! {
+    static TASK_LOCALS: OnceCell<TaskLocals>;
+}
+
+/// Marker type that plugs this crate's runtime into `pyo3_async_runtimes::generic`.
+pub(crate) struct SiftRuntime;
+
+impl GenericRuntime for SiftRuntime {
+    type JoinError = tokio::task::JoinError;
+    type JoinHandle = tokio::task::JoinHandle<()>;
+
+    fn spawn<F>(fut: F) -> Self::JoinHandle
+    where
+        F: Future<Output = ()> + Send + 'static,
+    {
+        handle().spawn(fut)
+    }
+
+    fn spawn_blocking<F>(f: F) -> Self::JoinHandle
+    where
+        F: FnOnce() + Send + 'static,
+    {
+        handle().spawn_blocking(f)
+    }
+}
+
+impl ContextExt for SiftRuntime {
+    fn scope<F, R>(locals: TaskLocals, fut: F) -> Pin<Box<dyn Future<Output = R> + Send>>
+    where
+        F: Future<Output = R> + Send + 'static,
+    {
+        let cell = OnceCell::new();
+        cell.set(locals)
+            .unwrap_or_else(|_| unreachable!("fresh OnceCell is empty"));
+        Box::pin(TASK_LOCALS.scope(cell, fut))
+    }
+
+    fn get_task_locals() -> Option<TaskLocals> {
+        TASK_LOCALS
+            .try_with(|c| c.get().cloned())
+            .unwrap_or_default()
+    }
+}
+
+/// Converts a Rust future into a Python awaitable that runs on this crate's runtime.
+///
+/// Drop-in replacement for `pyo3_async_runtimes::tokio::future_into_py`. Raises
+/// `RuntimeError` once [`shutdown`] has run, since there is no runtime left to
+/// drive the future.
+pub(crate) fn future_into_py<'py, F, T>(py: Python<'py>, fut: F) -> PyResult<Bound<'py, PyAny>>
+where
+    F: Future<Output = PyResult<T>> + Send + 'static,
+    T: for<'a> IntoPyObject<'a> + Send + 'static,
+{
+    ensure_started()?;
+    let locals = generic::get_current_locals::<SiftRuntime>(py)?;
+    generic::future_into_py_with_locals::<SiftRuntime, _, _>(py, locals, fut)
+}
+
+/// Whether [`shutdown`] has run.
+pub(crate) fn is_shut_down() -> bool {
+    SHUT_DOWN.load(Ordering::SeqCst)
+}
+
+/// Stops the runtime and waits up to `timeout` for its threads to finish.
+///
+/// Work still in flight is cancelled at its next await point, awaitables that
+/// have not completed never resolve, and every later call that needs the
+/// runtime raises `RuntimeError`. The GIL is released while waiting so that a
+/// task which is already resolving a Python future can finish. Calling this more
+/// than once is harmless.
+///
+/// Returns the number of runtime threads still alive when the wait ended. Tokio
+/// gives no signal when `shutdown_timeout` gives up, so this is how a caller
+/// learns that threads outlived the timeout and could still re-enter Python.
+pub(crate) fn shutdown(py: Python<'_>, timeout: Duration) -> usize {
+    SHUT_DOWN.store(true, Ordering::SeqCst);
+    let runtime = RUNTIME
+        .lock()
+        .unwrap_or_else(PoisonError::into_inner)
+        .take();
+    match runtime {
+        Some(runtime) => py.detach(|| {
+            runtime.shutdown_timeout(timeout);
+            LIVE_THREADS.load(Ordering::SeqCst)
+        }),
+        None => 0,
+    }
+}
+
+fn shut_down_err() -> PyErr {
+    PyRuntimeError::new_err(
+        "the sift_stream_bindings runtime has been shut down; it stops at interpreter exit \
+         or when sift_stream_bindings.shutdown() is called",
+    )
+}
+
+/// Builds the runtime on first use. Fails with `RuntimeError` after [`shutdown`].
+fn ensure_started() -> PyResult<()> {
+    if is_shut_down() {
+        return Err(shut_down_err());
+    }
+    if HANDLE.get().is_some() {
+        return Ok(());
+    }
+
+    let mut slot = RUNTIME.lock().unwrap_or_else(PoisonError::into_inner);
+    if HANDLE.get().is_some() {
+        // Another thread built it while we waited for the lock.
+        return Ok(());
+    }
+    let runtime = Builder::new_multi_thread()
+        .enable_all()
+        .thread_name("sift-stream-worker")
+        .on_thread_start(|| {
+            LIVE_THREADS.fetch_add(1, Ordering::SeqCst);
+        })
+        .on_thread_stop(|| {
+            LIVE_THREADS.fetch_sub(1, Ordering::SeqCst);
+        })
+        .build()
+        .map_err(|e| {
+            PyRuntimeError::new_err(format!(
+                "failed to start the sift_stream_bindings runtime: {e}"
+            ))
+        })?;
+    let _ = HANDLE.set(runtime.handle().clone());
+    // `shutdown` sets the flag before draining the slot, so if it ran while we
+    // were building, drop this runtime now rather than leave its threads alive.
+    if is_shut_down() {
+        drop(slot);
+        drop(runtime);
+        return Err(shut_down_err());
+    }
+    *slot = Some(runtime);
+    Ok(())
+}
+
+/// The spawn handle. Every caller goes through [`future_into_py`], which runs
+/// [`ensure_started`] first, so the handle exists by the time this is reached.
+fn handle() -> &'static Handle {
+    HANDLE
+        .get()
+        .expect("sift_stream_bindings runtime is started before work is spawned on it")
+}

rust/crates/sift_stream_bindings/src/stream/auto_register.rs

@@ -1,9 +1,9 @@
 use super::{FlowPy, SiftStreamInner, SiftStreamPy, stream_consumed_err};
 use crate::error::SiftErrorWrapper;
 use crate::metrics::SiftStreamMetricsSnapshotPy;
+use crate::runtime::future_into_py;
 use crate::stream::config::{FlowConfigPy, FlowDescriptorPy, RunSelectorPy};
 use pyo3::prelude::*;
-use pyo3_async_runtimes::tokio::future_into_py;
 use pyo3_stub_gen::derive::*;
 use sift_stream::{
     AutoRegisterStream, FileBackup, Flow, LiveStreamingOnly, LiveStreamingWithBackups,

rust/crates/sift_stream_bindings/src/stream/builder.rs

@@ -128,7 +128,7 @@ impl SiftStreamBuilderPy {
 
         let mode_builder = config_builder.live_only();
 
-        let awaitable = pyo3_async_runtimes::tokio::future_into_py(py, async move {
+        let awaitable = crate::runtime::future_into_py(py, async move {
             match mode_builder.build().await {
                 Ok(stream) => Ok(SiftStreamPy::from(stream)),
                 Err(e) => Err(crate::error::SiftErrorWrapper(e).into()),
@@ -321,7 +321,7 @@ impl LiveOnlyBuilderPy {
         let ingestion_cap = self.ingestion_data_channel_capacity;
         let control_cap = self.control_channel_capacity;
 
-        let awaitable = pyo3_async_runtimes::tokio::future_into_py(py, async move {
+        let awaitable = crate::runtime::future_into_py(py, async move {
             let config_builder = make_stream_config_builder(base)?;
 
             let mode_builder = config_builder
@@ -408,7 +408,7 @@ impl LiveWithBackupsBuilderPy {
         let backup_cap = self.backup_data_channel_capacity;
         let control_cap = self.control_channel_capacity;
 
-        let awaitable = pyo3_async_runtimes::tokio::future_into_py(py, async move {
+        let awaitable = crate::runtime::future_into_py(py, async move {
             let config_builder = make_stream_config_builder(base)?;
 
             let mode_builder = config_builder
@@ -481,7 +481,7 @@ impl FileBackupBuilderPy {
             .metrics_streaming_interval
             .map(std::time::Duration::from);
 
-        let awaitable = pyo3_async_runtimes::tokio::future_into_py(py, async move {
+        let awaitable = crate::runtime::future_into_py(py, async move {
             let config_builder = make_stream_config_builder(base)?;
 
             let mode_builder = config_builder

rust/crates/sift_stream_bindings/src/stream/mod.rs

@@ -8,12 +8,12 @@ pub mod time;
 
 use crate::error::SiftErrorWrapper;
 use crate::metrics::SiftStreamMetricsSnapshotPy;
+use crate::runtime::future_into_py;
 use crate::stream::channel::ChannelValuePy;
 use crate::stream::config::{FlowConfigPy, FlowDescriptorPy, RunSelectorPy};
 use crate::stream::request::IngestWithConfigDataStreamRequestWrapperPy;
 use crate::stream::time::TimeValuePy;
 use pyo3::{prelude::*, types::PyIterator};
-use pyo3_async_runtimes::tokio::future_into_py;
 use pyo3_stub_gen::derive::*;
 use sift_rs::ingest::v1::IngestWithConfigDataStreamRequest;
 use sift_stream::{