diff --git a/rust/CHANGELOG.md b/rust/CHANGELOG.md index 622bcf318f..53ec9e902d 100644 --- a/rust/CHANGELOG.md +++ b/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 diff --git a/rust/crates/sift_stream_bindings/Cargo.toml b/rust/crates/sift_stream_bindings/Cargo.toml index cf09e4655c..0415506bc2 100644 --- a/rust/crates/sift_stream_bindings/Cargo.toml +++ b/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 } diff --git a/rust/crates/sift_stream_bindings/src/lib.rs b/rust/crates/sift_stream_bindings/src/lib.rs index dc87d67caa..8d94d6f900 100644 --- a/rust/crates/sift_stream_bindings/src/lib.rs +++ b/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::().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(()) } diff --git a/rust/crates/sift_stream_bindings/src/runtime.rs b/rust/crates/sift_stream_bindings/src/runtime.rs new file mode 100644 index 0000000000..8b1dca5f0c --- /dev/null +++ b/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> = Mutex::new(None); +/// Handle used to spawn work. Set once, when the runtime is first built. +static HANDLE: OnceLock = 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; +} + +/// 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(fut: F) -> Self::JoinHandle + where + F: Future + Send + 'static, + { + handle().spawn(fut) + } + + fn spawn_blocking(f: F) -> Self::JoinHandle + where + F: FnOnce() + Send + 'static, + { + handle().spawn_blocking(f) + } +} + +impl ContextExt for SiftRuntime { + fn scope(locals: TaskLocals, fut: F) -> Pin + Send>> + where + F: Future + 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 { + 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> +where + F: Future> + Send + 'static, + T: for<'a> IntoPyObject<'a> + Send + 'static, +{ + ensure_started()?; + let locals = generic::get_current_locals::(py)?; + generic::future_into_py_with_locals::(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") +} diff --git a/rust/crates/sift_stream_bindings/src/stream/auto_register.rs b/rust/crates/sift_stream_bindings/src/stream/auto_register.rs index ed139fa04e..e54866be10 100644 --- a/rust/crates/sift_stream_bindings/src/stream/auto_register.rs +++ b/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, diff --git a/rust/crates/sift_stream_bindings/src/stream/builder.rs b/rust/crates/sift_stream_bindings/src/stream/builder.rs index 0b84b2f397..c06fe62cb8 100644 --- a/rust/crates/sift_stream_bindings/src/stream/builder.rs +++ b/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 diff --git a/rust/crates/sift_stream_bindings/src/stream/mod.rs b/rust/crates/sift_stream_bindings/src/stream/mod.rs index 8a5cfc96d9..28717288c3 100644 --- a/rust/crates/sift_stream_bindings/src/stream/mod.rs +++ b/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::{ diff --git a/rust/crates/sift_stream_bindings/tests/test_runtime_shutdown.py b/rust/crates/sift_stream_bindings/tests/test_runtime_shutdown.py new file mode 100644 index 0000000000..82ef5109e3 --- /dev/null +++ b/rust/crates/sift_stream_bindings/tests/test_runtime_shutdown.py @@ -0,0 +1,211 @@ +"""Regression tests for the crate-owned tokio runtime and its shutdown. + +Runtime threads that finish work while the interpreter is finalizing used to +re-enter Python and crash the process (SIGSEGV/SIGABRT on Linux). The module +now registers ``shutdown()`` with ``atexit`` so the runtime is stopped and its +threads joined before finalization starts. + +Every scenario runs in a subprocess: ``shutdown()`` is process-wide and cannot +be undone, and the crash under test happens at interpreter exit. +""" + +import os +import socket +import struct +import subprocess +import sys +import textwrap +import threading + +import pytest +from sift_stream_bindings import is_shut_down, shutdown + +CHILD_TIMEOUT_S = 60 + + +class ResettingListener: + """Accepts and holds connections; resets them all when told to over a control socket. + + Each ``build()`` against ``uri`` connects and then waits for a gRPC reply that never + comes. Resetting the connection fails the build a few milliseconds later, which makes a + runtime thread re-enter Python to reject the awaitable. The control socket lets the child + choose the exact moment, including from inside interpreter finalization. + """ + + def __init__(self) -> None: + self._data = socket.create_server(("127.0.0.1", 0), backlog=128) + self._ctrl = socket.create_server(("127.0.0.1", 0), backlog=16) + self.uri = f"http://127.0.0.1:{self._data.getsockname()[1]}" + self.ctrl_port = self._ctrl.getsockname()[1] + self._lock = threading.Lock() + self._conns: list[socket.socket] = [] + threading.Thread(target=self._accept_data, daemon=True).start() + threading.Thread(target=self._accept_ctrl, daemon=True).start() + + def _accept_data(self) -> None: + while True: + conn, _ = self._data.accept() + with self._lock: + self._conns.append(conn) + + def _accept_ctrl(self) -> None: + while True: + conn, _ = self._ctrl.accept() + threading.Thread(target=self._reset_on_signal, args=(conn,), daemon=True).start() + + def _reset_on_signal(self, ctrl: socket.socket) -> None: + try: + ctrl.recv(2) + finally: + with self._lock: + conns, self._conns = self._conns, [] + for conn in conns: + # SO_LINGER 0 turns close() into a TCP RST. + conn.setsockopt(socket.SOL_SOCKET, socket.SO_LINGER, struct.pack("ii", 1, 0)) + conn.close() + + +def run_child(source: str, *args: str) -> subprocess.CompletedProcess: + return subprocess.run( + [sys.executable, "-c", textwrap.dedent(source), *args], + capture_output=True, + text=True, + timeout=CHILD_TIMEOUT_S, + env={**os.environ, "RUST_BACKTRACE": "0"}, + ) + + +def assert_exited_cleanly(child: subprocess.CompletedProcess) -> None: + assert child.returncode == 0, f"child exited {child.returncode}\n{child.stderr}" + assert "panicked" not in child.stderr, child.stderr + assert "Fatal Python error" not in child.stderr, child.stderr + + +# Starts a few builds that stay pending, then arranges for their connections to be reset +# from a __del__ that runs during module teardown, after CPython has marked itself +# finalizing. Without the atexit shutdown the runtime threads re-enter Python at that +# point: pyo3 panics on every worker, and on Linux the process dies by signal. +PENDING_BUILDS_THEN_EXIT = """ + import asyncio, os, socket, sys, time + from sift_stream_bindings import IngestionConfigFormPy, SiftStreamBuilderPy + + uri, ctrl_port, mode = sys.argv[1], int(sys.argv[2]), sys.argv[3] + ctrl = socket.create_connection(("127.0.0.1", ctrl_port)) + ctrl_fd = os.dup(ctrl.fileno()) # a raw fd survives module teardown; the socket object may not + + async def start_builds(): + for _ in range(5): + builder = SiftStreamBuilderPy(uri=uri, apikey="not-a-real-key") + builder.enable_tls = False + builder.ingestion_config_form = IngestionConfigFormPy( + asset_name="repro", client_key="repro", flows=[] + ) + builder.build() + await asyncio.sleep(0.6) # let every build connect and block on its first request + + asyncio.run(start_builds()) + + class ResetDuringFinalization: + def __init__(self, fd, write, sleep): + self.fd, self.write, self.sleep = fd, write, sleep + + def __del__(self): + self.write(self.fd, b"go") + self.sleep(0.5) # releases the GIL so runtime threads can try to re-enter + + _keep = ResetDuringFinalization(ctrl_fd, os.write, time.sleep) + if mode == "explicit-shutdown": + import sift_stream_bindings + sift_stream_bindings.shutdown() + assert sift_stream_bindings.is_shut_down() +""" + + +@pytest.mark.parametrize("mode", ["atexit", "explicit-shutdown"]) +def test_runtime_threads_do_not_reenter_python_during_finalization(mode: str) -> None: + listener = ResettingListener() + child = run_child(PENDING_BUILDS_THEN_EXIT, listener.uri, str(listener.ctrl_port), mode) + assert_exited_cleanly(child) + + +def test_shutdown_is_idempotent_and_blocks_new_work() -> None: + child = run_child( + """ + import asyncio + import sift_stream_bindings + from sift_stream_bindings import IngestionConfigFormPy, SiftStreamBuilderPy + + assert not sift_stream_bindings.is_shut_down() + sift_stream_bindings.shutdown(timeout=1.0) + sift_stream_bindings.shutdown() # second call is a no-op + assert sift_stream_bindings.is_shut_down() + + async def main(): + builder = SiftStreamBuilderPy(uri="http://127.0.0.1:1", apikey="k") + builder.ingestion_config_form = IngestionConfigFormPy( + asset_name="a", client_key="a", flows=[] + ) + try: + builder.build() + except RuntimeError as e: + assert "shut down" in str(e), e + else: + raise AssertionError("build() should fail after shutdown()") + + asyncio.run(main()) + """ + ) + assert_exited_cleanly(child) + + +def test_shutdown_before_any_use_is_safe() -> None: + child = run_child( + """ + import sift_stream_bindings + sift_stream_bindings.shutdown() + assert sift_stream_bindings.is_shut_down() + """ + ) + assert_exited_cleanly(child) + + +def test_shutdown_warns_when_threads_outlive_timeout() -> None: + # Builds that are blocked on a server that never replies keep runtime threads busy. A zero + # timeout cannot wait for them, so shutdown() must say that threads are still alive. + listener = ResettingListener() + child = run_child( + """ + import asyncio, sys, warnings + import sift_stream_bindings + from sift_stream_bindings import IngestionConfigFormPy, SiftStreamBuilderPy + + async def main(): + for _ in range(3): + builder = SiftStreamBuilderPy(uri=sys.argv[1], apikey="not-a-real-key") + builder.enable_tls = False + builder.ingestion_config_form = IngestionConfigFormPy( + asset_name="repro", client_key="repro", flows=[] + ) + builder.build() + await asyncio.sleep(0.3) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + sift_stream_bindings.shutdown(timeout=0.0) + messages = [str(w.message) for w in caught if w.category is RuntimeWarning] + assert any("timed out" in m and "still running" in m for m in messages), caught + + asyncio.run(main()) + """, + listener.uri, + ) + assert_exited_cleanly(child) + + +def test_shutdown_rejects_bad_timeout() -> None: + # Runs in-process: an invalid timeout is rejected before the runtime is touched. + assert not is_shut_down() + with pytest.raises(ValueError): + shutdown(timeout=-1.0) + with pytest.raises(ValueError): + shutdown(timeout=float("nan")) + assert not is_shut_down()