diff --git a/libdd-common/src/lib.rs b/libdd-common/src/lib.rs index b10df80d75..295889ee17 100644 --- a/libdd-common/src/lib.rs +++ b/libdd-common/src/lib.rs @@ -79,6 +79,17 @@ pub mod unix_utils; /// ``` pub trait MutexExt { fn lock_or_panic(&self) -> MutexGuard<'_, T>; + + /// Acquires the lock, recovering the guard instead of panicking if poisoned. + /// Reserve for terminal teardown (`shutdown`/`Drop`); elsewhere prefer + /// [`MutexExt::lock_or_panic`] since poisoning there indicates a real bug. + /// + /// Default implementation just delegates to `lock_or_panic` (a safe, behavior- + /// preserving default for existing implementors); `Mutex`'s impl below is the + /// one that actually recovers. + fn lock_or_recover(&self) -> MutexGuard<'_, T> { + self.lock_or_panic() + } } impl MutexExt for Mutex { @@ -88,6 +99,11 @@ impl MutexExt for Mutex { #[allow(clippy::unwrap_used)] self.lock().unwrap() } + + #[inline(always)] + fn lock_or_recover(&self) -> MutexGuard<'_, T> { + self.lock().unwrap_or_else(|poisoned| poisoned.into_inner()) + } } /// Extension trait for `RwLock` to provide methods that acquire read/write locks, panicking if diff --git a/libdd-data-pipeline/src/trace_exporter/mod.rs b/libdd-data-pipeline/src/trace_exporter/mod.rs index 5cdefa3db9..ce62f07241 100644 --- a/libdd-data-pipeline/src/trace_exporter/mod.rs +++ b/libdd-data-pipeline/src/trace_exporter/mod.rs @@ -258,10 +258,11 @@ impl< #[cfg(not(target_arch = "wasm32"))] pub fn shutdown(self, timeout: Option) -> Result<(), TraceExporterError> where - R: BlockingRuntime, + R: BlockingRuntime + Send + Sync + 'static, + Self: Send, { let runtime = self.shared_runtime.clone(); - runtime.block_on(self.shutdown_async(timeout))? + runtime.block_on_send(self.shutdown_async(timeout))? } /// Async version of [`Self::shutdown`]. diff --git a/libdd-shared-runtime/src/shared_runtime/fork_safe.rs b/libdd-shared-runtime/src/shared_runtime/fork_safe.rs index 30d5124e8d..e48b1f4bfc 100644 --- a/libdd-shared-runtime/src/shared_runtime/fork_safe.rs +++ b/libdd-shared-runtime/src/shared_runtime/fork_safe.rs @@ -7,6 +7,7 @@ use libdd_common::MutexExt; use std::io; use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::{Arc, Mutex}; +use std::time::Duration; use tokio::runtime::{Builder, Runtime}; use tracing::{debug, error}; @@ -129,27 +130,92 @@ impl ForkSafeRuntime { /// Shuts down all workers synchronously. Returns `ShutdownTimedOut` if `timeout` is /// exceeded. - pub fn shutdown(&self, timeout: Option) -> Result<(), SharedRuntimeError> { + /// + /// This always drives each worker's shutdown to completion so in-flight work (e.g. a + /// final trace flush) is not lost. Normally that happens by calling `block_on` on the + /// current thread. But `block_on` needs a live Tokio context, which requires the + /// calling thread's CONTEXT thread-local — and during (embedded) interpreter + /// finalization that thread-local may already be destroyed (see + /// [`Self::can_block_on`]). The runtime's own worker threads are unaffected by that; + /// only the *calling* thread's TLS is gone. So in that case we hand the same graceful + /// shutdown off to a freshly spawned thread, whose TLS was never touched, and just wait + /// for it to finish (bounded by `timeout`, if given) using plain OS-level + /// synchronization that doesn't touch Tokio's TLS. + pub fn shutdown(&self, timeout: Option) -> Result<(), SharedRuntimeError> { debug!(?timeout, "Shutting down ForkSafeRuntime"); - match self.runtime.lock_or_panic().take() { - Some(runtime) => { - if let Some(timeout) = timeout { - match runtime.block_on(async { - tokio::time::timeout(timeout, ::shutdown_async(self)) - .await - }) { - Ok(()) => Ok(()), - Err(_) => Err(SharedRuntimeError::ShutdownTimedOut(timeout)), - } - } else { - runtime.block_on(::shutdown_async(self)); - Ok(()) + let Some(runtime) = self.runtime.lock_or_recover().take() else { + return Ok(()); + }; + + if Self::can_block_on() { + return Self::block_on_shutdown(&runtime, &self.workers, timeout); + } + + debug!("No live Tokio context (finalization); flushing from a helper thread"); + Self::shutdown_from_helper_thread(runtime, self.workers.clone(), timeout) + } + + /// Drives `shutdown_workers` to completion via `runtime.block_on` on the calling + /// thread. Only safe to call when [`Self::can_block_on`] is true. + fn block_on_shutdown( + runtime: &Runtime, + workers: &Arc>>, + timeout: Option, + ) -> Result<(), SharedRuntimeError> { + match timeout { + Some(timeout) => { + match runtime.block_on(async { + tokio::time::timeout(timeout, shutdown_workers(workers)).await + }) { + Ok(()) => Ok(()), + Err(_) => Err(SharedRuntimeError::ShutdownTimedOut(timeout)), } } - None => Ok(()), + None => { + runtime.block_on(shutdown_workers(workers)); + Ok(()) + } + } + } + + /// Runs [`Self::block_on_shutdown`] on a fresh, detached thread (whose Tokio TLS was + /// never touched) and blocks the calling thread on its completion via a channel — + /// plain OS-level synchronization, safe even with the calling thread's Tokio TLS gone. + /// + /// A missing `timeout` waits indefinitely for completion, matching the behavior of + /// [`Self::block_on_shutdown`] itself when called with `None`. + fn shutdown_from_helper_thread( + runtime: Arc, + workers: Arc>>, + timeout: Option, + ) -> Result<(), SharedRuntimeError> { + let (done_tx, done_rx) = std::sync::mpsc::channel(); + std::thread::spawn(move || { + let result = Self::block_on_shutdown(&runtime, &workers, timeout); + let _ = done_tx.send(result); + }); + + match timeout { + Some(timeout) => done_rx + .recv_timeout(timeout) + .unwrap_or(Err(SharedRuntimeError::ShutdownTimedOut(timeout))), + None => done_rx.recv().unwrap_or(Ok(())), } } + /// Whether the current thread can safely enter a Tokio context (i.e. call `block_on`). + /// + /// Returns `false` only when the Tokio CONTEXT thread-local has been *destroyed*, + /// which happens during interpreter/thread finalization. A merely-missing context + /// (the normal case for a thread that never entered a runtime) still returns `true`, + /// because `block_on` establishes its own context in that case. + fn can_block_on() -> bool { + !matches!( + tokio::runtime::Handle::try_current(), + Err(ref e) if e.is_thread_local_destroyed() + ) + } + fn push_worker( &self, workers_guard: &mut std::sync::MutexGuard>, @@ -169,6 +235,39 @@ impl ForkSafeRuntime { } } +impl Drop for ForkSafeRuntime { + /// Terminal teardown for the owned Tokio runtime. + /// + /// A normal `Runtime` drop blocks the current thread to join its worker threads, + /// and that join path touches the Tokio CONTEXT thread-local. During (embedded) + /// interpreter finalization — e.g. a uWSGI worker exiting — that thread-local may + /// already be destroyed, so a normal drop panics with "The Tokio context + /// thread-local variable has been destroyed" (and, via a poisoned mutex, cascades). + /// + /// `shutdown_background` instead signals shutdown and returns immediately, detaching + /// the worker threads; the OS reclaims them on process exit. This is done + /// *unconditionally* — we do not try to detect whether we are in finalization — + /// so there is no ordering-dependent branch that a future teardown scenario can + /// slip past. + /// + /// [`Self::shutdown`] always takes the runtime out (even during finalization, via a + /// helper thread — see its docs), so this only fires when `shutdown` was never called + /// at all before drop; that caller-error case is the one scenario this crate can't + /// flush on behalf of the caller, so it just detaches cleanly instead of risking a + /// panic. + fn drop(&mut self) { + if let Some(runtime) = self.runtime.lock_or_recover().take() { + match Arc::try_unwrap(runtime) { + Ok(runtime) => runtime.shutdown_background(), + // Another Arc clone is still alive (e.g. an in-flight block_on holds + // one). We cannot take ownership to detach; drop our reference and let + // the last owner's drop run. This is not the finalization path. + Err(_runtime) => {} + } + } + } +} + impl SharedRuntime for ForkSafeRuntime { fn new() -> Result { Self::with_worker_threads(1) @@ -197,25 +296,33 @@ impl SharedRuntime for ForkSafeRuntime { } async fn shutdown_async(&self) { - debug!("Shutting down all workers asynchronously"); - let workers = { - let mut workers_lock = self.workers.lock_or_panic(); - std::mem::take(&mut *workers_lock) - }; + shutdown_workers(&self.workers).await; + } +} - let futures: FuturesUnordered<_> = workers - .into_iter() - .map(|mut worker_entry| async move { - if let Err(e) = worker_entry.worker.pause().await { - error!("Worker failed to shutdown: {:?}", e); - return; - } - worker_entry.worker.shutdown().await; - }) - .collect(); +/// Drains and gracefully shuts down every registered worker. Standalone (rather than a +/// `&self` method) so it can run identically whether driven from the calling thread +/// ([`ForkSafeRuntime::block_on_shutdown`]) or from the helper thread spawned by +/// [`ForkSafeRuntime::shutdown_from_helper_thread`]. +async fn shutdown_workers(workers: &Mutex>) { + debug!("Shutting down all workers asynchronously"); + let workers = { + let mut workers_lock = workers.lock_or_recover(); + std::mem::take(&mut *workers_lock) + }; + + let futures: FuturesUnordered<_> = workers + .into_iter() + .map(|mut worker_entry| async move { + if let Err(e) = worker_entry.worker.pause().await { + error!("Worker failed to shutdown: {:?}", e); + return; + } + worker_entry.worker.shutdown().await; + }) + .collect(); - futures.collect::<()>().await; - } + futures.collect::<()>().await; } impl BlockingRuntime for ForkSafeRuntime { @@ -227,6 +334,53 @@ impl BlockingRuntime for ForkSafeRuntime { }; Ok(runtime.block_on(f)) } + + /// Entering a Tokio context, or building a new runtime, touches thread-locals that + /// panic if the calling thread's TLS is already destroyed (same hazard as + /// [`Self::shutdown`]). When that's the case, run `f` on a helper thread instead. + fn block_on_send( + &self, + f: F, + ) -> Result + where + F::Output: Send + 'static, + { + if Self::can_block_on() { + return self.block_on(f); + } + + debug!("No live Tokio context (finalization); running block_on from a helper thread"); + let runtime = self.runtime.lock_or_recover().clone(); + Self::block_on_from_helper_thread(runtime, f) + } +} + +impl ForkSafeRuntime { + /// Runs `f` on a fresh thread (whose Tokio TLS was never touched) and blocks the + /// calling thread on its result via a plain OS channel, safe even with the calling + /// thread's TLS gone. Mirrors [`Self::shutdown_from_helper_thread`]. + fn block_on_from_helper_thread( + runtime: Option>, + f: F, + ) -> Result + where + F::Output: Send + 'static, + { + let (done_tx, done_rx) = std::sync::mpsc::channel(); + std::thread::Builder::new().spawn(move || { + let result = (|| -> Result { + let runtime = match runtime { + Some(runtime) => runtime, + None => Arc::new(Builder::new_current_thread().enable_all().build()?), + }; + Ok(runtime.block_on(f)) + })(); + let _ = done_tx.send(result); + })?; + done_rx + .recv() + .map_err(|_| io::Error::other("block_on helper thread died without a result"))? + } } #[cfg(test)] @@ -421,4 +575,131 @@ mod tests { "worker should not run or shut down after fork in child when restart_on_fork is false" ); } + + #[test] + fn test_shutdown_is_idempotent() { + // Calling shutdown() twice must not panic or error. The second call hits the + // None-guard (runtime already taken). This covers the same early-return path as + // the TLS-destroyed guard added for CPython atexit finalization ordering. + let shared_runtime = ForkSafeRuntime::new().unwrap(); + let (worker, receiver) = make_test_worker(); + + let _ = shared_runtime.spawn_worker(worker, true).unwrap(); + receiver + .recv_timeout(Duration::from_secs(1)) + .expect("worker did not run"); + + assert!(shared_runtime.shutdown(None).is_ok()); + assert!( + shared_runtime.shutdown(None).is_ok(), + "second shutdown must not panic" + ); + } + + #[test] + fn test_shutdown_from_helper_thread_flushes_workers() { + // Exercises the destroyed-TLS fallback mechanism directly (this test's thread has + // a perfectly fine Tokio context, but shutdown_from_helper_thread is the exact + // code path shutdown() delegates to when the calling thread's context is gone). + // Confirms the worker's shutdown() -- the "final flush" -- actually runs to + // completion on the helper thread, not just that teardown avoids a panic. + let shared_runtime = ForkSafeRuntime::new().unwrap(); + let (worker, receiver) = make_test_worker(); + + let _ = shared_runtime.spawn_worker(worker, true).unwrap(); + receiver + .recv_timeout(Duration::from_secs(1)) + .expect("worker did not run"); + + let runtime = shared_runtime.runtime.lock_or_recover().take().unwrap(); + let result = ForkSafeRuntime::shutdown_from_helper_thread( + runtime, + shared_runtime.workers.clone(), + Some(Duration::from_secs(1)), + ); + assert!(result.is_ok()); + + let mut last = receiver + .recv_timeout(Duration::from_secs(1)) + .expect("helper thread did not flush worker shutdown"); + while let Ok(v) = receiver.try_recv() { + last = v; + } + assert_eq!( + last, -1, + "worker shutdown (the flush) must actually run on the helper thread" + ); + } + + #[test] + fn test_drop_without_shutdown_detaches_cleanly() { + // Dropping a ForkSafeRuntime that still owns a live runtime + worker (i.e. + // shutdown() was never called) must not panic or hang. Drop routes through + // shutdown_background(), which detaches worker threads without blocking or + // entering the Tokio context — the structural terminal-teardown path that + // makes finalization (destroyed-TLS) drops safe. + let shared_runtime = ForkSafeRuntime::new().unwrap(); + let (worker, receiver) = make_test_worker(); + + let _ = shared_runtime.spawn_worker(worker, true).unwrap(); + receiver + .recv_timeout(Duration::from_secs(1)) + .expect("worker did not run"); + + // No explicit shutdown() — just drop. Must return promptly without panicking. + drop(shared_runtime); + } + + #[test] + fn test_lock_recovers_from_poison() { + // A panic while holding one of the runtime's locks must not cascade into + // subsequent lifecycle calls (the original PoisonError second-panic). After a + // poisoning panic, shutdown()/Drop must still succeed. + use std::panic::{catch_unwind, AssertUnwindSafe}; + + let shared_runtime = ForkSafeRuntime::new().unwrap(); + let (worker, _receiver) = make_test_worker(); + let _ = shared_runtime.spawn_worker(worker, true).unwrap(); + + let _ = catch_unwind(AssertUnwindSafe(|| { + let _guard = shared_runtime.workers.lock_or_recover(); + panic!("poison the workers mutex while holding it"); + })); + + // The mutex is now poisoned; lock_or_recover must still work rather than panic. + assert!( + shared_runtime.shutdown(None).is_ok(), + "shutdown must recover from a poisoned lock" + ); + } + + #[test] + fn test_block_on_send_runs_future_on_calling_thread() { + // Sanity check for the live-TLS path: block_on_send must still just drive `f` to + // completion on the calling thread and return its output, matching block_on. + let shared_runtime = ForkSafeRuntime::new().unwrap(); + let result = shared_runtime.block_on_send(async { 21 * 2 }).unwrap(); + assert_eq!(result, 42); + } + + #[test] + fn test_block_on_from_helper_thread_runs_future_and_returns_value() { + // Exercises the fallback directly; real TLS destruction can't be triggered + // deterministically from a unit test. Mirrors + // test_shutdown_from_helper_thread_flushes_workers. + let shared_runtime = ForkSafeRuntime::new().unwrap(); + let runtime = shared_runtime.runtime.lock_or_recover().clone(); + + let result = ForkSafeRuntime::block_on_from_helper_thread(runtime, async { 1 + 1 }); + assert_eq!(result.unwrap(), 2); + } + + #[test] + fn test_block_on_from_helper_thread_builds_runtime_when_none() { + // Same as above but for the (runtime already taken) case block_on's own None + // branch handles unsafely — the helper thread must build its own fresh runtime + // instead, since building on the *calling* thread is what panics once TLS is gone. + let result = ForkSafeRuntime::block_on_from_helper_thread(None, async { 40 + 2 }); + assert_eq!(result.unwrap(), 42); + } } diff --git a/libdd-shared-runtime/src/shared_runtime/mod.rs b/libdd-shared-runtime/src/shared_runtime/mod.rs index 852283179f..f99d2cddef 100644 --- a/libdd-shared-runtime/src/shared_runtime/mod.rs +++ b/libdd-shared-runtime/src/shared_runtime/mod.rs @@ -82,6 +82,20 @@ pub trait BlockingRuntime: SharedRuntime { /// /// Returns an [`io::Error`] if the executor cannot be accessed or constructed. fn block_on(&self, f: F) -> Result; + + /// Like [`Self::block_on`], but for callers that can guarantee `F`/its output are + /// `Send + 'static`, letting fork-safe impls hand `f` off to a helper thread when the + /// calling thread can't safely enter a Tokio context (see + /// [`ForkSafeRuntime::block_on_send`]). Defaults to forwarding to [`Self::block_on`]. + fn block_on_send( + &self, + f: F, + ) -> Result + where + F::Output: Send + 'static, + { + self.block_on(f) + } } /// Handle to a worker registered on a [`SharedRuntime`].