fix(python): don't crash when the interpreter exits during a result delivery

impit-node/test/basics.test.ts

@@ -13,14 +13,7 @@ import { Server as ProxyServer } from 'proxy-chain';
 function getHttpBinUrl(path: string, https?: boolean): string {
     https ??= true;
 
-    let url: URL;
-    if (process.env.APIFY_HTTPBIN_TOKEN) {
-        url = new URL(path, 'https://httpbin.apify.actor');
-        url.searchParams.set('token', process.env.APIFY_HTTPBIN_TOKEN);
-    } else {
-        url = new URL(path, 'https://httpbin.org');
-    }
-
+    const url = new URL(path, 'https://httpbin.apify.actor');
     url.protocol = https ? 'https:' : 'http:';
     return url.href;
 }

impit-python/clippy.toml

@@ -1 +1,5 @@
 too-many-arguments-threshold = 20
+
+disallowed-methods = [
+    { path = "pyo3_async_runtimes::tokio::future_into_py", reason = "use crate::exit_gate::future_into_py, it keeps result deliveries from racing interpreter shutdown" },
+]

impit-python/src/async_client.rs

@@ -27,7 +27,7 @@ impl AsyncClient {
         slf: Py<Self>,
         py: Python<'_>,
     ) -> Result<pyo3::Bound<'_, pyo3::PyAny>, pyo3::PyErr> {
-        pyo3_async_runtimes::tokio::future_into_py::<_, Py<AsyncClient>>(py, async { Ok(slf) })
+        crate::exit_gate::future_into_py::<_, Py<AsyncClient>>(py, async { Ok(slf) })
     }
 
     pub fn __aexit__<'python>(
@@ -37,7 +37,7 @@ impl AsyncClient {
         _traceback: &crate::Bound<'_, crate::PyAny>,
         py: Python<'python>,
     ) -> Result<pyo3::Bound<'python, pyo3::PyAny>, pyo3::PyErr> {
-        pyo3_async_runtimes::tokio::future_into_py::<_, ()>(py, async { Ok(()) })
+        crate::exit_gate::future_into_py::<_, ()>(py, async { Ok(()) })
     }
 
     #[new]
