fix(client): catch PyO3 attach panics during interpreter shutdown

Cargo.toml

@@ -92,5 +92,4 @@ debug = false
 incremental = false
 lto = "fat"
 opt-level = 3
-panic = "abort"
 strip = true

src/client/body/stream.rs

@@ -15,7 +15,12 @@ use pyo3::{
 };
 use tokio::{sync::Mutex, task::JoinHandle};
 
-use crate::{buffer::PyBuffer, client::nogil::NoGIL, error::Error, header::HeaderMap};
+use crate::{
+    buffer::PyBuffer,
+    client::nogil::NoGIL,
+    error::{self, Error},
+    header::HeaderMap,
+};
 
 type Pending = Option<JoinHandle<Option<PyResult<PyBytesLike>>>>;
 
@@ -224,19 +229,21 @@ impl Stream for PyStream {
                     PyStreamSource::Sync(ref ob) => {
                         let ob = ob.clone();
                         runtime.spawn_blocking(move || {
-                            Python::attach(|py| {
+                            error::attach(|py| {
                                 ob.call_method0(py, intern!(py, "__next__"))
                                     .ok()
                                     .map(|ob| ob.extract(py))
                             })
+                            .unwrap_or_else(|err| Some(Err(err.into())))
                         })
                     }
                     PyStreamSource::Async(ref stream) => {
                         let stream = stream.clone();
                         runtime.spawn(async move {
                             let ob = stream.lock().await.next().await;
                             tokio::task::spawn_blocking(move || {
-                                Python::attach(|py| ob.map(|ob| ob.extract(py)))
+                                error::attach(|py| ob.map(|ob| ob.extract(py)))
+                                    .unwrap_or_else(|err| Some(Err(err.into())))
                             })
                             .await
                             .ok()?

src/client/nogil.rs

@@ -1,7 +1,8 @@
 use std::{
     future::Future,
     pin::Pin,
-    task::{Context, Poll},
+    sync::{Arc, Once},
+    task::{Context, Poll, Wake, Waker},
 };
 
 use pin_project_lite::pin_project;
@@ -12,6 +13,8 @@ use pyo3::{
 };
 use tokio_util::{sync::CancellationToken, task::AbortOnDropHandle};
 
+use crate::error;
+
 pin_project! {
     /// A future that allows Python threads to run while it is being polled or executed.
     /// It also handles cancellation and spawns the task in tokio runtime.
@@ -77,13 +80,36 @@ where
             })));
         }
 
-        let waker = cx.waker();
+        let waker = Waker::from(Arc::new(GuardedWaker(cx.waker().clone())));
         Python::attach(|py| {
-            py.detach(|| match this.handle.poll(&mut Context::from_waker(waker)) {
-                Poll::Ready(Ok(result)) => Poll::Ready(result),
-                Poll::Ready(Err(e)) => Poll::Ready(Err(PyRuntimeError::new_err(e.to_string()))),
-                Poll::Pending => Poll::Pending,
+            py.detach(|| {
+                let mut cx = Context::from_waker(&waker);
+                match this.handle.poll(&mut cx) {
+                    Poll::Ready(Ok(result)) => Poll::Ready(result),
+                    Poll::Ready(Err(e)) => Poll::Ready(Err(PyRuntimeError::new_err(e.to_string()))),
+                    Poll::Pending => Poll::Pending,
+                }
             })
         })
     }
 }
+
+/// Wakes the Python coroutine from Tokio threads.
+///
+/// PyO3's coroutine waker calls `Python::attach`, which panics once the interpreter has
+/// shut down. Waking inside [`error::attach`] lets it reuse that attachment instead; when
+/// Python is gone the wake is dropped and reported once, as no coroutine is left to resume.
+struct GuardedWaker(Waker);
+
+impl Wake for GuardedWaker {
+    fn wake(self: Arc<Self>) {
+        self.wake_by_ref();
+    }
+
+    fn wake_by_ref(self: &Arc<Self>) {
+        if let Err(err) = error::attach(|_| self.0.wake_by_ref()) {
+            static REPORTED: Once = Once::new();
+            REPORTED.call_once(|| eprintln!("wreq: failed to wake a Python coroutine: {err}"));
+        }
+    }
+}

src/error.rs

@@ -1,5 +1,11 @@
+use std::{
+    any::Any,
+    fmt,
+    panic::{self, AssertUnwindSafe},
+};
+
 use pyo3::{
-    PyErr, create_exception,
+    PyErr, Python, create_exception,
     exceptions::{PyException, PyRuntimeError, PyStopAsyncIteration, PyStopIteration},
 };
 use wreq::header;
@@ -57,6 +63,9 @@ macro_rules! wrap_error {
     };
 }
 
+const INTERPRETER_UNAVAILABLE_MSG: &str =
+    "The Python interpreter is not available (not initialized or shutting down)";
+
 /// Unified error enum
 #[derive(Debug)]
 pub enum Error {
@@ -74,6 +83,8 @@ pub enum Error {
     Json(serde_json::Error),
     Form(serde_urlencoded::ser::Error),
     Library(wreq::Error),
+    InterpreterUnavailable,
+    Panic(Box<str>),
 }
 
 impl From<Error> for PyErr {
@@ -100,6 +111,8 @@ impl From<Error> for PyErr {
             Error::Dns(err) => BuilderError::new_err(format!("DNS resolver error: {err:?}")),
             Error::Json(err) => PyRuntimeError::new_err(format!("JSON error: {err:?}")),
             Error::Form(err) => PyRuntimeError::new_err(format!("Form error: {err:?}")),
+            Error::InterpreterUnavailable => PyRuntimeError::new_err(INTERPRETER_UNAVAILABLE_MSG),
+            Error::Panic(msg) => RustPanic::new_err(msg.into_string()),
             Error::Library(err) => wrap_error!(err,
                 is_body => BodyError,
                 is_tls => TlsError,
@@ -118,6 +131,41 @@ impl From<Error> for PyErr {
     }
 }
 
+impl fmt::Display for Error {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            Error::InterpreterUnavailable => f.write_str(INTERPRETER_UNAVAILABLE_MSG),
+            Error::Panic(msg) => f.write_str(msg),
+            err => fmt::Debug::fmt(err, f),
+        }
+    }
+}
+
+impl Error {
+    fn panic(payload: Box<dyn Any + Send>) -> Self {
+        let msg = payload
+            .downcast_ref::<&str>()
+            .copied()
+            .or_else(|| payload.downcast_ref::<String>().map(String::as_str))
+            .unwrap_or("unknown panic");
+        Error::Panic(format!("Rust panic: {msg}").into_boxed_str())
+    }
+}
+
+/// Attaches to Python from a non-Python thread.
+///
+/// `Python::attach` panics once the interpreter has shut down (see discussions/305), so
+/// this uses `Python::try_attach` and returns [`Error::InterpreterUnavailable`] instead.
+/// Any other panic in `f` is caught as [`Error::Panic`] to keep Tokio threads alive.
+pub fn attach<F, R>(f: F) -> Result<R, Error>
+where
+    F: for<'py> FnOnce(Python<'py>) -> R,
+{
+    panic::catch_unwind(AssertUnwindSafe(|| Python::try_attach(f)))
+        .map_err(Error::panic)?
+        .ok_or(Error::InterpreterUnavailable)
+}
+
 impl From<header::InvalidHeaderName> for Error {
     fn from(err: header::InvalidHeaderName) -> Self {
         Error::InvalidHeaderName(err)

src/redirect.rs

@@ -2,7 +2,7 @@ use std::{fmt::Display, sync::Arc};
 
 use pyo3::prelude::*;
 
-use crate::{header::HeaderMap, http::StatusCode};
+use crate::{error, header::HeaderMap, http::StatusCode};
 
 /// Represents the redirect policy for HTTP requests.
 #[derive(Clone)]
@@ -103,13 +103,14 @@ impl Policy {
             attempt.pending(|attempt| async move {
                 let args = Attempt::from(&attempt);
                 let kind = tokio::task::spawn_blocking(move || {
-                    Python::attach(|py| {
+                    error::attach(|py| {
                         callback
                             .call1(py, (args,))
                             .and_then(|result| result.extract::<Action>(py).map_err(PyErr::from))
                             .map(|action| action.kind)
                             .unwrap_or_else(|err| ActionKind::Error(err.to_string()))
                     })
+                    .unwrap_or_else(|err| ActionKind::Error(err.to_string()))
                 })
                 .await;
 

tests/shutdown_test.py

@@ -0,0 +1,66 @@
+import platform
+import subprocess
+import sys
+
+import pytest
+
+# https://github.com/0x676e67/wreq-python/discussions/305: a Tokio worker that attaches
+# to Python during interpreter shutdown used to hit PyO3's not-initialized panic. wreq
+# now detects the unavailable interpreter and reports a wreq error without panicking.
+
+# A request is left pending on a listener that never accepts, then an uncaught error
+# shuts the interpreter down. Module teardown runs after Py_IsInitialized() drops to 0;
+# `HoldTeardown.__del__` closes the listener there, failing the request so the coroutine
+# waker attaches from a Tokio worker, and stalls teardown until that happens.
+SCRIPT = """
+import asyncio
+import socket
+import time
+
+import wreq
+
+listener = socket.socket()
+listener.bind(("127.0.0.1", 0))
+listener.listen(8)
+url = f"http://127.0.0.1:{listener.getsockname()[1]}/"
+
+
+class HoldTeardown:
+    def __init__(self, listener, task):
+        self.listener = listener
+        self.task = task
+
+    def __del__(self):
+        self.listener.close()
+        time.sleep(1)
+
+
+loop = asyncio.new_event_loop()
+task = loop.create_task(wreq.Client().get(url))
+loop.run_until_complete(asyncio.sleep(0.05))
+hold = HoldTeardown(listener, task)
+del listener, task
+raise RuntimeError("uncaught error while a request is in flight")
+"""
+
+# PyPy doesn't guarantee `__del__` runs during interpreter exit, so the trigger may not fire.
+pytestmark = pytest.mark.skipif(
+    platform.python_implementation() != "CPython",
+    reason="relies on CPython running __del__ during module teardown",
+)
+
+
+def test_shutdown_wake_reports_error_without_panic():
+    proc = subprocess.run(
+        [sys.executable, "-c", SCRIPT],
+        capture_output=True,
+        text=True,
+        timeout=60,
+    )
+    # Exits through the uncaught RuntimeError, not an abort.
+    assert proc.returncode == 1, proc.stderr
+    assert "panicked" not in proc.stderr
+    assert (
+        "wreq: failed to wake a Python coroutine: "
+        "The Python interpreter is not available" in proc.stderr
+    )