From 578398780cf5e27205d133724da69d6da979226c Mon Sep 17 00:00:00 2001 From: "sergei.romanchuk" Date: Fri, 3 Jul 2026 21:53:05 +0200 Subject: [PATCH] perf(pyo3): cache async future completion helpers Cache Python asyncio future setter callables on RawClient and reuse them across buffered, stream, and async stream-read completion paths. Refs: #202 --- src/py/client/async_requests/spawn.rs | 7 +- src/py/client/async_requests/stream_spawn.rs | 8 +- src/py/client/future.rs | 67 ++++++----- src/py/client/mod.rs | 9 ++ src/py/client/streams/parts.rs | 2 + src/py/client/streams/response.rs | 43 +------ src/py/client/transport/context.rs | 2 + src/py/client/transport/response.rs | 1 + src/py/client/transport/streaming.rs | 3 + tests/pyo3_boundary/async_future_server.py | 71 +++++++++++ .../test_async_future_boundary.py | 111 ++++++++++++++++++ 11 files changed, 250 insertions(+), 74 deletions(-) create mode 100644 tests/pyo3_boundary/async_future_server.py diff --git a/src/py/client/async_requests/spawn.rs b/src/py/client/async_requests/spawn.rs index a8f4d21..91f5854 100644 --- a/src/py/client/async_requests/spawn.rs +++ b/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, 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); } }); diff --git a/src/py/client/async_requests/stream_spawn.rs b/src/py/client/async_requests/stream_spawn.rs index 02b7aa7..79e80f7 100644 --- a/src/py/client/async_requests/stream_spawn.rs +++ b/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, 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( diff --git a/src/py/client/future.rs b/src/py/client/future.rs index 83267ab..dd7540a 100644 --- a/src/py/client/future.rs +++ b/src/py/client/future.rs @@ -37,14 +37,21 @@ pub fn cancel_python_future(loop_: &Py, future: &Py) { pub fn complete_python_future( loop_: &Py, future: &Py, + setters: &PythonFutureSetters, result: PyResult, ) { 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, future: &Py, + setters: &PythonFutureSetters, result: PyResult, ) { 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, - future: &Py, - exception: Py, -) -> 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, @@ -179,15 +178,18 @@ fn schedule_set_stream_result( py: Python<'_>, loop_: &Py, future: &Py, + 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(()) } @@ -195,14 +197,17 @@ fn schedule_set_result( py: Python<'_>, loop_: &Py, future: &Py, + 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(()) } diff --git a/src/py/client/mod.rs b/src/py/client/mod.rs index ae8fb34..f4f4e48 100644 --- a/src/py/client/mod.rs +++ b/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, active_async_requests: AsyncRequestRegistry, active_streams: StreamRegistry, + future_setters: PythonFutureSetters, max_response_body_size: Option, 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, max_connections: Option, @@ -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, diff --git a/src/py/client/streams/parts.rs b/src/py/client/streams/parts.rs index f2ebde2..e7916c4 100644 --- a/src/py/client/streams/parts.rs +++ b/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, diff --git a/src/py/client/streams/response.rs b/src/py/client/streams/response.rs index 387d9da..007aa51 100644 --- a/src/py/client/streams/response.rs +++ b/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, state: StreamState, runtime_handle: Handle, - future_setters: Mutex>, + 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 { - // 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 { - 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> { - 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 { diff --git a/src/py/client/transport/context.rs b/src/py/client/transport/context.rs index 4bf5bf0..61e3090 100644 --- a/src/py/client/transport/context.rs +++ b/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, } diff --git a/src/py/client/transport/response.rs b/src/py/client/transport/response.rs index aaefbfb..cbf5ad1 100644 --- a/src/py/client/transport/response.rs +++ b/src/py/client/transport/response.rs @@ -104,6 +104,7 @@ pub(super) fn raw_stream_response( completion: context.completion, registry: context.active_streams, runtime_handle: context.runtime_handle, + future_setters: context.future_setters, read_timeout, read_timeout_secs: context.read_timeout, origin: context.origin, diff --git a/src/py/client/transport/streaming.rs b/src/py/client/transport/streaming.rs index 56274d5..f741445 100644 --- a/src/py/client/transport/streaming.rs +++ b/src/py/client/transport/streaming.rs @@ -11,6 +11,7 @@ use crate::errors::FogHttpError; use crate::messages::{redirect_limit_exceeded, REQUEST_TOTAL_TIMEOUT}; use crate::py::client::acquire::AcquireGate; use crate::py::client::async_requests::RequestCompletion; +use crate::py::client::future::PythonFutureSetters; use crate::py::client::redirects::{redirect_decision, RedirectDecision}; use crate::py::client::streams::{RawStreamResponse, StreamRegistry}; use crate::py::client::timeout_diagnostics::{ @@ -31,6 +32,7 @@ pub async fn send_stream_request( active_streams: StreamRegistry, runtime_handle: Handle, pool_timeout: f64, + future_setters: PythonFutureSetters, parts: TransportRequest, completion: RequestCompletion, ) -> PyResult { @@ -86,6 +88,7 @@ pub async fn send_stream_request( runtime_handle, completion, permit, + future_setters, redirect_hop, history, }, diff --git a/tests/pyo3_boundary/async_future_server.py b/tests/pyo3_boundary/async_future_server.py new file mode 100644 index 0000000..e5f9ef6 --- /dev/null +++ b/tests/pyo3_boundary/async_future_server.py @@ -0,0 +1,71 @@ +__all__ = ( + "DelayedResponseServer", + "cache_async_client_future_setters", + "delayed_response_server", +) + +import asyncio +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager, suppress +from dataclasses import dataclass + +import foghttp +from foghttp.status_codes.success import OK + +from .constants import WAIT_TIMEOUT + + +EMPTY_HTTP_RESPONSE = b"HTTP/1.1 200 OK\r\nContent-Length: 0\r\n\r\n" + + +@dataclass(frozen=True, slots=True) +class DelayedResponseServer: + base_url: str + request_seen: asyncio.Event + release_response: asyncio.Event + + +@asynccontextmanager +async def delayed_response_server( + response: bytes | None, +) -> AsyncIterator[DelayedResponseServer]: + request_seen = asyncio.Event() + release_response = asyncio.Event() + + async def handle(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + try: + await reader.readuntil(b"\r\n\r\n") + request_seen.set() + await release_response.wait() + if response is not None: + writer.write(response) + await writer.drain() + finally: + writer.close() + with suppress(asyncio.CancelledError, OSError): + await writer.wait_closed() + + server = await asyncio.start_server(handle, "127.0.0.1", 0) + host, port = server.sockets[0].getsockname() + try: + yield DelayedResponseServer( + base_url=f"http://{host}:{port}", + request_seen=request_seen, + release_response=release_response, + ) + finally: + release_response.set() + server.close() + await server.wait_closed() + + +async def cache_async_client_future_setters(client: foghttp.AsyncClient) -> None: + async with delayed_response_server(EMPTY_HTTP_RESPONSE) as server: + task = asyncio.create_task(client.get(server.base_url)) + await asyncio.wait_for(server.request_seen.wait(), timeout=WAIT_TIMEOUT) + server.release_response.set() + response = await asyncio.wait_for(task, timeout=WAIT_TIMEOUT) + + if response.status_code != OK: + msg = "warm-up request did not complete successfully" + raise AssertionError(msg) diff --git a/tests/pyo3_boundary/test_async_future_boundary.py b/tests/pyo3_boundary/test_async_future_boundary.py index fbe0dd0..66d5b44 100644 --- a/tests/pyo3_boundary/test_async_future_boundary.py +++ b/tests/pyo3_boundary/test_async_future_boundary.py @@ -1,11 +1,19 @@ import asyncio from typing import Any +import pytest + import foghttp +from foghttp._client import asyncio_futures from foghttp.status_codes.success import OK from tests.client_streaming.constants import GATED_STREAM_PATH from tests.client_streaming.server import start_async_streaming_server +from .async_future_server import ( + EMPTY_HTTP_RESPONSE, + cache_async_client_future_setters, + delayed_response_server, +) from .constants import ASYNC_BOUNDARY_REQUESTS, WAIT_TIMEOUT @@ -54,3 +62,106 @@ def capture_loop_error(_loop: asyncio.AbstractEventLoop, context: dict[str, Any] assert responses assert cancellations assert all(response.status_code == OK for response in responses) + + +async def test_async_buffered_completion_uses_cached_result_helper( + monkeypatch: pytest.MonkeyPatch, +) -> None: + async with ( + start_async_streaming_server() as server, + foghttp.AsyncClient() as client, + ): + await cache_async_client_future_setters(client) + + def fail_set_result(_future: asyncio.Future[object], _result: object) -> None: + msg = "request completion should use the cached result helper" + raise AssertionError(msg) + + monkeypatch.setattr(asyncio_futures, "set_result_if_pending", fail_set_result) + task = asyncio.create_task(client.get(f"{server.base_url}{GATED_STREAM_PATH}")) + await asyncio.wait_for(server.first_chunk_sent.wait(), timeout=WAIT_TIMEOUT) + server.release_tail.set() + + response = await asyncio.wait_for(task, timeout=WAIT_TIMEOUT) + + assert response.status_code == OK + + +async def test_async_stream_completion_uses_cached_result_helper( + monkeypatch: pytest.MonkeyPatch, +) -> None: + async with delayed_response_server(EMPTY_HTTP_RESPONSE) as server, foghttp.AsyncClient() as client: + await cache_async_client_future_setters(client) + + async def read_stream_status() -> int: + async with client.stream("GET", f"{server.base_url}/stream") as response: + return response.status_code + + def fail_set_result(_future: asyncio.Future[object], _result: object) -> None: + msg = "stream completion should use the cached result helper" + raise AssertionError(msg) + + monkeypatch.setattr(asyncio_futures, "set_result_if_pending", fail_set_result) + task = asyncio.create_task(read_stream_status()) + await asyncio.wait_for(server.request_seen.wait(), timeout=WAIT_TIMEOUT) + server.release_response.set() + + status_code = await asyncio.wait_for(task, timeout=WAIT_TIMEOUT) + + assert status_code == OK + + +async def test_async_error_completion_uses_cached_exception_helper( + monkeypatch: pytest.MonkeyPatch, +) -> None: + async with delayed_response_server(None) as server, foghttp.AsyncClient() as client: + await cache_async_client_future_setters(client) + + def fail_set_exception( + _future: asyncio.Future[object], + _exception: BaseException, + ) -> None: + msg = "request completion should use the cached exception helper" + raise AssertionError(msg) + + monkeypatch.setattr( + asyncio_futures, + "set_exception_if_pending", + fail_set_exception, + ) + task = asyncio.create_task(client.get(f"{server.base_url}/broken")) + await asyncio.wait_for(server.request_seen.wait(), timeout=WAIT_TIMEOUT) + server.release_response.set() + + with pytest.raises(foghttp.RequestError): + await asyncio.wait_for(task, timeout=WAIT_TIMEOUT) + + +async def test_async_stream_error_completion_uses_cached_exception_helper( + monkeypatch: pytest.MonkeyPatch, +) -> None: + async with delayed_response_server(None) as server, foghttp.AsyncClient() as client: + await cache_async_client_future_setters(client) + + async def open_stream() -> None: + async with client.stream("GET", f"{server.base_url}/broken-stream"): + pass + + def fail_set_exception( + _future: asyncio.Future[object], + _exception: BaseException, + ) -> None: + msg = "stream error completion should use the cached exception helper" + raise AssertionError(msg) + + monkeypatch.setattr( + asyncio_futures, + "set_exception_if_pending", + fail_set_exception, + ) + task = asyncio.create_task(open_stream()) + await asyncio.wait_for(server.request_seen.wait(), timeout=WAIT_TIMEOUT) + server.release_response.set() + + with pytest.raises(foghttp.RequestError): + await asyncio.wait_for(task, timeout=WAIT_TIMEOUT)