@@ -400,7 +400,7 @@ impl AsyncClient {
         let stream_value = stream.unwrap_or(false);
         let impit = Arc::clone(&self.impit);
 
-        pyo3_async_runtimes::tokio::future_into_py::<_, ImpitPyResponse>(py, async move {
+        crate::exit_gate::future_into_py::<_, ImpitPyResponse>(py, async move {
             let response = match method_str.to_lowercase().as_str() {
                 "get" => impit.get(url, Some(body), Some(options)).await,
                 "post" => impit.post(url, Some(body), Some(options)).await,

impit-python/src/exit_gate.rs

@@ -0,0 +1,138 @@
+//! Keeps tokio threads out of the interpreter once it starts shutting down.
+//!
+//! On CPython < 3.13.8, a non-Python thread that re-takes the GIL during finalization is killed
+//! with `pthread_exit`, which aborts the process when there are Rust frames on its stack. A result
+//! delivery from `future_into_py` releases the GIL inside `call_soon_threadsafe` (`socket.send`),
+//! so if the awaiting coroutine was the last thing the program did, the interpreter can start
+//! finalizing before the delivering thread gets the GIL back. `atexit` hooks run before the
+//! interpreter is marked as finalizing, so we close the gate there and wait for in-flight
+//! deliveries to leave Python.
+
+use std::cell::OnceCell;
+use std::future::Future;
+use std::pin::Pin;
+use std::sync::{Condvar, Mutex};
+use std::time::Duration;
+
+use pyo3::prelude::*;
+use pyo3_async_runtimes::generic::{ContextExt, Runtime};
+use pyo3_async_runtimes::TaskLocals;
+
+const DRAIN_TIMEOUT: Duration = Duration::from_secs(5);
+
+struct State {
+    closed: bool,
+    in_flight: usize,
+}
+
+static STATE: Mutex<State> = Mutex::new(State {
+    closed: false,
+    in_flight: 0,
+});
+static DRAINED: Condvar = Condvar::new();
+
+struct InFlight;
+
+impl InFlight {
+    fn enter() -> Option<Self> {
+        let mut state = STATE.lock().unwrap();
+        if state.closed {
+            return None;
+        }
+        state.in_flight += 1;
+        Some(InFlight)
+    }
+}
+
+impl Drop for InFlight {
+    fn drop(&mut self) {
+        let mut state = STATE.lock().unwrap();
+        state.in_flight -= 1;
+        if state.in_flight == 0 {
+            DRAINED.notify_all();
+        }
+    }
+}
+
+#[pyfunction]
+fn close(py: Python<'_>) {
+    py.detach(|| {
+        let mut state = STATE.lock().unwrap();
+        state.closed = true;
+        let _ = DRAINED.wait_timeout_while(state, DRAIN_TIMEOUT, |s| s.in_flight > 0);
+    });
+}
+
+pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> {
+    // Newer versions hang such threads instead of killing them.
+    if m.py().version_info() >= (3, 13, 8) {
+        return Ok(());
+    }
+    let close = wrap_pyfunction!(close, m)?;
+    m.py()
+        .import("atexit")?
+        .call_method1("register", (close,))?;
+    Ok(())
+}
+
+struct GatedTokio;
+
+tokio::task_local! {
+    static TASK_LOCALS: OnceCell<TaskLocals>;
+}
+
+impl Runtime for GatedTokio {
+    type JoinError = tokio::task::JoinError;
+    type JoinHandle = tokio::task::JoinHandle<()>;
+
+    fn spawn<F>(fut: F) -> Self::JoinHandle
+    where
+        F: Future<Output = ()> + Send + 'static,
+    {
+        pyo3_async_runtimes::tokio::get_runtime().spawn(fut)
+    }
+
+    // pyo3-async-runtimes delivers every result (and panic) to Python from a closure passed here.
+    // Once the gate is closed, the closure is dropped without attaching, and pyo3 defers the
+    // decrefs of the Python objects it holds.
+    fn spawn_blocking<F>(f: F) -> Self::JoinHandle
+    where
+        F: FnOnce() + Send + 'static,
+    {
+        pyo3_async_runtimes::tokio::get_runtime().spawn_blocking(move || {
+            if let Some(_in_flight) = InFlight::enter() {
+                f();
+            }
+        })
+    }
+}
+
+impl ContextExt for GatedTokio {
+    fn scope<F, R>(locals: TaskLocals, fut: F) -> Pin<Box<dyn Future<Output = R> + Send>>
+    where
+        F: Future<Output = R> + Send + 'static,
+    {
+        let cell = OnceCell::new();
+        cell.set(locals).unwrap();
+        Box::pin(TASK_LOCALS.scope(cell, fut))
+    }
+
+    fn get_task_locals() -> Option<TaskLocals> {
+        TASK_LOCALS
+            .try_with(|c| c.get().cloned())
+            .unwrap_or_default()
+    }
+}
+
+/// Drop-in replacement for `pyo3_async_runtimes::tokio::future_into_py` that goes through the
+/// exit gate.
+pub fn future_into_py<F, T>(py: Python<'_>, fut: F) -> PyResult<Bound<'_, PyAny>>
+where
+    F: Future<Output = PyResult<T>> + Send + 'static,
+    T: for<'py> IntoPyObject<'py> + Send + 'static,
+{
+    // Python code is still running, so the interpreter isn't finalizing yet. This reopens the gate
+    // for atexit hooks that run after ours, whose results would never arrive otherwise.
+    STATE.lock().unwrap().closed = false;
+    pyo3_async_runtimes::generic::future_into_py::<GatedTokio, _, T>(py, fut)
+}

impit-python/src/lib.rs

@@ -5,6 +5,7 @@ mod async_client;
 mod client;
 mod cookies;
 mod errors;
+mod exit_gate;
 mod fingerprint;
 mod request;
 mod response;
@@ -20,6 +21,8 @@ fn impit(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<AsyncClient>()?;
     m.add_class::<response::ImpitPyResponse>()?;
 
+    exit_gate::register(m)?;
+
     m.add("HTTPError", m.py().get_type::<errors::HTTPError>())?;
     m.add("RequestError", m.py().get_type::<errors::RequestError>())?;
     m.add(

impit-python/src/response.rs

@@ -109,9 +109,7 @@ impl PyResponseAsyncBytesIterator {
             self.content_returned = true;
             let content = PyBytes::new(py, &content).unbind();
             let future =
-                pyo3_async_runtimes::tokio::future_into_py::<_, Py<PyBytes>>(py, async move {
-                    Ok(content)
-                })?;
+                crate::exit_gate::future_into_py::<_, Py<PyBytes>>(py, async move { Ok(content) })?;
             return Ok(future);
         }
 
@@ -141,44 +139,43 @@ impl PyResponseAsyncBytesIterator {
 
         let parent_response = self.parent_response.as_ref().map(|p| p.clone_ref(py));
 
-        let future =
-            pyo3_async_runtimes::tokio::future_into_py::<_, Py<PyBytes>>(py, async move {
-                let chunk_result = {
-                    let mut stream_guard = stream_arc.lock().await;
-                    if let Some(stream) = stream_guard.as_mut() {
-                        stream.next().await
-                    } else {
-                        None
-                    }
-                };
-                match chunk_result {
-                    Some(Ok(chunk)) => Ok(Python::attach(|py| PyBytes::new(py, &chunk).unbind())),
-                    Some(Err(e)) => {
-                        if let Some(parent) = parent_response {
-                            Python::attach(|py| {
-                                if let Ok(mut parent_ref) = parent.try_borrow_mut(py) {
-                                    parent_ref.inner_state = InnerResponseState::StreamingClosed;
-                                    parent_ref.is_stream_consumed = true;
-                                    parent_ref.is_closed = true;
-                                }
-                            });
-                        }
-                        Err(ImpitPyError(ImpitError::from(e, None)).into())
+        let future = crate::exit_gate::future_into_py::<_, Py<PyBytes>>(py, async move {
+            let chunk_result = {
+                let mut stream_guard = stream_arc.lock().await;
+                if let Some(stream) = stream_guard.as_mut() {
+                    stream.next().await
+                } else {
+                    None
+                }
+            };
+            match chunk_result {
+                Some(Ok(chunk)) => Ok(Python::attach(|py| PyBytes::new(py, &chunk).unbind())),
+                Some(Err(e)) => {
+                    if let Some(parent) = parent_response {
+                        Python::attach(|py| {
+                            if let Ok(mut parent_ref) = parent.try_borrow_mut(py) {
+                                parent_ref.inner_state = InnerResponseState::StreamingClosed;
+                                parent_ref.is_stream_consumed = true;
+                                parent_ref.is_closed = true;
+                            }
+                        });
                     }
-                    None => {
-                        if let Some(parent) = parent_response {
-                            Python::attach(|py| {
-                                if let Ok(mut parent_ref) = parent.try_borrow_mut(py) {
-                                    parent_ref.inner_state = InnerResponseState::StreamingClosed;
-                                    parent_ref.is_stream_consumed = true;
-                                    parent_ref.is_closed = true;
-                                }
-                            });
-                        }
-                        Err(pyo3::exceptions::PyStopAsyncIteration::new_err(""))
+                    Err(ImpitPyError(ImpitError::from(e, None)).into())
+                }
+                None => {
+                    if let Some(parent) = parent_response {
+                        Python::attach(|py| {
+                            if let Ok(mut parent_ref) = parent.try_borrow_mut(py) {
+                                parent_ref.inner_state = InnerResponseState::StreamingClosed;
+                                parent_ref.is_stream_consumed = true;
+                                parent_ref.is_closed = true;
+                            }
+                        });
                     }
+                    Err(pyo3::exceptions::PyStopAsyncIteration::new_err(""))
                 }
-            })?;
+            }
+        })?;
 
         Ok(future)
     }
@@ -295,7 +292,7 @@ impl ImpitPyResponse {
     }
 
     fn aclose(slf: Py<Self>, py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+        crate::exit_gate::future_into_py(py, async move {
             Self::close_async_impl(slf)?;
             Ok(())
         })
@@ -390,7 +387,7 @@ impl ImpitPyResponse {
     }
 
     fn aread(slf: Py<Self>, py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+        crate::exit_gate::future_into_py(py, async move {
             let (state, response_option, content_option, is_stream_consumed) =
                 Python::attach(|py| {
                     let mut slf_ref = slf.borrow_mut(py);

impit-python/test/httpbin.py

@@ -1,20 +1,12 @@
 from __future__ import annotations
 
-import os
 import urllib.parse
 
 
 def get_httpbin_url(path: str, *, query: dict[str, list[str] | str] | None = None, https: bool = True) -> str:
-    query = query or {}
-    url = None
-    if os.environ.get('APIFY_HTTPBIN_TOKEN'):
-        url = urllib.parse.urlparse('https://httpbin.apify.actor')
-        query['token'] = os.environ['APIFY_HTTPBIN_TOKEN']
+    url = urllib.parse.urlparse('https://httpbin.apify.actor')
+    if query:
         url = url._replace(query=urllib.parse.urlencode(query, doseq=True))
-    else:
-        url = urllib.parse.urlparse('https://httpbin.org')
-        if query:
-            url = url._replace(query=urllib.parse.urlencode(query, doseq=True))
 
     scheme = 'https' if https else 'http'
     url = url._replace(scheme=scheme)