perf(pyo3): cache async future completion helpers

src/py/client/async_requests/spawn.rs

@@ -3,7 +3,7 @@ use super::callback::PythonFutureCancellation;
 use super::registry::AsyncRequestRegistry;
 use crate::core::metrics::Metrics;
 use crate::py::client::acquire::AcquireGate;
-use crate::py::client::future::complete_python_future;
+use crate::py::client::future::{complete_python_future, PythonFutureSetters};
 use crate::py::client::transport::{send_request, TransportClients, TransportRequest};
 use pyo3::prelude::*;
 use pyo3::types::PyAny;
@@ -15,6 +15,7 @@ pub struct AsyncRequestSpawn {
     pub clients: TransportClients,
     pub metrics: Arc<Metrics>,
     pub pool_timeout: f64,
+    pub future_setters: PythonFutureSetters,
     pub request: TransportRequest,
 }
 
@@ -29,6 +30,7 @@ pub fn spawn_async_request(
         clients,
         metrics,
         pool_timeout,
+        future_setters,
         request,
     } = spawn;
     let loop_ = py
@@ -41,6 +43,7 @@ pub fn spawn_async_request(
 
     let task_loop = loop_.clone_ref(py);
     let task_future = future.clone_ref(py);
+    let task_future_setters = future_setters.clone_ref(py);
     let task_registry = registry.clone();
     let task_metrics = Arc::clone(&metrics);
     let task_completion = completion.clone();
@@ -58,7 +61,7 @@ pub fn spawn_async_request(
         if task_completion.finish() {
             task_registry.remove(request_id);
             task_metrics.request_finished(result.is_err());
-            complete_python_future(&task_loop, &task_future, result);
+            complete_python_future(&task_loop, &task_future, &task_future_setters, result);
         }
     });
 

src/py/client/async_requests/stream_spawn.rs

@@ -5,7 +5,7 @@ use crate::core::metrics::Metrics;
 use crate::errors::FogHttpError;
 use crate::messages::STREAM_REQUEST_TASK_START_FAILED;
 use crate::py::client::acquire::AcquireGate;
-use crate::py::client::future::complete_python_stream_future;
+use crate::py::client::future::{complete_python_stream_future, PythonFutureSetters};
 use crate::py::client::streams::StreamRegistry;
 use crate::py::client::transport::{send_stream_request, TransportClients, TransportRequest};
 use pyo3::prelude::*;
@@ -20,6 +20,7 @@ pub struct AsyncStreamRequestSpawn {
     pub metrics: Arc<Metrics>,
     pub active_streams: StreamRegistry,
     pub pool_timeout: f64,
+    pub future_setters: PythonFutureSetters,
     pub request: TransportRequest,
 }
 
@@ -35,6 +36,7 @@ pub fn spawn_async_stream_request(
         metrics,
         active_streams,
         pool_timeout,
+        future_setters,
         request,
     } = spawn;
     let loop_ = py
