diff --git a/impit-node/test/basics.test.ts b/impit-node/test/basics.test.ts index aa9cac99..2794cdf7 100644 --- a/impit-node/test/basics.test.ts +++ b/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; } diff --git a/impit-python/clippy.toml b/impit-python/clippy.toml index 0d4e02f0..9421bd94 100644 --- a/impit-python/clippy.toml +++ b/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" }, +] diff --git a/impit-python/src/async_client.rs b/impit-python/src/async_client.rs index 9a5a5004..990936b2 100644 --- a/impit-python/src/async_client.rs +++ b/impit-python/src/async_client.rs @@ -27,7 +27,7 @@ impl AsyncClient { slf: Py, py: Python<'_>, ) -> Result, pyo3::PyErr> { - pyo3_async_runtimes::tokio::future_into_py::<_, Py>(py, async { Ok(slf) }) + crate::exit_gate::future_into_py::<_, Py>(py, async { Ok(slf) }) } pub fn __aexit__<'python>( @@ -37,7 +37,7 @@ impl AsyncClient { _traceback: &crate::Bound<'_, crate::PyAny>, py: Python<'python>, ) -> Result, 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, diff --git a/impit-python/src/exit_gate.rs b/impit-python/src/exit_gate.rs new file mode 100644 index 00000000..af285c94 --- /dev/null +++ b/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 = Mutex::new(State { + closed: false, + in_flight: 0, +}); +static DRAINED: Condvar = Condvar::new(); + +struct InFlight; + +impl InFlight { + fn enter() -> Option { + 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; +} + +impl Runtime for GatedTokio { + type JoinError = tokio::task::JoinError; + type JoinHandle = tokio::task::JoinHandle<()>; + + fn spawn(fut: F) -> Self::JoinHandle + where + F: Future + 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) -> 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(locals: TaskLocals, fut: F) -> Pin + Send>> + where + F: Future + Send + 'static, + { + let cell = OnceCell::new(); + cell.set(locals).unwrap(); + Box::pin(TASK_LOCALS.scope(cell, fut)) + } + + fn get_task_locals() -> Option { + 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(py: Python<'_>, fut: F) -> PyResult> +where + F: Future> + 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::(py, fut) +} diff --git a/impit-python/src/lib.rs b/impit-python/src/lib.rs index e6b89936..2617c6ba 100644 --- a/impit-python/src/lib.rs +++ b/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::()?; m.add_class::()?; + exit_gate::register(m)?; + m.add("HTTPError", m.py().get_type::())?; m.add("RequestError", m.py().get_type::())?; m.add( diff --git a/impit-python/src/response.rs b/impit-python/src/response.rs index 49e833dc..911b3d69 100644 --- a/impit-python/src/response.rs +++ b/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>(py, async move { - Ok(content) - })?; + crate::exit_gate::future_into_py::<_, Py>(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>(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>(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, py: Python<'_>) -> PyResult> { - 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, py: Python<'_>) -> PyResult> { - 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); diff --git a/impit-python/test/httpbin.py b/impit-python/test/httpbin.py index aeb20499..780d30b7 100644 --- a/impit-python/test/httpbin.py +++ b/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)