fix: Replace atexit with async-compatible join_pending_shutdown

Remove atexit handler mechanism and replace with explicit join function:
- Add join_pending_shutdown(py) that blocks with GIL released
- Remove PENDING_CLEANUPS, ATEXIT_REGISTERED, and related code
- Add PENDING_SHUTDOWN to store single thread handle for joining

The new approach allows Python code to await the shutdown via
asyncio.to_thread(), ensuring cleanup completes within the async
context before the event loop closes.

Author: achimnol

src/tokio.rs

@@ -14,7 +14,6 @@
 //! ```
 
 use std::cell::OnceCell;
-use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::{Arc, Mutex, OnceLock};
 use std::thread::{self, JoinHandle};
 use std::{future::Future, pin::Pin};
@@ -27,7 +26,6 @@ use ::tokio::{
 use once_cell::sync::Lazy;
 use parking_lot::RwLock;
 use pyo3::prelude::*;
-use pyo3::types::PyCFunction;
 
 use crate::{
     generic::{self, ContextExt, LocalContextExt, Runtime as GenericRuntime, SpawnLocalExt},
@@ -184,70 +182,9 @@ static TOKIO_BUILDER: Lazy<Mutex<Builder>> = Lazy::new(|| Mutex::new(multi_threa
 static RUNTIME_WRAPPER: RwLock<Option<Arc<RuntimeWrapper>>> = RwLock::new(None);
 static TOKIO_RUNTIME: OnceLock<&'static Runtime> = OnceLock::new();
 
-/// Pending runtime threads that need cleanup at process exit.
-/// Used by request_shutdown_background to defer thread joining until Python's atexit.
-static PENDING_CLEANUPS: Lazy<Mutex<Vec<JoinHandle<()>>>> = Lazy::new(|| Mutex::new(Vec::new()));
-/// Flag to track if we've registered the atexit handler.
-static ATEXIT_REGISTERED: AtomicBool = AtomicBool::new(false);
-
-/// Join all pending runtime threads.
-///
-/// This function is called from Python's atexit to ensure all runtime threads
-/// are properly joined before Python finalizes. It should be registered via
-/// [`register_atexit_cleanup`] during module initialization.
-fn join_pending_cleanups() {
-    let threads: Vec<JoinHandle<()>> = {
-        let mut guard = PENDING_CLEANUPS.lock().unwrap();
-        std::mem::take(&mut *guard)
-    };
-
-    for thread in threads {
-        let _ = thread.join();
-    }
-}
-
-/// Register the atexit cleanup handler for this module.
-///
-/// Call this function during your `#[pymodule]` initialization to ensure
-/// tokio runtime threads are properly cleaned up before Python finalizes.
-///
-/// # Example
-///
-/// ```rust,ignore
-/// use pyo3::prelude::*;
-///
-/// #[pymodule]
-/// fn my_module(py: Python, m: &Bound<PyModule>) -> PyResult<()> {
-///     // Register atexit cleanup for tokio runtime
-///     pyo3_async_runtimes::tokio::register_atexit_cleanup(py)?;
-///
-///     // ... rest of module init
-///     Ok(())
-/// }
-/// ```
-pub fn register_atexit_cleanup(py: Python<'_>) -> PyResult<()> {
-    if ATEXIT_REGISTERED
-        .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
-        .is_ok()
-    {
-        // Create a Python function that calls our Rust cleanup
-        let cleanup_fn = PyCFunction::new_closure(
-            py,
-            None,
-            None,
-            |_args: &Bound<'_, pyo3::types::PyTuple>,
-             _kwargs: Option<&Bound<'_, pyo3::types::PyDict>>| {
-                join_pending_cleanups();
-                Ok::<(), PyErr>(())
-            },
-        )?;
-
-        // Register with Python's atexit module
-        let atexit = py.import("atexit")?;
-        atexit.call_method1("register", (cleanup_fn,))?;
-    }
-    Ok(())
-}
+/// Pending runtime thread that needs cleanup.
+/// Stored when request_shutdown_background is called, joined by join_pending_shutdown.
+static PENDING_SHUTDOWN: Mutex<Option<JoinHandle<()>>> = Mutex::new(None);
 
 /// Get or create the runtime wrapper.
 ///
@@ -639,16 +576,12 @@ pub fn request_shutdown_background(timeout_ms: u64) -> bool {
             let _ = sender.send(timeout_ms);
         }
 
-        // Take the thread handle and store it for cleanup at process exit
+        // Store the thread handle for later joining via join_pending_shutdown()
         if let Some(thread) = wrapper.runtime_thread.lock().unwrap().take() {
-            // Store the thread handle for later cleanup via atexit.
-            // This ensures the thread is properly joined before Python finalizes,
-            // preventing SIGSEGV from tokio threads running during interpreter shutdown.
-            PENDING_CLEANUPS.lock().unwrap().push(thread);
+            *PENDING_SHUTDOWN.lock().unwrap() = Some(thread);
         }
 
         // The wrapper is dropped here, releasing all Arc references.
-        // The thread will be joined later via join_pending_cleanups() called from atexit.
 
         true
     } else {
@@ -657,6 +590,50 @@ pub fn request_shutdown_background(timeout_ms: u64) -> bool {
     }
 }
 
+/// Join any pending runtime shutdown thread.
+///
+/// This function should be called from Python (via `asyncio.to_thread`) after
+/// `request_shutdown_background` signals shutdown. It blocks until the runtime
+/// thread completes, with the GIL released to allow other Python threads to run.
+///
+/// This is the key to proper cleanup: by calling this from Python's async context
+/// (via `asyncio.to_thread`), we ensure the runtime thread is fully terminated
+/// before Python's event loop closes.
+///
+/// # Arguments
+///
+/// * `py` - Python GIL token (will be released during blocking wait)
+///
+/// # Returns
+///
+/// Returns `true` if a thread was joined, `false` if no pending shutdown.
+///
+/// # Example
+///
+/// ```python
+/// import asyncio
+/// from etcd_client import _join_pending_shutdown
+///
+/// async def cleanup():
+///     # ... signal shutdown ...
+///     # Wait for runtime to fully terminate
+///     await asyncio.to_thread(_join_pending_shutdown)
+/// ```
+pub fn join_pending_shutdown(py: Python<'_>) -> bool {
+    let thread = PENDING_SHUTDOWN.lock().unwrap().take();
+
+    if let Some(thread) = thread {
+        // Release GIL while blocking on thread join
+        #[allow(deprecated)] // py.allow_threads is deprecated but detach doesn't fit our use case
+        py.allow_threads(|| {
+            let _ = thread.join();
+        });
+        true
+    } else {
+        false
+    }
+}
+
 // ============================================================================
 // Public API - Future Conversion
 // ============================================================================