@@ -48,6 +50,7 @@ pub fn spawn_async_stream_request(
 
     let task_loop = loop_.clone_ref(py);
     let task_future = future.clone_ref(py);
+    let task_future_setters = future_setters.clone_ref(py);
     let task_registry = registry.clone();
     let task_metrics = Arc::clone(&metrics);
     let task_completion = completion.clone();
@@ -65,6 +68,7 @@ pub fn spawn_async_stream_request(
             active_streams,
             runtime_handle,
             pool_timeout,
+            future_setters,
             request,
             task_completion.clone(),
         )
@@ -73,7 +77,7 @@ pub fn spawn_async_stream_request(
         if result.is_err() && task_completion.finish() {
             task_metrics.request_finished(true);
         }
-        complete_python_stream_future(&task_loop, &task_future, result);
+        complete_python_stream_future(&task_loop, &task_future, &task_future_setters, result);
     });
 
     registry.insert(

src/py/client/future.rs

@@ -37,14 +37,21 @@ pub fn cancel_python_future(loop_: &Py<PyAny>, future: &Py<PyAny>) {
 pub fn complete_python_future(
     loop_: &Py<PyAny>,
     future: &Py<PyAny>,
+    setters: &PythonFutureSetters,
     result: PyResult<RawResponse>,
 ) {
     Python::attach(|py| {
         let call_result: PyResult<()> = match result {
-            Ok(response) => schedule_set_result(py, loop_, future, response),
+            Ok(response) => schedule_set_result(py, loop_, future, setters, response),
             Err(err) => {
                 let exception = err.into_value(py);
-                schedule_set_exception(py, loop_, future, exception)
+                schedule_set_exception_with_helper(
+                    py,
+                    loop_,
+                    future,
+                    &setters.set_exception_if_pending,
+                    exception,
+                )
             }
         };
 
@@ -57,14 +64,21 @@ pub fn complete_python_future(
 pub fn complete_python_stream_future(
     loop_: &Py<PyAny>,
     future: &Py<PyAny>,
+    setters: &PythonFutureSetters,
     result: PyResult<RawStreamResponse>,
 ) {
     Python::attach(|py| {
         let call_result: PyResult<()> = match result {
-            Ok(response) => schedule_set_stream_result(py, loop_, future, response),
+            Ok(response) => schedule_set_stream_result(py, loop_, future, setters, response),
             Err(err) => {
                 let exception = err.into_value(py);
-                schedule_set_exception(py, loop_, future, exception)
+                schedule_set_exception_with_helper(
+                    py,
+                    loop_,
+                    future,
+                    &setters.set_exception_if_pending,
+                    exception,
+                )
             }
         };
 
@@ -146,21 +160,6 @@ fn schedule_set_none_result(
     Ok(())
 }
 
-fn schedule_set_exception(
-    py: Python<'_>,
-    loop_: &Py<PyAny>,
-    future: &Py<PyAny>,
-    exception: Py<PyBaseException>,
-) -> PyResult<()> {
-    let helper = py
-        .import("foghttp._client.asyncio_futures")?
-        .getattr("set_exception_if_pending")?;
-    loop_
-        .bind(py)
-        .call_method1("call_soon_threadsafe", (helper, future.bind(py), exception))?;
-    Ok(())
-}
-
 fn schedule_set_exception_with_helper(
     py: Python<'_>,
     loop_: &Py<PyAny>,
@@ -179,30 +178,36 @@ fn schedule_set_stream_result(
     py: Python<'_>,
     loop_: &Py<PyAny>,
     future: &Py<PyAny>,
+    setters: &PythonFutureSetters,
     response: RawStreamResponse,
 ) -> PyResult<()> {
-    let helper = py
-        .import("foghttp._client.asyncio_futures")?
-        .getattr("set_result_if_pending")?;
     let response = Py::new(py, response)?;
-    loop_
-        .bind(py)
-        .call_method1("call_soon_threadsafe", (helper, future.bind(py), response))?;
+    loop_.bind(py).call_method1(
+        "call_soon_threadsafe",
+        (
+            setters.set_result_if_pending.bind(py),
+            future.bind(py),
+            response,
+        ),
+    )?;
     Ok(())
 }
 
 fn schedule_set_result(
     py: Python<'_>,
     loop_: &Py<PyAny>,
     future: &Py<PyAny>,
+    setters: &PythonFutureSetters,
     response: RawResponse,
 ) -> PyResult<()> {
-    let helper = py
-        .import("foghttp._client.asyncio_futures")?
-        .getattr("set_result_if_pending")?;
     let response = Py::new(py, response)?;
-    loop_
-        .bind(py)
-        .call_method1("call_soon_threadsafe", (helper, future.bind(py), response))?;
+    loop_.bind(py).call_method1(
+        "call_soon_threadsafe",
+        (
+            setters.set_result_if_pending.bind(py),
+            future.bind(py),
+            response,
+        ),
+    )?;
     Ok(())
 }

src/py/client/mod.rs

@@ -24,6 +24,7 @@ use crate::py::client::async_requests::{
     spawn_async_request, spawn_async_stream_request, AsyncRequestRegistry, AsyncRequestSpawn,
     AsyncStreamRequestSpawn, RequestCompletion,
 };
+use crate::py::client::future::PythonFutureSetters;
 use crate::py::client::options::{
     validate_numeric_client_options, validate_request_timeouts, NumericClientOptions,
 };
@@ -50,6 +51,7 @@ pub struct RawClient {
     metrics: Arc<Metrics>,
     active_async_requests: AsyncRequestRegistry,
     active_streams: StreamRegistry,
+    future_setters: PythonFutureSetters,
     max_response_body_size: Option<usize>,
     buffered_body_budget: BufferedBodyBudget,
     follow_redirects: bool,
@@ -91,6 +93,7 @@ impl RawClient {
         reason = "PyO3 constructor mirrors Python client options before transport grouping."
     )]
     fn new(
+        py: Python<'_>,
         max_active_requests: usize,
         max_active_requests_per_origin: Option<usize>,
         max_connections: Option<usize>,
@@ -193,6 +196,7 @@ impl RawClient {
             Arc::clone(&metrics),
         );
         let runtime = ClientRuntime::build(max_active_requests, runtime_mode, runtime_workers)?;
+        let future_setters = PythonFutureSetters::new(py)?;
 
         Ok(Self {
             clients: Some(TransportClients::new(
@@ -206,6 +210,7 @@ impl RawClient {
             metrics,
             active_async_requests: AsyncRequestRegistry::default(),
             active_streams: StreamRegistry::default(),
+            future_setters,
             max_response_body_size,
             buffered_body_budget,
             follow_redirects,
@@ -343,6 +348,7 @@ impl RawClient {
                 clients,
                 metrics: Arc::clone(&self.metrics),
                 pool_timeout,
+                future_setters: self.future_setters.clone_ref(py),
                 request: TransportRequest {
                     method,
                     url,
@@ -412,6 +418,7 @@ impl RawClient {
         let proxy_authorization = self.proxy_authorization.clone();
         let completion = RequestCompletion::default();
         let request_completion = completion.clone();
+        let future_setters = self.future_setters.clone_ref(py);
         self.metrics.request_started();
 
         let result = py.detach(|| {
@@ -423,6 +430,7 @@ impl RawClient {
                     active_streams,
                     runtime_handle,
                     pool_timeout,
+                    future_setters,
                     TransportRequest {
                         method,
                         url,
@@ -504,6 +512,7 @@ impl RawClient {
                 metrics: Arc::clone(&self.metrics),
                 active_streams: self.active_streams.clone(),
                 pool_timeout,
+                future_setters: self.future_setters.clone_ref(py),
                 request: TransportRequest {
                     method,
                     url,

src/py/client/streams/parts.rs

@@ -3,6 +3,7 @@ use crate::core::headers::HeaderPairs;
 use crate::core::metrics::{Metrics, ResponseBodyLifecycleOutcome};
 use crate::py::client::acquire::AcquirePermit;
 use crate::py::client::async_requests::RequestCompletion;
+use crate::py::client::future::PythonFutureSetters;
 use crate::py::client::lifecycle::ResponseBodyLifecycle;
 use crate::py::response::{RawRequestInfo, RawResponse};
 use hyper::body::Incoming;
@@ -27,6 +28,7 @@ pub(crate) struct RawStreamResponseParts {
     pub(crate) completion: RequestCompletion,
     pub(crate) registry: super::registry::StreamRegistry,
     pub(crate) runtime_handle: Handle,
+    pub(crate) future_setters: PythonFutureSetters,
     pub(crate) read_timeout: Duration,
     pub(crate) read_timeout_secs: f64,
     pub(crate) origin: String,

src/py/client/streams/response.rs

@@ -8,7 +8,6 @@ use crate::py::client::future::{complete_python_bytes_future, PythonFutureSetter
 use crate::py::response::{RawRequestInfo, RawResponse};
 use pyo3::prelude::*;
 use pyo3::types::PyAny;
-use std::sync::{Mutex, MutexGuard};
 use tokio::runtime::Handle;
 use tokio::sync::oneshot;
 
@@ -29,7 +28,7 @@ pub struct RawStreamResponse {
     history: Vec<RawResponse>,
     state: StreamState,
     runtime_handle: Handle,
-    future_setters: Mutex<Option<PythonFutureSetters>>,
+    future_setters: PythonFutureSetters,
 }
 
 impl RawStreamResponse {
@@ -51,6 +50,7 @@ impl RawStreamResponse {
             completion,
             registry,
             runtime_handle,
+            future_setters,
             read_timeout,
             read_timeout_secs,
             origin,
@@ -79,7 +79,7 @@ impl RawStreamResponse {
                 redirect_hop,
             }),
             runtime_handle,
-            future_setters: Mutex::new(None),
+            future_setters,
         }
     }
 
@@ -88,41 +88,6 @@ impl RawStreamResponse {
             response.release_body_reservations();
         }
     }
-
-    fn future_setters(&self, py: Python<'_>) -> PyResult<PythonFutureSetters> {
-        // PyO3 handle clones touch Python refcounts, so keep them outside the
-        // Rust mutex protecting the cached setters.
-        if let Some(cached_setters) = self.take_future_setters() {
-            let setters = cached_setters.clone_ref(py);
-            self.store_future_setters_if_empty(cached_setters);
-            return Ok(setters);
-        }
-
-        let setters = PythonFutureSetters::new(py)?;
-        let cached_setters = setters.clone_ref(py);
-        self.store_future_setters_if_empty(cached_setters);
-        Ok(setters)
-    }
-
-    fn take_future_setters(&self) -> Option<PythonFutureSetters> {
-        self.future_setters_guard().take()
-    }
-
-    fn store_future_setters_if_empty(&self, setters: PythonFutureSetters) {
-        let mut setters_to_store = Some(setters);
-        {
-            let mut guard = self.future_setters_guard();
-            if guard.is_none() {
-                *guard = setters_to_store.take();
-            }
-        }
-    }
-
-    fn future_setters_guard(&self) -> MutexGuard<'_, Option<PythonFutureSetters>> {
-        self.future_setters
-            .lock()
-            .unwrap_or_else(std::sync::PoisonError::into_inner)
-    }
 }
 
 #[pymethods]
@@ -176,7 +141,7 @@ impl RawStreamResponse {
         };
         let task_loop = loop_.clone_ref(py);
         let task_future = future.clone_ref(py);
-        let task_future_setters = self.future_setters(py)?;
+        let task_future_setters = self.future_setters.clone_ref(py);
         let (start_sender, start_receiver) = oneshot::channel();
 
         let handle = self.runtime_handle.spawn(async move {

src/py/client/transport/context.rs

@@ -2,6 +2,7 @@ use crate::core::metrics::{Metrics, OriginMetrics};
 use crate::core::response::BufferedBodyBudget;
 use crate::py::client::acquire::AcquirePermit;
 use crate::py::client::async_requests::RequestCompletion;
+use crate::py::client::future::PythonFutureSetters;
 use crate::py::client::streams::StreamRegistry;
 use crate::py::response::RawResponse;
 use std::sync::Arc;
@@ -30,6 +31,7 @@ pub(super) struct RawStreamResponseContext {
     pub(super) runtime_handle: Handle,
     pub(super) completion: RequestCompletion,
     pub(super) permit: AcquirePermit,
+    pub(super) future_setters: PythonFutureSetters,
     pub(super) redirect_hop: usize,
     pub(super) history: Vec<RawResponse>,
 }