fix(shared-runtime): make ForkSafeRuntime teardown structurally TLS/poison-safe

libdd-common/src/lib.rs

@@ -79,6 +79,17 @@ pub mod unix_utils;
 /// ```
 pub trait MutexExt<T> {
     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<T>`'s impl below is the
+    /// one that actually recovers.
+    fn lock_or_recover(&self) -> MutexGuard<'_, T> {
+        self.lock_or_panic()
+    }
 }
 
 impl<T> MutexExt<T> for Mutex<T> {
@@ -88,6 +99,11 @@ impl<T> MutexExt<T> for Mutex<T> {
         #[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

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<std::time::Duration>) -> 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<Duration>) -> 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, <Self as SharedRuntime>::shutdown_async(self))
-                            .await
-                    }) {
-                        Ok(()) => Ok(()),
-                        Err(_) => Err(SharedRuntimeError::ShutdownTimedOut(timeout)),
-                    }
-                } else {
-                    runtime.block_on(<Self as SharedRuntime>::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<Mutex<Vec<WorkerEntry>>>,
+        timeout: Option<Duration>,
+    ) -> 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<Runtime>,
+        workers: Arc<Mutex<Vec<WorkerEntry>>>,
+        timeout: Option<Duration>,
+    ) -> 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<Vec<WorkerEntry>>,
@@ -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, SharedRuntimeError> {
         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<Vec<WorkerEntry>>) {
+    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 {
@@ -421,4 +528,101 @@ 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"
+        );
+    }
 }