fix(native): shutdown in thread for uwsgi

Author: VianneyRuhlmann

ddtrace/internal/native/_native.pyi

@@ -180,6 +180,15 @@ class SharedRuntime:
     def shutdown(self, timeout_ms: Optional[int] = None) -> None:
         """Gracefully shut down the shared runtime.
 
+        Args:
+            timeout_ms: Maximum time in milliseconds to wait for shutdown.
+                If None, waits indefinitely.
+        """
+        ...
+    def shutdown_in_thread(self, timeout_ms: Optional[int] = None) -> None:
+        """Gracefully shut down the shared runtime.
+        The code is run in a separate thread to bypass a stale thread local storage.
+
         Args:
             timeout_ms: Maximum time in milliseconds to wait for shutdown.
                 If None, waits indefinitely.

ddtrace/internal/native_runtime.py

@@ -1,3 +1,4 @@
+import importlib.util
 import logging
 from typing import Optional
 
@@ -28,6 +29,7 @@ def __init__(self) -> None:
         forksafe.register_after_parent(self.after_fork_parent)
         forksafe.register(self.after_fork_child)
         atexit.register(self._atexit)
+        atexit.register_on_exit_signal(self._atexit)
 
     def _atexit(self) -> None:
         try:
@@ -43,6 +45,11 @@ def shutdown(self, timeout_ms: Optional[int] = None) -> None:
                 If None, waits indefinitely — only safe if all workers have
                 already been stopped (e.g. via TraceExporter.shutdown).
         """
+        using_uwsgi = importlib.util.find_spec("uwsgi") is not None
+        if not using_uwsgi:
+            super().shutdown(timeout_ms=_DEFAULT_SHUTDOWN_TIMEOUT_MS)
+        else:
+            super().shutdown_in_thread(timeout_ms=_DEFAULT_SHUTDOWN_TIMEOUT_MS)
         super().shutdown(timeout_ms=timeout_ms)
         atexit.unregister(self._atexit)
         forksafe.unregister_before_fork(self.before_fork)

src/native/data_pipeline/mod.rs

@@ -269,14 +269,6 @@ impl TraceExporterPy {
     }
 }
 
-impl Drop for TraceExporterPy {
-    fn drop(&mut self) {
-        if let Some(exporter) = self.inner.take() {
-            let _ = exporter.shutdown(Some(Duration::from_secs(3)));
-        }
-    }
-}
-
 #[pymodule]
 pub fn register_data_pipeline(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<TraceExporterBuilderPy>()?;

src/native/shared_runtime/mod.rs

@@ -1,6 +1,7 @@
 use libdd_shared_runtime::SharedRuntime;
 use pyo3::prelude::*;
 use std::sync::Arc;
+use std::thread;
 use std::time::Duration;
 
 mod exceptions;
@@ -43,12 +44,27 @@ impl SharedRuntimePy {
             .map_err(shared_runtime_error_to_pyerr)
     }
 
-    fn shutdown(&self, timeout_ms: Option<u64>) -> PyResult<()> {
+    fn shutdown(&self, py: Python<'_>, timeout_ms: Option<u64>) -> PyResult<()> {
         let timeout = timeout_ms.map(Duration::from_millis);
-        self.inner
-            .clone()
-            .shutdown(timeout)
-            .map_err(shared_runtime_error_to_pyerr)
+        let inner = self.inner.clone();
+        let result = py.detach(move || inner.shutdown(timeout));
+        result.map_err(shared_runtime_error_to_pyerr)
+    }
+
+    fn shutdown_in_thread(&self, py: Python<'_>, timeout_ms: Option<u64>) -> PyResult<()> {
+        let timeout = timeout_ms.map(Duration::from_millis);
+        let inner = self.inner.clone();
+        // Run shutdown in a dedicated OS thread so that Tokio's `block_on`
+        // does not execute on the Python main thread (which avoids conflicts
+        // with thread-local Tokio runtime state).  We release the GIL while
+        // waiting for the thread to join so that other Python threads can
+        // continue running during the (potentially slow) shutdown.
+        let result = py.detach(move || {
+            thread::spawn(move || inner.shutdown(timeout))
+                .join()
+                .unwrap_or_else(|_| panic!("SharedRuntime shutdown thread panicked"))
+        });
+        result.map_err(shared_runtime_error_to_pyerr)
     }
 
     fn debug(&self) -> String {