From d45461e8df99c22f1cd11ab90808b27261ca16b4 Mon Sep 17 00:00:00 2001 From: 0x676e67 Date: Tue, 6 Oct 2026 06:20:29 +0000 Subject: [PATCH 01/25] perf(resp): never poll HTTP/2 bodies on the event loop, even empty ones --- src/client/resp.rs | 13 ++---- src/client/resp/http.rs | 21 +++++---- src/client/resp/stream.rs | 3 +- tests/response_test.py | 97 +++++++++++++++++++++++++++++++++++++++ 4 files changed, 117 insertions(+), 17 deletions(-) diff --git a/src/client/resp.rs b/src/client/resp.rs index 2ca9a0fd..0312082f 100644 --- a/src/client/resp.rs +++ b/src/client/resp.rs @@ -14,12 +14,9 @@ pub use self::{ const READ_ATTACHED: u64 = 64 * 1024; /// The largest unread body the event loop polls for a response of `version`: `limit` up to -/// HTTP/1.1, else 0. An HTTP/2 stream shares its connection's state behind a lock that the -/// connection task holds while it handles frames, so polling one would stall the loop. -fn loop_limit(version: wreq::Version, limit: u64) -> u64 { - if version <= wreq::Version::HTTP_11 { - limit - } else { - 0 - } +/// HTTP/1.1, and none after it, not even an empty one. An HTTP/2 stream shares its +/// connection's state behind a lock that the connection task holds while it handles frames, +/// so polling one would stall the loop. +fn loop_limit(version: wreq::Version, limit: u64) -> Option { + (version <= wreq::Version::HTTP_11).then_some(limit) } diff --git a/src/client/resp/http.rs b/src/client/resp/http.rs index 42332175..0080d1ad 100644 --- a/src/client/resp/http.rs +++ b/src/client/resp/http.rs @@ -112,9 +112,10 @@ impl Response { } /// Take the body for reading. Cached bytes are shared at once, and a body of known - /// length up to `limit` that is already buffered is read now on the calling thread; - /// overlapping reads and `stream()` fail with [`Error::Memory`] while a read runs. - fn take_bytes(&self, limit: u64) -> Result { + /// length up to `limit` that is already buffered is read now on the calling thread, none + /// without a `limit`; overlapping reads and `stream()` fail with [`Error::Memory`] while + /// a read runs. + fn take_bytes(&self, limit: Option) -> Result { let mut slot = self.slot(); let body = match mem::replace(&mut *slot, Body::Taken) { Body::Unread(body) => body, @@ -128,7 +129,11 @@ impl Response { } }; drop(slot); - let attached = body.size_hint().exact().is_some_and(|len| len <= limit); + let attached = body + .size_hint() + .exact() + .zip(limit) + .is_some_and(|(len, limit)| len <= limit); let mut collect = body.collect(); let ready = if attached { // A read timeout starts a timer on its first poll, which needs the runtime. @@ -176,7 +181,7 @@ impl Response { /// the body is still arriving or too large to decode on the caller. fn read_body( &self, - limit: u64, + limit: Option, read: F, ) -> Result<(impl Future> + Send + 'static, bool), Error> where @@ -185,7 +190,7 @@ impl Response { { let (decode, inline) = match self.take_bytes(limit)? { BodyRead::Ready(bytes) => { - let inline = bytes.len() as u64 <= limit; + let inline = limit.is_some_and(|limit| bytes.len() as u64 <= limit); (Either::Left(read(self.build_response(bytes))), inline) } BodyRead::Pending(collect) => ( @@ -484,7 +489,7 @@ impl BlockingResponse { T: Send, { let runtime = &self.0.runtime; - let (read, inline) = self.0.read_body(READ_ATTACHED, read)?; + let (read, inline) = self.0.read_body(Some(READ_ATTACHED), read)?; if inline { nogil::run(py, runtime, read) } else { @@ -579,7 +584,7 @@ impl BlockingResponse { /// Read the body as a read-only memoryview, retaining its data after the response closes. pub fn bytes(&self, py: Python) -> PyResult { let response = &self.0; - match response.take_bytes(READ_ATTACHED)? { + match response.take_bytes(Some(READ_ATTACHED))? { BodyRead::Ready(bytes) => Ok(PyBuffer::from(bytes)), BodyRead::Pending(collect) => { let read = response diff --git a/src/client/resp/stream.rs b/src/client/resp/stream.rs index 8c8d621c..b4f9dfb6 100644 --- a/src/client/resp/stream.rs +++ b/src/client/resp/stream.rs @@ -171,7 +171,8 @@ impl Streamer { && !resp .size_hint() .exact() - .is_some_and(|len| len <= loop_limit(resp.version(), limit)) + .zip(loop_limit(resp.version(), limit)) + .is_some_and(|(len, limit)| len <= limit) { return None; } diff --git a/tests/response_test.py b/tests/response_test.py index 24aadfb5..1dfbb405 100644 --- a/tests/response_test.py +++ b/tests/response_test.py @@ -3,6 +3,7 @@ import socket import threading import time +from contextlib import asynccontextmanager from datetime import timedelta import pytest @@ -251,6 +252,102 @@ async def test_pending_read_holds_the_body_until_it_ends(): await response.text() +async def steps(coroutine): + """Drive `coroutine` by hand; return its result and how often it suspended.""" + suspended = 0 + while True: + try: + waiter = coroutine.send(None) + except StopIteration as stop: + return stop.value, suspended + except StopAsyncIteration: + return None, suspended + suspended += 1 + if waiter is None: + await asyncio.sleep(0) + else: + await asyncio.wait({waiter}, timeout=5) + + +@asynccontextmanager +async def body_server(http2, body): + """Answer every request with `body` and its Content-Length, over HTTP/1 or h2c.""" + length = str(len(body)).encode() + handlers, writers = set(), [] + + def frame(kind, flags, stream, payload=b""): + size = len(payload).to_bytes(3, "big") + return size + bytes((kind, flags)) + stream.to_bytes(4, "big") + payload + + async def accept(reader, writer): + handlers.add(asyncio.current_task()) + writers.append(writer) + try: + if not http2: + await reader.readuntil(b"\r\n\r\n") + head = b"HTTP/1.1 200 OK\r\nConnection: close\r\nContent-Length: " + writer.write(head + length + b"\r\n\r\n" + body) + await writer.drain() + await reader.read() + return + assert await reader.readexactly(24) == b"PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n" + writer.write(frame(4, 0, 0)) + while True: + header = await reader.readexactly(9) + payload = await reader.readexactly(int.from_bytes(header[:3], "big")) + kind, flags = header[3:5] + stream = int.from_bytes(header[5:], "big") & 0x7FFFFFFF + if kind == 4 and not flags & 1: + writer.write(frame(4, 1, 0)) + elif kind == 1: + # `:status: 200` and `content-length`, then the body ending the stream. + block = b"\x88\x0f\x0d" + bytes((len(length),)) + length + if body: + writer.write( + frame(1, 4, stream, block) + frame(0, 1, stream, body) + ) + else: + writer.write(frame(1, 5, stream, block)) + await writer.drain() + except (asyncio.IncompleteReadError, ConnectionError): + pass + finally: + handlers.discard(asyncio.current_task()) + writer.close() + + server = await asyncio.start_server(accept, "127.0.0.1", 0) + try: + yield f"http://127.0.0.1:{server.sockets[0].getsockname()[1]}/" + finally: + server.close() + for writer in writers: + writer.close() + await asyncio.gather(*handlers, return_exceptions=True) + await server.wait_closed() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("read", ["bytes", "stream"]) +@pytest.mark.parametrize("body", [b"ok", b""], ids=["body", "empty"]) +@pytest.mark.parametrize("http2", [False, True], ids=["http1", "http2"]) +async def test_only_http1_bodies_are_read_on_the_event_loop(http2, body, read): + # A buffered HTTP/1 body is read in the first step of the coroutine. An HTTP/2 body, + # even an empty one, is always read on the runtime, so the coroutine suspends first. + async with body_server(http2, body) as url: + async with wreq.Client(http2_only=http2, proxies=[]) as client: + response = await asyncio.wait_for(client.get(url), 5) + assert response.version == (Version.HTTP_2 if http2 else Version.HTTP_11) + # Let the whole body arrive before the read starts. + await asyncio.sleep(0.05) + if read == "bytes": + result, suspended = await steps(response.bytes()) + assert bytes(result) == body + else: + result, suspended = await steps(response.stream().__anext__()) + assert (bytes(result) if result is not None else b"") == body + assert (suspended > 0) == http2 + + @pytest.mark.asyncio async def test_stream_read_ahead_yields_and_closes_waiting_readers(): stalled = threading.Event() From ca35827994f8cdf5cbefe8e469b09ae313f2b7fe Mon Sep 17 00:00:00 2001 From: 0x676e67 Date: Tue, 6 Oct 2026 06:20:29 +0000 Subject: [PATCH 02/25] test(ws): assert send and close resolve to None --- tests/websocket_test.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tests/websocket_test.py b/tests/websocket_test.py index 4d9ff0e9..2d5e0654 100644 --- a/tests/websocket_test.py +++ b/tests/websocket_test.py @@ -73,7 +73,7 @@ async def test_websocket_close_frame(code, reason, expected): async with wreq.Client(proxies=[]) as client: # Leaving the block after an explicit close must not fail. async with client.websocket(url) as ws: - await ws.close(code, reason) + assert await ws.close(code, reason) is None assert await asyncio.wait_for(frames.get(), 5) == expected finally: server.close() @@ -137,7 +137,10 @@ async def send_later(): await asyncio.wait_for(ws.recv(), 0.2) # A pending receive does not hold up a send. pending = asyncio.ensure_future(ws.recv()) - await asyncio.wait_for(ws.send(wreq.Message.from_text("ping")), 5) + assert ( + await asyncio.wait_for(ws.send(wreq.Message.from_text("ping")), 5) + is None + ) assert await asyncio.wait_for(received.get(), 5) == (0x1, b"ping") release.set() assert (await asyncio.wait_for(pending, 5)).text == "first" From bd911fed0301b708e603c8f6eb68f13dccd3dd72 Mon Sep 17 00:00:00 2001 From: 0x676e67 Date: Tue, 6 Oct 2026 12:40:36 +0000 Subject: [PATCH 03/25] feat(runtime)!: replace work_steal with Scheduler and add CURRENT_THREAD --- bench/README.md | 5 +- bench/clients.py | 10 +- bench/registry.py | 4 +- bench/report.py | 3 + bench/tests/test_benchmark.py | 5 +- docs/source/api/runtime.md | 8 +- docs/source/guide/advanced.md | 28 ++-- docs/source/guide/blocking.md | 37 +++++- docs/templates/benchmark.md | 4 +- python/wreq/runtime.py | 34 ++++- python/wreq/wreq.py | 5 +- src/client.rs | 33 ++++- src/client/body/multipart.rs | 12 +- src/client/body/stream.rs | 106 +++++++++++++-- src/client/nogil.rs | 29 +++-- src/client/resp/http.rs | 9 +- src/client/resp/stream.rs | 22 +++- src/lib.rs | 3 +- src/runtime.rs | 234 +++++++++++++++++++++++++++++----- tests/runtime_test.py | 222 ++++++++++++++++++++++++++++---- tests/sync_upload_test.py | 72 +++++++++-- 21 files changed, 748 insertions(+), 137 deletions(-) diff --git a/bench/README.md b/bench/README.md index e587298c..33c4492d 100644 --- a/bench/README.md +++ b/bench/README.md @@ -205,8 +205,9 @@ shows each run's revision and collection time separately. ## Blocking wreq runtimes `wreq_blocking` uses the default shared multi-thread network runtime (MT). -`wreq_blocking_st` shares one `Runtime(workers=1, work_steal=False)` across -all logical workers in each case (ST). Both retain one client per logical +`wreq_blocking_st` shares one +`Runtime(scheduler=Scheduler.PER_WORKER, workers=1)` across all logical workers +in each case (ST). Both retain one client per logical worker and the same Python thread pool for concurrent blocking calls. The default full run includes both async and blocking wreq MT/ST, for 17 diff --git a/bench/clients.py b/bench/clients.py index 60fdd077..b3ee666a 100644 --- a/bench/clients.py +++ b/bench/clients.py @@ -112,9 +112,13 @@ def network_runtime(client): return None from wreq.runtime import Runtime - return Runtime( - workers=spec.runtime["workers"], work_steal=spec.runtime["work_steal"] - ) + workers, scheduler = spec.runtime["workers"], spec.runtime["scheduler"] + try: + from wreq.runtime import Scheduler + except ImportError: + # Releases before Scheduler choose it with `work_steal`. + return Runtime(workers=workers, work_steal=scheduler == "WORK_STEALING") + return Runtime(scheduler=getattr(Scheduler, scheduler), workers=workers) @asynccontextmanager diff --git a/bench/registry.py b/bench/registry.py index 6fe404f0..48886bb3 100644 --- a/bench/registry.py +++ b/bench/registry.py @@ -59,7 +59,7 @@ def settings(self): "wreq_st": ClientSpec( "wreq", "wreq (ST)", - runtime={"kind": "custom", "workers": 1, "work_steal": False}, + runtime={"kind": "custom", "workers": 1, "scheduler": "PER_WORKER"}, color="wreq_st", ), "pyreqwest_st": ClientSpec( @@ -117,7 +117,7 @@ def settings(self): runtime={ "kind": "custom", "workers": 1, - "work_steal": False, + "scheduler": "PER_WORKER", "scope": "shared per case", }, color="wreq_st", diff --git a/bench/report.py b/bench/report.py index 01858289..77a66bf9 100644 --- a/bench/report.py +++ b/bench/report.py @@ -44,6 +44,9 @@ def runtime_label(runtime: dict) -> str: return "Multi-thread runtime" if kind == "custom": workers = runtime.get("workers", "not recorded") + if "scheduler" in runtime: + return f"Custom: workers={workers}, scheduler={runtime['scheduler']}" + # Snapshots recorded before Scheduler. steal = runtime.get("work_steal", "not recorded") return f"Custom: workers={workers}, work_steal={steal}" if kind == "thread_pool": diff --git a/bench/tests/test_benchmark.py b/bench/tests/test_benchmark.py index 27021e1f..b5f20d6f 100644 --- a/bench/tests/test_benchmark.py +++ b/bench/tests/test_benchmark.py @@ -7,7 +7,7 @@ import json import time import sys -from types import ModuleType +from types import ModuleType, SimpleNamespace import pytest @@ -370,12 +370,13 @@ def test_blocking_st_shares_one_network_runtime(monkeypatch): runtime_module = ModuleType("wreq.runtime") def create_runtime(**options): - assert options == {"workers": 1, "work_steal": False} + assert options == {"scheduler": "per-worker", "workers": 1} runtime = object() runtimes.append(runtime) return runtime runtime_module.Runtime = create_runtime + runtime_module.Scheduler = SimpleNamespace(PER_WORKER="per-worker") monkeypatch.setitem(sys.modules, "wreq.runtime", runtime_module) @contextmanager diff --git a/docs/source/api/runtime.md b/docs/source/api/runtime.md index dfda7fb5..565c75f2 100644 --- a/docs/source/api/runtime.md +++ b/docs/source/api/runtime.md @@ -1,7 +1,11 @@ # wreq.runtime -Runtime configuration for asynchronous and blocking clients. `Runtime` is also -available as `wreq.Runtime`. +Runtime configuration for asynchronous and blocking clients. `Runtime` and +`Scheduler` are also available as `wreq.Runtime` and `wreq.Scheduler`. + +::: wreq.runtime.Scheduler + options: + show_root_heading: true ::: wreq.runtime.Runtime options: diff --git a/docs/source/guide/advanced.md b/docs/source/guide/advanced.md index caf48503..4f56ceff 100644 --- a/docs/source/guide/advanced.md +++ b/docs/source/guide/advanced.md @@ -105,11 +105,11 @@ start a separate worker pool for an async or blocking client: from datetime import timedelta from wreq import Client -from wreq.runtime import Runtime +from wreq.runtime import Runtime, Scheduler runtime = Runtime( + scheduler=Scheduler.PER_WORKER, workers=1, - work_steal=False, thread_name="http-client", max_blocking_threads=8, thread_keep_alive=timedelta(seconds=10), @@ -117,15 +117,25 @@ runtime = Runtime( client = Client(runtime=runtime) ``` -With `work_steal=False`, workers use independent single-thread Tokio runtimes. -Each client is assigned one worker for its lifetime; requests, response reads, -streams and WebSocket operations use that worker. Async reads of small HTTP/1 +`scheduler` selects how the runtime runs client work: + +- `Scheduler.WORK_STEALING` (the default) uses one multi-thread pool whose + workers steal work from each other. +- `Scheduler.PER_WORKER` gives each worker its own single-thread Tokio runtime. +- `Scheduler.CURRENT_THREAD` has no workers: blocking calls drive its IO, which + is fastest with a client and runtime per thread. It serves only blocking + clients; see the [blocking guide](blocking.md#current-thread-runtime). + +With `Scheduler.PER_WORKER`, each client is assigned one worker for its +lifetime; requests, response reads, streams and WebSocket operations use that +worker. Async reads of small HTTP/1 bodies that have already arrived finish on the event loop thread instead. With multiple workers, newly created clients select a worker randomly and keep that selection. This is not CPU pinning. Sharing the same `Runtime` between clients is supported, and `client.runtime` returns the shared runtime object. -`workers=None` uses the available CPU parallelism, or 1 if it cannot be determined. +`workers=None` uses the available CPU parallelism, or 1 if it cannot be determined; +`Scheduler.CURRENT_THREAD` requires it. Custom runtimes start their threads during construction, before any client is bound or request is sent. @@ -133,9 +143,9 @@ bound or request is sent. `thread_keep_alive` accepts a nonnegative `datetime.timedelta`. `max_blocking_threads` and `thread_keep_alive` default to Tokio's settings -(512 and 10 seconds). In -no-steal mode these limits apply to **each worker's** blocking pool, not the pool -as a whole. Python async upload generators still run on the caller's event loop. +(512 and 10 seconds). With +`Scheduler.PER_WORKER` these limits apply to **each worker's** blocking pool, not +the pool as a whole. Python async upload generators still run on the caller's event loop. Standalone multipart file preparation and upload-task cleanup can use the shared runtime; a dedicated client runtime does not isolate Python's GIL or every process resource. DNS resolvers are owned by individual clients so their diff --git a/docs/source/guide/blocking.md b/docs/source/guide/blocking.md index e731f43d..46f13a28 100644 --- a/docs/source/guide/blocking.md +++ b/docs/source/guide/blocking.md @@ -72,9 +72,9 @@ it uses the shared global multi-thread runtime. ```python from wreq.blocking import Client -from wreq.runtime import Runtime +from wreq.runtime import Runtime, Scheduler -runtime = Runtime(workers=1, work_steal=False) +runtime = Runtime(scheduler=Scheduler.PER_WORKER, workers=1) with Client(runtime=runtime) as client: with client.get("https://httpbin.io/get") as response: print(response.text()) @@ -86,6 +86,39 @@ runtime; it cancels pending requests and rejects new ones with `asyncio.CancelledError`. See [custom runtimes](advanced.md#custom-runtimes) for configuration and lifetime details. +### Current-thread runtime + +A `Scheduler.CURRENT_THREAD` runtime has no worker threads: the thread in a +blocking call drives its IO, with no handoff to another thread. It is fastest +with a client and runtime per thread, like one curl handle per thread: + +```python +import threading + +from wreq.blocking import Client +from wreq.runtime import Runtime, Scheduler + +local = threading.local() + + +def client() -> Client: + if not hasattr(local, "client"): + local.client = Client(runtime=Runtime(scheduler=Scheduler.CURRENT_THREAD)) + return local.client +``` + +- Nothing runs between calls. HTTP/2 pings, idle-connection cleanup, closing + cancelled connections and WebSocket keepalive progress only during a call. +- Threads sharing the runtime take turns driving it, and the driving thread + also runs the other threads' connections and upload iterators. A client + shared by many threads is faster with `Scheduler.PER_WORKER`. +- Upload iterators run on the driving thread. A slow `__next__` holds back the + call's timeout and every thread sharing the runtime, and an iterator that + waits for another call on the same runtime never returns. Blocking wreq calls + from an iterator, including reading a response or stream, raise `RuntimeError`. +- Async clients reject the runtime, and its response streams support only + blocking iteration. + ### Cookies ```python diff --git a/docs/templates/benchmark.md b/docs/templates/benchmark.md index ae58263a..b81e66d7 100644 --- a/docs/templates/benchmark.md +++ b/docs/templates/benchmark.md @@ -60,8 +60,8 @@ both the upload and download. `wreq (MT)` uses the shared multithreaded runtime, with its worker count set by available CPU parallelism. `wreq (ST)` uses -`Runtime(workers=1, work_steal=False)`. The ST and MT labels refer to the Rust -runtime, not Python threads. `ry` uses its default runtime; the two `pyreqwest` +`Runtime(scheduler=Scheduler.PER_WORKER, workers=1)`. The ST and MT labels refer +to the Rust runtime, not Python threads. `ry` uses its default runtime; the two `pyreqwest` variants use single-threaded and multithreaded runtimes. New runs also include `wreq (blocking MT)` and `wreq (blocking ST)`. Blocking diff --git a/python/wreq/runtime.py b/python/wreq/runtime.py index 1a3c6ab7..88c072ad 100644 --- a/python/wreq/runtime.py +++ b/python/wreq/runtime.py @@ -1,7 +1,28 @@ import datetime +from enum import Enum, auto from typing import final -__all__ = ["Runtime"] +__all__ = ["Runtime", "Scheduler"] + + +@final +class Scheduler(Enum): + """How a runtime schedules client work.""" + + WORK_STEALING = auto() + """Worker threads share client work and steal it from each other.""" + + PER_WORKER = auto() + """Each worker runs its own single-thread runtime; a client stays on one worker.""" + + CURRENT_THREAD = auto() + """No worker threads: blocking calls drive its IO, taking turns when threads + share it, and nothing runs between calls. + + Best with a client and runtime per thread; only blocking clients can use it. + Upload iterators run on the driving thread and must not wait for, or make, + other calls on this runtime. + """ @final @@ -10,26 +31,25 @@ class Runtime: Clients, responses and active work keep it alive; the last owner releases it automatically. - No-steal clients keep a fixed worker, without CPU pinning. + Per-worker clients keep a fixed worker, without CPU pinning. """ def __init__( self, *, + scheduler: Scheduler = Scheduler.WORK_STEALING, workers: int | None = None, - work_steal: bool = True, thread_name: str | None = None, max_blocking_threads: int | None = None, thread_keep_alive: datetime.timedelta | None = None, ) -> None: """Thread counts must be positive; thread_keep_alive is a nonnegative timedelta. - workers defaults to available CPU parallelism, or 1 if unavailable. - work_steal=True uses a multi-thread pool; False uses independent - single-thread workers. + workers defaults to available CPU parallelism, or 1 if unavailable, and + must be None with CURRENT_THREAD. thread_name defaults to the package name, "wreq-python". Thread names cannot contain NUL characters. - Blocking-pool settings apply to each worker runtime in no-steal mode. + Blocking-pool settings apply to each worker runtime with PER_WORKER. None preserves Tokio's blocking-pool defaults. """ ... diff --git a/python/wreq/wreq.py b/python/wreq/wreq.py index 04109214..01757b17 100644 --- a/python/wreq/wreq.py +++ b/python/wreq/wreq.py @@ -290,6 +290,8 @@ class Streamer: When streaming a response, each iteration yields either a memoryview (for body data) or a HeaderMap (for HTTP trailers, if the server sends them). This allows you to access HTTP/1.1 or HTTP/2 trailers in addition to the main body. Data views retain their backing data after the stream is closed. + A stream on a `CURRENT_THREAD` runtime supports only blocking iteration; + its async entry points raise `RuntimeError`. # Examples @@ -538,7 +540,8 @@ def __str__(self) -> str: ... class ClientConfig(TypedDict): runtime: NotRequired[Runtime | None] - """Runtime for this client and its responses; None uses the shared default.""" + """Runtime for this client and its responses; None uses the shared default. + Async clients reject a `CURRENT_THREAD` runtime with `ValueError`.""" emulation: NotRequired[emulation.Emulation | emulation.Profile] """Emulation config.""" diff --git a/src/client.rs b/src/client.rs index 32324568..67745582 100644 --- a/src/client.rs +++ b/src/client.rs @@ -11,7 +11,9 @@ use std::{ time::Duration, }; -use pyo3::{IntoPyObjectExt, prelude::*, pybacked::PyBackedStr, types::PyDict}; +use pyo3::{ + IntoPyObjectExt, exceptions::PyValueError, prelude::*, pybacked::PyBackedStr, types::PyDict, +}; use req::{Request, WebSocketRequest}; use tokio_util::sync::CancellationToken; use wreq::tls::trust::CertStore; @@ -312,12 +314,9 @@ impl Default for Client { } } -#[pymethods] impl Client { - /// Creates a new Client instance. - #[new] - #[pyo3(signature = (**kwds))] - fn new(py: Python, kwds: Option) -> PyResult { + /// Build a client for either API from its options. + fn build(py: Python, kwds: Option) -> PyResult { py.detach(|| { let runtime = match kwds.as_ref().and_then(|config| config.runtime.as_ref()) { Some(runtime) => runtime.select()?, @@ -541,6 +540,26 @@ impl Client { .map_err(Into::into) }) } +} + +#[pymethods] +impl Client { + /// Creates a new Client instance. + #[new] + #[pyo3(signature = (**kwds))] + fn new(py: Python, kwds: Option) -> PyResult { + // Nothing would drive its work while a coroutine waits. + if kwds + .as_ref() + .and_then(|config| config.runtime.as_ref()) + .is_some_and(runtime::Runtime::is_current_thread) + { + return Err(PyValueError::new_err( + "CURRENT_THREAD runtimes only serve blocking clients", + )); + } + Client::build(py, kwds) + } /// Cancel pending requests and reject new ones with asyncio.CancelledError. /// Existing responses, WebSockets and the shared runtime remain usable. @@ -703,7 +722,7 @@ impl BlockingClient { #[new] #[pyo3(signature = (**kwds))] fn new(py: Python, kwds: Option) -> PyResult { - Client::new(py, kwds).map(BlockingClient) + Client::build(py, kwds).map(BlockingClient) } /// Get the cookie jar of the client. diff --git a/src/client/body/multipart.rs b/src/client/body/multipart.rs index 854897dd..1427f52a 100644 --- a/src/client/body/multipart.rs +++ b/src/client/body/multipart.rs @@ -135,13 +135,13 @@ impl Part { Value::Text(text) => multipart::Part::stream(text.0), Value::Bytes(bytes) => multipart::Part::stream(bytes.0), // Opening the file blocks, so only that waits detached. - Value::File(path) => py - .detach(|| { - runtime::get() - .handle() - .block_on(multipart::Part::file(path)) + Value::File(path) => py.detach(|| { + runtime::get().block_on(async { + multipart::Part::file(path) + .await + .map_err(|err| Error::from(err).into()) }) - .map_err(Error::from)?, + })?, Value::Stream(stream) => { let stream = Body::wrap_stream(stream); match self.length { diff --git a/src/client/body/stream.rs b/src/client/body/stream.rs index bb051763..ea9e43e6 100644 --- a/src/client/body/stream.rs +++ b/src/client/body/stream.rs @@ -1,19 +1,21 @@ //! Request bodies streamed from Python iterators and async generators. use std::{ - pin::Pin, + mem, + pin::{Pin, pin}, sync::{ Arc, Mutex, MutexGuard, PoisonError, atomic::{AtomicBool, Ordering}, }, task::{Context, Poll}, + thread::{self, ThreadId}, time::Duration, }; use bytes::Bytes; use futures_util::Stream; use pyo3::{ - exceptions::{PyRuntimeError, PyStopIteration}, + exceptions::{PyException, PyKeyboardInterrupt, PyRuntimeError, PyStopIteration}, intern, prelude::*, sync::PyOnceLock, @@ -22,7 +24,7 @@ use pyo3::{ use tokio::{ runtime::Handle, sync::mpsc::{self, error::TrySendError}, - task::spawn_blocking, + task::{spawn_blocking, yield_now}, time, }; @@ -50,13 +52,35 @@ enum Source { Async(PyAsyncStream), } -/// A request body from a Python iterator, read on Tokio's blocking pool. +/// A request body from a Python iterator. /// -/// Neither a Tokio worker nor a blocked caller runs the iterator, so a slow `__next__` -/// cannot delay the response, a timeout or cancellation. A pump task reads an item only -/// once the channel has room, staying at most one item ahead of the upload. When the -/// upload stalls, it parks the iterator and frees its thread; the next item taken restarts it. -struct SyncStream { +/// Worker runtimes read the iterator on Tokio's blocking pool: neither a worker nor a +/// blocked caller runs it, so a slow `__next__` cannot delay the response, a timeout or +/// cancellation. A current-thread runtime reads it on its driving thread instead, which +/// waits for the request anyway. +enum SyncStream { + /// Not polled yet; the first poll picks where the iterator is read. `owner` passed the + /// iterator to the request. + Idle { + iter: Py, + owner: ThreadId, + }, + /// Read on the thread driving a current-thread runtime; `flushed` once a poll has + /// returned to the connection since the last item. + Inline { + iter: Py, + owner: ThreadId, + flushed: bool, + }, + Pumped(Pumped), + /// The inline iterator ended or raised. + Done, +} + +/// An iterator read by pump tasks on the blocking pool. A pump reads an item only once +/// the channel has room, staying at most one item ahead of the upload. When the upload +/// stalls, it parks the iterator and frees its thread; the next item taken restarts it. +struct Pumped { rx: mpsc::Receiver, /// The iterator and sender while no pump runs: before the first poll, or once a pump /// parked at a full channel. @@ -121,7 +145,10 @@ impl FromPyObject<'_, '_> for PyStream { let source = if ob.cast::().is_err() && ob.hasattr(intern!(ob.py(), "asend"))? { Source::Async(PyAsyncStream::new(ob.to_owned())?) } else { - Source::Sync(SyncStream::new(ob.to_owned().unbind())) + Source::Sync(SyncStream::Idle { + iter: ob.to_owned().unbind(), + owner: thread::current().id(), + }) }; Ok(PyStream(source)) } @@ -141,9 +168,66 @@ impl Stream for PyStream { // ===== impl SyncStream ===== impl SyncStream { + fn poll_next(&mut self, cx: &mut Context<'_>) -> Poll> { + if let SyncStream::Idle { .. } = self + && let SyncStream::Idle { iter, owner } = mem::replace(self, SyncStream::Done) + { + *self = if runtime::driving() { + SyncStream::Inline { + iter, + owner, + flushed: false, + } + } else { + SyncStream::Pumped(Pumped::new(iter)) + }; + } + match self { + SyncStream::Inline { + iter, + owner, + flushed, + } => { + // Return to the connection before each read, so the request head and the + // chunks already yielded go out before `__next__` can block. The deferred + // wake lets the runtime poll its IO and the caller first. + if !mem::replace(flushed, true) { + let _ = pin!(yield_now()).poll(cx); + return Poll::Pending; + } + *flushed = false; + // Like a pump, stop reading once Python is unavailable. + let item = Python::try_attach(|py| { + let item = next_item(py, iter); + // A KeyboardInterrupt or SystemExit reaches the caller as itself: the + // owner's call, or for Ctrl+C the main thread signals reach. + if let Some(Err(err)) = &item + && !err.is_instance_of::(py) + && (*owner == thread::current().id() + || err.is_instance_of::(py)) + { + runtime::interrupt(err.clone_ref(py)); + } + item + }) + .flatten(); + if !matches!(item, Some(Ok(_))) { + *self = SyncStream::Done; + } + Poll::Ready(item) + } + SyncStream::Pumped(pumped) => pumped.poll_next(cx), + SyncStream::Idle { .. } | SyncStream::Done => Poll::Ready(None), + } + } +} + +// ===== impl Pumped ===== + +impl Pumped { fn new(iter: Py) -> Self { let (tx, rx) = mpsc::channel(1); - SyncStream { + Pumped { rx, parked: Arc::new(Mutex::new(Some(Pump { iter, tx }))), } diff --git a/src/client/nogil.rs b/src/client/nogil.rs index bb834ff6..4aa6c53e 100644 --- a/src/client/nogil.rs +++ b/src/client/nogil.rs @@ -6,20 +6,26 @@ use std::{future::Future, pin::pin}; use futures_util::FutureExt; use pyo3::{exceptions::PyRuntimeError, prelude::*}; -use crate::runtime::Runtime; +use crate::runtime::{self, Runtime}; -/// Run network work on its selected worker; only wait for completion on the caller. -/// The caller must be detached from Python and outside an async Tokio context. -/// Its runtime borrow keeps the workers alive until the join completes. +/// Run network work on its selected worker and only wait on the caller, or drive it here +/// on a current-thread runtime. The caller must be detached from Python and outside an +/// async Tokio context. Its runtime borrow keeps the workers alive until the join completes. pub fn block_on(runtime: &Runtime, future: F) -> PyResult where F: Future> + Send + 'static, T: Send + 'static, { - runtime - .handle() - .block_on(runtime.handle().spawn(future)) - .map_err(|err| PyRuntimeError::new_err(err.to_string()))? + if runtime.is_current_thread() { + return runtime.block_on(future); + } + runtime.block_on(async { + runtime + .handle() + .spawn(future) + .await + .map_err(|err| PyRuntimeError::new_err(err.to_string()))? + }) } /// Poll `future` on the caller instead of spawning it: return still attached when it is @@ -30,14 +36,15 @@ where F: Future> + Send, T: Send, { - let handle = runtime.handle(); + // Refused before the first poll, which a later refusal would leave half done. + runtime::refuse_nested()?; let mut future = pin!(future); let ready = { - let _runtime = handle.enter(); + let _runtime = runtime.handle().enter(); future.as_mut().now_or_never() }; match ready { Some(output) => output, - None => py.detach(|| handle.block_on(future)), + None => py.detach(|| runtime.block_on(future)), } } diff --git a/src/client/resp/http.rs b/src/client/resp/http.rs index 0080d1ad..b5cadc54 100644 --- a/src/client/resp/http.rs +++ b/src/client/resp/http.rs @@ -23,7 +23,7 @@ use crate::{ header::HeaderMap, http::{StatusCode, Version}, redirect::History, - runtime::Runtime, + runtime::{self, Runtime}, tls::TlsInfo, }; @@ -488,12 +488,14 @@ impl BlockingResponse { Fut: Future> + Send + 'static, T: Send, { + // Refused before the body is taken, so it can still be read afterwards. + runtime::refuse_nested()?; let runtime = &self.0.runtime; let (read, inline) = self.0.read_body(Some(READ_ATTACHED), read)?; if inline { nogil::run(py, runtime, read) } else { - py.detach(|| runtime.handle().block_on(read)) + py.detach(|| runtime.block_on(read)) } } } @@ -583,6 +585,7 @@ impl BlockingResponse { /// Read the body as a read-only memoryview, retaining its data after the response closes. pub fn bytes(&self, py: Python) -> PyResult { + runtime::refuse_nested()?; let response = &self.0; match response.take_bytes(Some(READ_ATTACHED))? { BodyRead::Ready(bytes) => Ok(PyBuffer::from(bytes)), @@ -591,7 +594,7 @@ impl BlockingResponse { .collect_later(collect) .map_ok(|(_, bytes)| PyBuffer::from(bytes)) .map_err(Into::into); - py.detach(|| response.runtime.handle().block_on(read)) + py.detach(|| response.runtime.block_on(read)) } } } diff --git a/src/client/resp/stream.rs b/src/client/resp/stream.rs index b4f9dfb6..46530778 100644 --- a/src/client/resp/stream.rs +++ b/src/client/resp/stream.rs @@ -29,7 +29,7 @@ use crate::{ coroutine::{self, Coroutine}, error::Error, header::HeaderMap, - runtime::Runtime, + runtime::{self, Runtime}, }; /// A response frame exposed as a read-only memoryview or a header map. @@ -254,6 +254,17 @@ impl Streamer { } } + /// Refuse async use on a current-thread runtime: nothing would drive the read while a + /// coroutine waits, so iteration would hang once the buffered frames run out. + fn refuse_async(&self) -> PyResult<()> { + if self.runtime.is_current_thread() { + return Err(PyRuntimeError::new_err( + "Streams on a CURRENT_THREAD runtime only support blocking iteration", + )); + } + Ok(()) + } + /// Wait for the next frame; `end` builds the error that ends iteration. async fn next(&self, end: fn() -> Error) -> PyResult { loop { @@ -276,6 +287,8 @@ impl Streamer { } fn __next__(&self, py: Python) -> PyResult { + // Refused before any frame, so a nested read never leaves a stream half consumed. + runtime::refuse_nested()?; // Frames already received are returned without releasing the GIL. if let Some(frame) = self.ready_frame(None, || Error::StopIteration) { return frame; @@ -300,12 +313,14 @@ impl Streamer { #[pymethods] impl Streamer { - fn __aiter__(slf: PyRef) -> PyRef { - slf + fn __aiter__(slf: PyRef) -> PyResult> { + slf.refuse_async()?; + Ok(slf) } /// Read the next frame when awaited. fn __anext__(slf: Bound<'_, Self>) -> PyResult> { + slf.get().refuse_async()?; let py = slf.py(); let slf = slf.unbind(); coroutine::local(py, "Streamer.__anext__", async move { @@ -326,6 +341,7 @@ impl Streamer { } fn __aenter__(slf: Bound<'_, Self>) -> PyResult> { + slf.get().refuse_async()?; coroutine::ready("Streamer.__aenter__", slf) } diff --git a/src/lib.rs b/src/lib.rs index f165edd4..6ec5ee70 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -42,7 +42,7 @@ use http2::{ use mimalloc as _; use proxy::Proxy; use pyo3::{intern, prelude::*, types::PyDict, wrap_pymodule}; -use runtime::Runtime; +use runtime::{Runtime, Scheduler}; #[cfg(feature = "jemalloc")] use tikv_jemallocator as _; use tls::{ @@ -398,6 +398,7 @@ fn proxy_module(m: &Bound<'_, PyModule>) -> PyResult<()> { #[pymodule(gil_used = false, name = "runtime")] fn runtime_module(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; + m.add_class::()?; Ok(()) } diff --git a/src/runtime.rs b/src/runtime.rs index 37cb05e8..81f9be14 100644 --- a/src/runtime.rs +++ b/src/runtime.rs @@ -1,4 +1,6 @@ use std::{ + cell::{Cell, RefCell}, + future::Future, sync::{Arc, OnceLock}, thread, time::Duration, @@ -9,24 +11,89 @@ use pyo3::{ exceptions::{PyRuntimeError, PyValueError}, prelude::*, }; -use tokio::runtime::Handle; +use tokio::runtime::{Builder, Handle}; + +/// How a [`Runtime`] schedules client work. +#[pyclass(eq, eq_int, frozen, from_py_object)] +#[derive(Clone, Copy, PartialEq, Eq)] +#[allow(non_camel_case_types)] +pub enum Scheduler { + /// Worker threads share client work and steal it from each other. + WORK_STEALING, + /// Each worker runs its own single-thread runtime, and a client stays on one worker. + PER_WORKER, + /// No worker threads: blocking calls drive its IO, taking turns when threads share it, + /// and nothing runs between calls. Best with a client and runtime per thread; only + /// blocking clients can use it. + CURRENT_THREAD, +} /// Shared Tokio runtime, released after its clients and active work are dropped. #[derive(Clone)] #[pyclass(frozen, skip_from_py_object)] pub struct Runtime { /// Shared owner; `None` only once `Drop` has taken it to shut the runtime down. - inner: Option>, - /// The worker this copy runs on; without work stealing its client's tasks stay there. + inner: Option>, + /// The worker this copy runs on; per-worker clients keep their tasks there. handle: Handle, } +/// The runtime behind [`Runtime`]. +enum Shared { + Workers(PingoraRuntime), + /// Driven only by the threads blocked in [`Runtime::block_on`]. + CurrentThread(tokio::runtime::Runtime), +} + +/// Marks the calling thread as driving a current-thread runtime until dropped. +struct Driving; + +thread_local! { + static DRIVING: Cell = const { Cell::new(false) }; + /// A KeyboardInterrupt or other non-`Exception` error that this thread's upload + /// iterator raised while it drove the runtime, re-raised by its call instead of the + /// wrapped request error. + static INTERRUPT: RefCell> = const { RefCell::new(None) }; +} + +/// Whether this thread drives a current-thread runtime, which then runs the Python code +/// its tasks call, such as upload iterators. +pub fn driving() -> bool { + DRIVING.get() +} + +/// Re-raise `err` from the blocking call this thread is driving, once it returns. +pub fn interrupt(err: PyErr) { + INTERRUPT.set(Some(err)); +} + +// ===== impl Runtime ===== + impl Runtime { /// Borrow the selected worker's handle without changing the selection. pub fn handle(&self) -> &Handle { &self.handle } + /// Whether blocking calls drive this runtime on their own thread. + pub fn is_current_thread(&self) -> bool { + matches!(self.inner.as_deref(), Some(Shared::CurrentThread(_))) + } + + /// Poll `future` to completion on the calling thread, which must be detached from + /// Python. A current-thread runtime runs its tasks and IO here meanwhile. + pub fn block_on(&self, future: impl Future>) -> PyResult { + refuse_nested()?; + match self.inner.as_deref() { + Some(Shared::CurrentThread(runtime)) => { + let _driving = Driving::enter(); + let output = runtime.block_on(future); + INTERRUPT.take().map_or(output, Err) + } + _ => self.handle.block_on(future), + } + } + /// Share the runtime and select a worker for a new client. pub fn select(&self) -> PyResult { let inner = self @@ -34,16 +101,15 @@ impl Runtime { .as_ref() .ok_or_else(|| PyRuntimeError::new_err("Runtime is unavailable"))?; Ok(Self { + handle: inner.handle(), inner: Some(inner.clone()), - handle: inner.get_handle().clone(), }) } } -impl From for Runtime { - fn from(runtime: PingoraRuntime) -> Self { - // No-steal workers are lazy upstream. Start them before sharing the runtime. - let handle = runtime.get_handle().clone(); +impl From for Runtime { + fn from(runtime: Shared) -> Self { + let handle = runtime.handle(); Self { inner: Some(Arc::new(runtime)), handle, @@ -51,6 +117,12 @@ impl From for Runtime { } } +impl From for Runtime { + fn from(runtime: PingoraRuntime) -> Self { + Shared::Workers(runtime).into() + } +} + impl FromPyObject<'_, '_> for Runtime { type Error = PyErr; @@ -61,52 +133,76 @@ impl FromPyObject<'_, '_> for Runtime { #[pymethods] impl Runtime { - /// Create and start the runtime's workers. - /// Without work stealing, each client stays on one worker (not CPU-pinned). - /// Workers default to CPU parallelism and names to the package name. + /// Create the runtime and start its workers. + /// Workers default to CPU parallelism and do not apply to CURRENT_THREAD. + /// Thread names default to the package name. /// Thread counts must be positive; thread_keep_alive is a nonnegative timedelta. #[new] #[pyo3(signature = ( *, + scheduler = Scheduler::WORK_STEALING, workers = None, - work_steal = true, thread_name = None, max_blocking_threads = None, thread_keep_alive = None, ))] fn new( py: Python<'_>, + scheduler: Scheduler, workers: Option, - work_steal: bool, thread_name: Option<&str>, max_blocking_threads: Option, thread_keep_alive: Option, ) -> PyResult { - let workers = workers.unwrap_or_else(parallelism); - if workers == 0 - || max_blocking_threads == Some(0) - || workers - .checked_add(max_blocking_threads.unwrap_or(512)) - .is_none() - { + if max_blocking_threads == Some(0) { return Err(PyValueError::new_err("Invalid runtime thread counts")); } - let thread_name = thread_name.unwrap_or(env!("CARGO_PKG_NAME")); if thread_name.contains('\0') { return Err(PyValueError::new_err("thread_name must not contain NUL")); } - Ok(py.detach(|| { - RuntimeBuilder::new(workers, thread_name) - .work_steal(work_steal) - .blocking_pool_opts(BlockingPoolOpts { - max_threads: max_blocking_threads, - thread_keep_alive, + Ok(match scheduler { + Scheduler::CURRENT_THREAD => { + if workers.is_some() { + return Err(PyValueError::new_err( + "CURRENT_THREAD runtimes have no workers", + )); + } + let mut builder = Builder::new_current_thread(); + builder.enable_all().thread_name(thread_name); + if let Some(max_threads) = max_blocking_threads { + builder.max_blocking_threads(max_threads); + } + if let Some(keep_alive) = thread_keep_alive { + builder.thread_keep_alive(keep_alive); + } + let runtime = builder + .build() + .map_err(|err| PyRuntimeError::new_err(err.to_string()))?; + Shared::CurrentThread(runtime).into() + } + _ => { + let workers = workers.unwrap_or_else(parallelism); + if workers == 0 + || workers + .checked_add(max_blocking_threads.unwrap_or(512)) + .is_none() + { + return Err(PyValueError::new_err("Invalid runtime thread counts")); + } + py.detach(|| { + RuntimeBuilder::new(workers, thread_name) + .work_steal(scheduler == Scheduler::WORK_STEALING) + .blocking_pool_opts(BlockingPoolOpts { + max_threads: max_blocking_threads, + thread_keep_alive, + }) + .build() + .into() }) - .build() - .into() - })) + } + }) } } @@ -115,13 +211,54 @@ impl Drop for Runtime { // The final owner may be released on a worker or while holding the GIL. if let Some(runtime) = self.inner.take().and_then(Arc::into_inner) { match runtime { - PingoraRuntime::Steal { runtime, .. } => runtime.shutdown_background(), - PingoraRuntime::NoSteal(runtime) => drop(runtime), + Shared::Workers(PingoraRuntime::Steal { runtime, .. }) + | Shared::CurrentThread(runtime) => runtime.shutdown_background(), + Shared::Workers(PingoraRuntime::NoSteal(runtime)) => drop(runtime), } } } } +// ===== impl Shared ===== + +impl Shared { + fn handle(&self) -> Handle { + match self { + // No-steal workers are lazy upstream; this starts them. + Shared::Workers(runtime) => runtime.get_handle().clone(), + Shared::CurrentThread(runtime) => runtime.handle().clone(), + } + } +} + +// ===== impl Driving ===== + +impl Driving { + fn enter() -> Self { + DRIVING.set(true); + Driving + } +} + +impl Drop for Driving { + fn drop(&mut self) { + DRIVING.set(false); + // Left only by an unwinding call; a later call must not raise it. + INTERRUPT.take(); + } +} + +/// Fail a blocking wait on a thread driving a current-thread runtime, where Tokio would +/// panic, as when an upload iterator it reads sends a blocking request. +pub fn refuse_nested() -> PyResult<()> { + if driving() { + return Err(PyRuntimeError::new_err( + "Cannot make a blocking call from code a CURRENT_THREAD runtime is running", + )); + } + Ok(()) +} + /// Workers used when none are given: one per available CPU. fn parallelism() -> usize { thread::available_parallelism().map_or(1, usize::from) @@ -182,6 +319,39 @@ mod tests { // Different clients may select the same worker; each selection stays fixed. } + #[test] + fn current_thread_runs_on_its_callers() { + let current_thread = || { + Runtime::from(Shared::CurrentThread( + Builder::new_current_thread().enable_all().build().unwrap(), + )) + }; + let runtime = current_thread(); + // Spawned tasks and timers run on the blocked caller. + let child = crate::client::nogil::block_on(&runtime, async { + Ok(tokio::spawn(async { + tokio::time::sleep(Duration::from_millis(1)).await; + thread::current().id() + }) + .await + .unwrap()) + }) + .unwrap(); + assert_eq!(child, thread::current().id()); + // A blocking call from code the runtime runs fails instead of panicking, and the + // last owner of another runtime can be released there. + let other = current_thread(); + let nested = runtime + .block_on(async { + let nested = other.block_on(async { Ok(()) }).is_err(); + drop(other); + Ok(nested) + }) + .unwrap(); + assert!(nested); + assert!(!driving()); + } + #[test] fn last_owner_can_be_released_on_a_worker() { struct NotifyOnDrop(std::sync::mpsc::Sender<()>); diff --git a/tests/runtime_test.py b/tests/runtime_test.py index 7e889a06..afbde27d 100644 --- a/tests/runtime_test.py +++ b/tests/runtime_test.py @@ -1,25 +1,36 @@ import asyncio import base64 import hashlib +import threading +from contextlib import asynccontextmanager from datetime import timedelta import pytest import wreq -from wreq.runtime import Runtime +from wreq.runtime import Runtime, Scheduler from cancellation_test import local_server from upload_test import read_chunked +def custom_runtime(scheduler, **options): + """A one-worker runtime, or a current-thread one, which has no workers.""" + if scheduler != Scheduler.CURRENT_THREAD: + options["workers"] = 1 + return Runtime(scheduler=scheduler, **options) + + def test_runtime_configuration(): assert Runtime is wreq.Runtime assert Runtime is wreq.runtime.Runtime - assert "Runtime" in wreq.__all__ + assert Scheduler is wreq.Scheduler + assert {"Runtime", "Scheduler"} <= set(wreq.__all__) for kwargs in ( {"workers": 0}, {"max_blocking_threads": 0}, {"thread_name": "bad\0name"}, {"thread_keep_alive": timedelta(microseconds=-1)}, + {"scheduler": Scheduler.CURRENT_THREAD, "workers": 1}, ): with pytest.raises(ValueError): Runtime(**kwargs) @@ -29,8 +40,8 @@ def test_runtime_configuration(): wreq.Client(runtime=object()) for duration in (None, timedelta(), timedelta(microseconds=250001)): runtime = Runtime( + scheduler=Scheduler.PER_WORKER, workers=1, - work_steal=False, thread_name=None if duration is None else "isolated", max_blocking_threads=3, thread_keep_alive=duration, @@ -52,9 +63,9 @@ def test_runtime_configuration(): @pytest.mark.asyncio -@pytest.mark.parametrize("steal", [False, True]) -async def test_response_and_stream_keep_runtime_alive(steal): - runtime = wreq.Runtime(workers=1, work_steal=steal) +@pytest.mark.parametrize("scheduler", [Scheduler.PER_WORKER, Scheduler.WORK_STEALING]) +async def test_response_and_stream_keep_runtime_alive(scheduler): + runtime = custom_runtime(scheduler) async with local_server() as (url, connections): client = wreq.Client(runtime=runtime, proxies=[]) task = asyncio.create_task(client.get(url)) @@ -77,7 +88,7 @@ async def test_response_and_stream_keep_runtime_alive(steal): @pytest.mark.asyncio async def test_shared_runtime_cancellation_and_upload(): - runtime = wreq.Runtime(workers=2, work_steal=False, max_blocking_threads=2) + runtime = Runtime(scheduler=Scheduler.PER_WORKER, workers=2, max_blocking_threads=2) async with local_server() as (url, connections): first = wreq.Client(runtime=runtime, proxies=[]) second = wreq.Client(runtime=runtime, proxies=[]) @@ -114,12 +125,13 @@ async def chunks(): @pytest.mark.asyncio -@pytest.mark.parametrize("steal", [None, False, True]) -async def test_blocking_client_runtime(steal): +@pytest.mark.parametrize( + "scheduler", + [None, Scheduler.PER_WORKER, Scheduler.WORK_STEALING, Scheduler.CURRENT_THREAD], +) +async def test_blocking_client_runtime(scheduler): runtime = ( - None - if steal is None - else wreq.Runtime(workers=1, work_steal=steal, max_blocking_threads=2) + None if scheduler is None else custom_runtime(scheduler, max_blocking_threads=2) ) def request(url): @@ -147,6 +159,163 @@ def request(url): del task +async def run_blocking(func, *args, timeout=5): + """Run a blocking call on a daemon thread, so a call that never returns fails the + test instead of hanging executor shutdown.""" + loop = asyncio.get_running_loop() + future = loop.create_future() + + def run(): + try: + outcome = (func(*args), None) + except BaseException as error: + outcome = (None, error) + + def settle(): + if not future.done(): + if outcome[1] is None: + future.set_result(outcome[0]) + else: + future.set_exception(outcome[1]) + + loop.call_soon_threadsafe(settle) + + threading.Thread(target=run, daemon=True).start() + return await asyncio.wait_for(future, timeout) + + +@asynccontextmanager +async def path_server(): + """Answer each request with its path, sending the body after the head so reading it + must drive the runtime. `/hold` waits for `release`; `arrived` records paths.""" + arrived, release, uploads, writers = asyncio.Queue(), asyncio.Event(), [], [] + + async def accept(reader, writer): + writers.append(writer) + try: + head = await reader.readuntil(b"\r\n\r\n") + path = head.split(b" ", 2)[1] + arrived.put_nowait(path.decode()) + if b"transfer-encoding: chunked" in head.lower(): + uploads.append(await read_chunked(reader)) + if path == b"/hold": + await release.wait() + writer.write( + b"HTTP/1.1 200 OK\r\nConnection: close\r\n" + b"Content-Length: %d\r\n\r\n" % len(path) + ) + await writer.drain() + await asyncio.sleep(0.05) + writer.write(path) + await writer.drain() + except (asyncio.IncompleteReadError, ConnectionError, ValueError): + # Interrupted uploads end mid-body. + pass + finally: + writer.close() + + server = await asyncio.start_server(accept, "127.0.0.1", 0) + try: + url = f"http://127.0.0.1:{server.sockets[0].getsockname()[1]}" + yield url, arrived, release, uploads + finally: + release.set() + server.close() + for writer in writers: + writer.close() + await server.wait_closed() + + +@pytest.mark.asyncio +async def test_current_thread_runtime(): + runtime = Runtime(scheduler=Scheduler.CURRENT_THREAD) + with pytest.raises(ValueError): + wreq.Client(runtime=runtime) + client = wreq.blocking.Client(runtime=runtime, proxies=[]) + + class Interrupt(BaseException): + pass + + def fetch(path): + with client.get(url + path) as response: + return response.text() + + def interrupted(): + yield b"x" + raise Interrupt + + def upload_interrupted(): + client.post(url + "/interrupted", body=interrupted()) + + def run(): + caller, readers = threading.get_ident(), [] + pending = client.get(url + "/pending") + streamed = client.get(url + "/streamed").stream() + + def chunks(): + # Read on the thread driving the runtime, which refuses blocking calls + # before they consume anything. + readers.append(threading.get_ident()) + yield b"sync" + for read in (lambda: fetch("/nested"), pending.text, pending.bytes): + with pytest.raises(RuntimeError, match="CURRENT_THREAD"): + read() + with pytest.raises(RuntimeError, match="CURRENT_THREAD"): + next(streamed) + yield b" upload" + + with client.post(url + "/upload", body=chunks()) as response: + uploaded = response.text() + # Nothing would drive an async read, so it is refused up front. + for entry in (streamed.__aiter__, streamed.__anext__, streamed.__aenter__): + with pytest.raises(RuntimeError): + entry() + + # A KeyboardInterrupt-like error from the caller's own iterator surfaces as + # itself. + with pytest.raises(Interrupt): + upload_interrupted() + return readers == [caller], uploaded, pending.text(), b"".join(streamed) + + try: + async with path_server() as (url, arrived, release, uploads): + assert await run_blocking(run) == ( + True, + "/upload", + "/pending", + b"/streamed", + ) + assert uploads == [b"sync upload"] + seen = [] + while not seen or seen[-1] != "/interrupted": + seen.append(await asyncio.wait_for(arrived.get(), 5)) + assert "/nested" not in seen + + # Threads sharing the runtime take turns driving it, so a request waiting + # on the server does not hold up another thread's request. + held = asyncio.create_task(run_blocking(fetch, "/hold")) + assert await asyncio.wait_for(arrived.get(), 5) == "/hold" + assert await run_blocking(fetch, "/other") == "/other" + # This upload's iterator runs on the held thread, which keeps its own + # result; the uploading call gets the wrapped error. + with pytest.raises(wreq.exceptions.RequestError, match="Interrupt"): + await run_blocking(upload_interrupted) + assert not held.done() + release.set() + assert await held == "/hold" + + # Closing the client from another thread ends a call driving the runtime. + release.clear() + held = asyncio.create_task(run_blocking(fetch, "/hold")) + while await asyncio.wait_for(arrived.get(), 5) != "/hold": + pass + client.close() + with pytest.raises(asyncio.CancelledError): + await held + finally: + client.close() + + @pytest.mark.asyncio @pytest.mark.parametrize( ("blocking", "operation"), @@ -155,7 +324,7 @@ def request(url): async def test_client_close_cancels_requests(blocking, operation): factory = wreq.blocking.Client if blocking else wreq.Client client = factory( - runtime=Runtime(workers=1, work_steal=False), + runtime=custom_runtime(Scheduler.PER_WORKER), proxies=[], timeout=timedelta(seconds=10), ) @@ -196,8 +365,15 @@ def start(): @pytest.mark.asyncio -@pytest.mark.parametrize("blocking", [False, True]) -async def test_websocket_outlives_client(blocking): +@pytest.mark.parametrize( + ("blocking", "scheduler"), + [ + (False, Scheduler.PER_WORKER), + (True, Scheduler.PER_WORKER), + (True, Scheduler.CURRENT_THREAD), + ], +) +async def test_websocket_outlives_client(blocking, scheduler): connections = asyncio.Queue() async def accept(reader, writer): @@ -218,7 +394,7 @@ async def accept(reader, writer): connections.put_nowait((reader, writer)) server = await asyncio.start_server(accept, "127.0.0.1", 0) - runtime = wreq.Runtime(workers=1, work_steal=False) + runtime = custom_runtime(scheduler) writer = None try: url = f"ws://127.0.0.1:{server.sockets[0].getsockname()[1]}/" @@ -226,7 +402,7 @@ async def accept(reader, writer): runtime=runtime, proxies=[] ) ws = ( - await asyncio.to_thread(client.websocket, url) + await run_blocking(client.websocket, url) if blocking else await client.websocket(url) ) @@ -235,11 +411,11 @@ async def accept(reader, writer): del client, runtime writer.write(b"\x81\x04pong") await writer.drain() - message = await asyncio.to_thread(ws.recv) if blocking else await ws.recv() + message = await run_blocking(ws.recv) if blocking else await ws.recv() assert message.text == "pong" outgoing = wreq.Message.from_text("ping") if blocking: - await asyncio.to_thread(ws.send, outgoing) + await run_blocking(ws.send, outgoing) else: await ws.send(outgoing) frame = await asyncio.wait_for(reader.readexactly(10), 5) @@ -249,7 +425,7 @@ async def accept(reader, writer): == b"ping" ) if blocking: - await asyncio.to_thread(ws.close) + await run_blocking(ws.close) else: await ws.close() del ws @@ -262,8 +438,8 @@ async def accept(reader, writer): @pytest.mark.asyncio -@pytest.mark.parametrize("steal", [False, True]) -async def test_http2_multiplexing_on_custom_runtime(steal): +@pytest.mark.parametrize("scheduler", [Scheduler.PER_WORKER, Scheduler.WORK_STEALING]) +async def test_http2_multiplexing_on_custom_runtime(scheduler): # Minimal h2c responder: indexed :status=200, then a two-byte DATA frame. # No HPACK decoder is needed because the test does not inspect request headers. def frame(kind, flags, stream, payload=b""): @@ -307,7 +483,7 @@ async def accept(reader, writer): finally: handlers.discard(asyncio.current_task()) - runtime = wreq.Runtime(workers=2, work_steal=steal) + runtime = Runtime(scheduler=scheduler, workers=2) server = await asyncio.start_server(accept, "127.0.0.1", 0) try: url = f"http://127.0.0.1:{server.sockets[0].getsockname()[1]}/" diff --git a/tests/sync_upload_test.py b/tests/sync_upload_test.py index 3d35043d..f6b69bd6 100644 --- a/tests/sync_upload_test.py +++ b/tests/sync_upload_test.py @@ -6,6 +6,7 @@ import pytest import wreq +from wreq.runtime import Runtime, Scheduler @asynccontextmanager @@ -65,12 +66,24 @@ async def accept(reader, writer): assert not errors, errors +def upload_client(blocking, current_thread): + """A client whose iterators run on the blocking pool, or inline on a current-thread + runtime's driving thread.""" + if current_thread: + return wreq.blocking.Client( + proxies=[], runtime=Runtime(scheduler=Scheduler.CURRENT_THREAD) + ) + return (wreq.blocking.Client if blocking else wreq.Client)(proxies=[]) + + +UPLOAD_CLIENTS = [(False, False), (True, False), (True, True)] + + @pytest.mark.asyncio -@pytest.mark.parametrize("blocking", [False, True]) +@pytest.mark.parametrize(("blocking", "current_thread"), UPLOAD_CLIENTS) @pytest.mark.parametrize("multipart", [False, True]) -async def test_sync_upload_iterator_errors(blocking, multipart): - factory = wreq.blocking.Client if blocking else wreq.Client - client = factory(proxies=[]) +async def test_sync_upload_iterator_errors(blocking, current_thread, multipart): + client = upload_client(blocking, current_thread) try: async with upload_server() as (url, bodies): for failure in (False, True): @@ -151,8 +164,8 @@ def send(): @pytest.mark.asyncio -@pytest.mark.parametrize("blocking", [False, True]) -async def test_sync_upload_sends_each_chunk_as_yielded(blocking): +@pytest.mark.parametrize(("blocking", "current_thread"), UPLOAD_CLIENTS) +async def test_sync_upload_sends_each_chunk_as_yielded(blocking, current_thread): # A chunk goes out once yielded, before the iterator produces the next one. first_seen = threading.Event() bodies = asyncio.Queue() @@ -182,8 +195,7 @@ def chunks(): server = await asyncio.start_server(accept, "127.0.0.1", 0) url = f"http://127.0.0.1:{server.sockets[0].getsockname()[1]}/" - factory = wreq.blocking.Client if blocking else wreq.Client - client = factory(proxies=[]) + client = upload_client(blocking, current_thread) try: if blocking: @@ -254,6 +266,50 @@ def send(): await server.wait_closed() +@pytest.mark.asyncio +async def test_current_thread_upload_sees_an_early_response(): + # An inline iterator yields to the runtime before each item, so a response sent + # before the body is read ends the call within an item or two. + writers, calls = [], [] + + async def accept(reader, writer): + writers.append(writer) + try: + await asyncio.wait_for(reader.readuntil(b"\r\n\r\n"), 5) + writer.write(b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\n\r\nok") + await writer.drain() + await asyncio.wait_for(reader.read(), 10) + except (asyncio.TimeoutError, ConnectionError): + pass + finally: + writer.close() + + def chunks(): + for _ in range(400): + calls.append(None) + time.sleep(0.005) + yield b"x" + + server = await asyncio.start_server(accept, "127.0.0.1", 0) + url = f"http://127.0.0.1:{server.sockets[0].getsockname()[1]}/" + client = upload_client(True, True) + try: + + def send(): + timeout = datetime.timedelta(seconds=1) + with client.post(url, body=chunks(), timeout=timeout) as response: + return response.status + + assert await asyncio.wait_for(asyncio.to_thread(send), 5) == 200 + assert len(calls) < 20, len(calls) + finally: + client.close() + server.close() + for writer in writers: + writer.close() + await server.wait_closed() + + @pytest.mark.asyncio async def test_blocking_upload_does_not_wait_for_the_iterator(): # A response or timeout must not wait for a `__next__` call that blocks. From 860b99a1116bd4f9969fbaf040dd6c0d8db8f5ad Mon Sep 17 00:00:00 2001 From: 0x676e67 Date: Tue, 6 Oct 2026 12:44:23 +0000 Subject: [PATCH 04/25] feat(bench): add a blocking CT variant with a current-thread runtime per client --- bench/README.md | 15 +++++++++------ bench/blocking_clients.py | 17 ++++++++++++++--- bench/charts.py | 2 ++ bench/registry.py | 8 ++++++++ bench/report.py | 2 ++ bench/results.py | 13 ++++++++----- bench/tests/test_run.py | 25 +++++++++++++++++-------- docs/templates/benchmark.md | 8 +++++--- 8 files changed, 65 insertions(+), 25 deletions(-) diff --git a/bench/README.md b/bench/README.md index 33c4492d..f8096c67 100644 --- a/bench/README.md +++ b/bench/README.md @@ -207,12 +207,15 @@ shows each run's revision and collection time separately. `wreq_blocking` uses the default shared multi-thread network runtime (MT). `wreq_blocking_st` shares one `Runtime(scheduler=Scheduler.PER_WORKER, workers=1)` across all logical workers -in each case (ST). Both retain one client per logical -worker and the same Python thread pool for concurrent blocking calls. - -The default full run includes both async and blocking wreq MT/ST, for 17 -variants and 1,792 supported cells. No `--clients` option is needed. Historical -published snapshots remain readable; new full/blocking publication requires ST. +in each case (ST). `wreq_blocking_ct` gives each logical worker its own +`Runtime(scheduler=Scheduler.CURRENT_THREAD)`, driven by the calling thread +(CT). All three retain one client per logical worker and the same Python thread +pool for concurrent blocking calls. + +The default full run includes async wreq MT/ST and blocking wreq MT/ST/CT, for +18 variants and 1,904 supported cells. No `--clients` option is needed. +Historical published snapshots remain readable; new full/blocking publication +requires ST and CT. CT needs a wreq release with `Scheduler`. ## Maintaining the suite diff --git a/bench/blocking_clients.py b/bench/blocking_clients.py index 3c2bad8c..e5b1fbb4 100644 --- a/bench/blocking_clients.py +++ b/bench/blocking_clients.py @@ -100,8 +100,18 @@ def operations(client_id, protocol, url, body, body_kind, *, chunks=None, runtim from wreq.blocking import Client from wreq.tls import TlsVersion - if SPECS[client_id].runtime["kind"] == "custom" and runtime is None: + kind = SPECS[client_id].runtime["kind"] + if kind == "custom" and runtime is None: raise ValueError("Blocking ST requires one runtime shared by the case") + if kind == "current_thread": + from wreq.runtime import Runtime + + try: + from wreq.runtime import Scheduler + except ImportError as error: + raise ValueError("Blocking CT requires wreq with Scheduler") from error + # Each logical worker drives its own runtime, like one curl handle each. + runtime = Runtime(scheduler=Scheduler.CURRENT_THREAD) with Client( runtime=runtime, tls_verify=False, @@ -117,8 +127,9 @@ def post(): response = client.post( url, body=iter(chunks) if body_kind == "stream" else body ) - status, version = response.status.as_int(), native_protocol( - response.version + status, version = ( + response.status.as_int(), + native_protocol(response.version), ) total = 0 with response.stream() as stream: diff --git a/bench/charts.py b/bench/charts.py index 6d45ea4d..66886265 100644 --- a/bench/charts.py +++ b/bench/charts.py @@ -24,6 +24,7 @@ "track": "#edf1f8", "wreq": "#4657d9", "wreq_st": "#0f8a76", + "wreq_ct": "#b45309", "peer": "#8393aa", }, "dark": { @@ -34,6 +35,7 @@ "track": "#223047", "wreq": "#a0aaff", "wreq_st": "#55ccb3", + "wreq_ct": "#f6b14b", "peer": "#8b9db8", }, } diff --git a/bench/registry.py b/bench/registry.py index 48886bb3..3d433b50 100644 --- a/bench/registry.py +++ b/bench/registry.py @@ -122,6 +122,14 @@ def settings(self): }, color="wreq_st", ), + "wreq_blocking_ct": ClientSpec( + "wreq", + "wreq (blocking CT)", + "blocking", + response_read="stream(): native transport chunks", + runtime={"kind": "current_thread", "scope": "one per client"}, + color="wreq_ct", + ), "ry_blocking": ClientSpec( "ry", "ry (blocking)", diff --git a/bench/report.py b/bench/report.py index 77a66bf9..4caea682 100644 --- a/bench/report.py +++ b/bench/report.py @@ -42,6 +42,8 @@ def runtime_label(runtime: dict) -> str: return "Single-thread runtime" if kind == "multi_thread": return "Multi-thread runtime" + if kind == "current_thread": + return "Current-thread runtime per client" if kind == "custom": workers = runtime.get("workers", "not recorded") if "scheduler" in runtime: diff --git a/bench/results.py b/bench/results.py index e75c05fa..06a427f3 100644 --- a/bench/results.py +++ b/bench/results.py @@ -44,11 +44,14 @@ def require_publish(config, *, blocking=False, allow_legacy_snapshot=False): """Require the complete default suite or its independent blocking matrix.""" expected_clients = BLOCKING_CLIENTS if blocking else CLIENTS if allow_legacy_snapshot: - legacy_clients = tuple( - client for client in expected_clients if client != "wreq_blocking_st" - ) - if set(config.get("clients", ())) == set(legacy_clients): - expected_clients = legacy_clients + # Snapshots from before blocking CT, and from before blocking ST as well. + for missing in ({"wreq_blocking_ct"}, {"wreq_blocking_st", "wreq_blocking_ct"}): + legacy_clients = tuple( + client for client in expected_clients if client not in missing + ) + if set(config.get("clients", ())) == set(legacy_clients): + expected_clients = legacy_clients + break axes = { "clients": expected_clients, "protocols": ("h1", "h2"), diff --git a/bench/tests/test_run.py b/bench/tests/test_run.py index 6befb367..d508352c 100644 --- a/bench/tests/test_run.py +++ b/bench/tests/test_run.py @@ -44,7 +44,7 @@ def candidate(tmp_path, full=False, blocking=False): def test_separate_publication_gates_reject_subsets_and_weak_batches(): - assert len(BLOCKING_CLIENTS) == 8 and len(CLIENTS) == 17 + assert len(BLOCKING_CLIENTS) == 9 and len(CLIENTS) == 18 for blocking in (False, True): config = publication_config(blocking) require_publish(config, blocking=blocking) @@ -76,15 +76,23 @@ def test_separate_publication_gates_reject_subsets_and_weak_batches(): def test_default_suite_includes_wreq_mt_st_and_reads_legacy_snapshots(): args = benchmark.parse_args(["--server", "server"]) - assert {"wreq", "wreq_st", "wreq_blocking", "wreq_blocking_st"} <= set(args.clients) - assert len(args.clients) == 17 + assert { + "wreq", + "wreq_st", + "wreq_blocking", + "wreq_blocking_st", + "wreq_blocking_ct", + } <= set(args.clients) + assert len(args.clients) == 18 assert args.requests == args.warmup_requests == 300 for blocking in (False, True): legacy = publication_config(blocking) - legacy["clients"].remove("wreq_blocking_st") - require_publish(legacy, blocking=blocking, allow_legacy_snapshot=True) - with pytest.raises(ValueError): - require_publish(legacy, blocking=blocking) + # Snapshots from before CT, then from before ST as well. + for client in ("wreq_blocking_ct", "wreq_blocking_st"): + legacy["clients"].remove(client) + require_publish(legacy, blocking=blocking, allow_legacy_snapshot=True) + with pytest.raises(ValueError): + require_publish(legacy, blocking=blocking) legacy["clients"].remove("wreq_blocking") with pytest.raises(ValueError): require_publish(legacy, blocking=blocking, allow_legacy_snapshot=True) @@ -216,7 +224,8 @@ def test_full_publish_freezes_existing_blocking_selection( latest.write_bytes(b"previous full selection") latest_blocking = tmp_path / "latest-blocking.json" legacy_configuration = publication_config(True) - legacy_configuration["clients"].remove("wreq_blocking_st") + for client in ("wreq_blocking_st", "wreq_blocking_ct"): + legacy_configuration["clients"].remove(client) blocking_raw = json.dumps(make_document(legacy_configuration), indent=3).encode() latest_blocking.write_bytes(blocking_raw) monkeypatch.setattr(run, "LATEST", latest) diff --git a/docs/templates/benchmark.md b/docs/templates/benchmark.md index b81e66d7..5de10013 100644 --- a/docs/templates/benchmark.md +++ b/docs/templates/benchmark.md @@ -64,9 +64,11 @@ available CPU parallelism. `wreq (ST)` uses to the Rust runtime, not Python threads. `ry` uses its default runtime; the two `pyreqwest` variants use single-threaded and multithreaded runtimes. -New runs also include `wreq (blocking MT)` and `wreq (blocking ST)`. Blocking -ST shares one single-worker runtime across every client in a case; the Python -thread pool still supplies concurrent blocking calls. Historical charts only +New runs also include `wreq (blocking MT)`, `wreq (blocking ST)` and +`wreq (blocking CT)`. Blocking ST shares one single-worker runtime across every +client in a case; blocking CT gives each client its own current-thread runtime, +driven by the calling thread. The Python thread pool still supplies concurrent +blocking calls. Historical charts only show variants present in their recorded data. Blocking clients run in a persistent thread pool. Each worker has its own client From 38484b2d7a7af23a085fb6fbdde6bcccac182bd7 Mon Sep 17 00:00:00 2001 From: 0x676e67 Date: Tue, 6 Oct 2026 13:59:42 +0000 Subject: [PATCH 05/25] perf(aio): read ready HTTP/1 stream frames on the loop and budget async uploads --- docs/source/guide/advanced.md | 15 ++-- src/client/body/stream.rs | 132 ++++++++++++++++++++++++++++------ src/client/resp.rs | 15 ++-- src/client/resp/stream.rs | 25 +++---- tests/response_test.py | 22 ++++++ tests/shutdown_test.py | 3 +- tests/upload_test.py | 18 +++++ 7 files changed, 181 insertions(+), 49 deletions(-) diff --git a/docs/source/guide/advanced.md b/docs/source/guide/advanced.md index 4f56ceff..7f9f9159 100644 --- a/docs/source/guide/advanced.md +++ b/docs/source/guide/advanced.md @@ -13,7 +13,8 @@ Their exceptions fail the request. When an upload ends, the generator is closed; cancelling or dropping the upload schedules producer cancellation and cleanup on that loop. Keep the loop running until generator cleanup has finished. This also applies to async multipart parts. Construct async-generator `Part` objects inside a running event loop; their producers -start at construction, with bounded buffering before the request consumes them. +start at construction. A producer runs at most 256 KiB, or 64 small chunks, ahead of the +upload. Use synchronous iterators for blocking uploads. A blocking call on the producer's event-loop thread prevents async generators from progressing. @@ -128,11 +129,13 @@ client = Client(runtime=runtime) With `Scheduler.PER_WORKER`, each client is assigned one worker for its lifetime; requests, response reads, streams and WebSocket operations use that -worker. Async reads of small HTTP/1 -bodies that have already arrived finish on the event loop thread instead. With -multiple workers, newly created clients select a worker randomly and keep that -selection. This is not CPU pinning. Sharing the same `Runtime` between clients -is supported, and `client.runtime` returns the shared runtime object. +worker. Async reads of HTTP/1 data that has already arrived, and is not +content-encoded, finish on the event loop thread instead: `bytes()` and `text()` +up to 64 KiB, `json()` up to 8 KiB, and the first frame of a stream whose length +is known; later frames are read on the worker. With multiple workers, newly +created clients select a worker randomly and keep that selection. This is not CPU +pinning. Sharing the same `Runtime` between clients is supported, and +`client.runtime` returns the shared runtime object. `workers=None` uses the available CPU parallelism, or 1 if it cannot be determined; `Scheduler.CURRENT_THREAD` requires it. diff --git a/src/client/body/stream.rs b/src/client/body/stream.rs index ea9e43e6..ae4a3060 100644 --- a/src/client/body/stream.rs +++ b/src/client/body/stream.rs @@ -15,6 +15,7 @@ use std::{ use bytes::Bytes; use futures_util::Stream; use pyo3::{ + IntoPyObjectExt, exceptions::{PyException, PyKeyboardInterrupt, PyRuntimeError, PyStopIteration}, intern, prelude::*, @@ -23,7 +24,10 @@ use pyo3::{ }; use tokio::{ runtime::Handle, - sync::mpsc::{self, error::TrySendError}, + sync::{ + OwnedSemaphorePermit, Semaphore, TryAcquireError, + mpsc::{self, error::TrySendError}, + }, task::{spawn_blocking, yield_now}, time, }; @@ -93,20 +97,35 @@ struct Pump { tx: mpsc::Sender, } -/// A request body from a Python async generator, forwarded with one chunk of buffering -/// by a task on the loop that was running at extraction. Dropping it cancels that task. +/// Bytes an async upload may queue ahead of the connection. +const UPLOAD_BUDGET: usize = 256 * 1024; + +/// The least budget a queued chunk holds, so small chunks queue at most 64 items. +const UPLOAD_CHARGE: usize = 4 * 1024; + +/// A request body from a Python async generator, forwarded by a task on the loop that was +/// running at extraction, up to [`UPLOAD_BUDGET`] ahead. Dropping it cancels that task. struct PyAsyncStream { - rx: mpsc::Receiver, + rx: mpsc::UnboundedReceiver, + budget: Arc, /// Set by [`Sender::finish`], so the channel closing reads as the end of the body. finished: Arc, task: Option<(Py, Py)>, } -/// The channel end given to the forwarding coroutine; awaiting `send` applies upload -/// backpressure. Closing it without `finish` fails the body. +/// A queued item and the share of the upload budget it holds until the body takes it. +struct Queued { + item: Item, + _share: Option, +} + +/// The channel end given to the forwarding coroutine: `try_send` queues a chunk within the +/// budget or waits for room, and `send` also queues the error that ends the body. Closing +/// it without `finish` fails the body. #[pyclass(frozen)] struct Sender { - tx: Mutex>>, + tx: Mutex>>, + budget: Arc, finished: Arc, } @@ -125,6 +144,18 @@ impl FromPyObject<'_, '_> for PyBytesLike { } } +impl PyBytesLike { + /// The share of the upload budget this chunk holds while queued; one larger than the + /// budget takes all of it. + fn share(&self) -> u32 { + let len = match self { + PyBytesLike::Bytes(b) => b.0.len(), + PyBytesLike::String(s) => s.0.len(), + }; + u32::try_from(len.clamp(UPLOAD_CHARGE, UPLOAD_BUDGET)).unwrap_or(u32::MAX) + } +} + impl From for Bytes { #[inline] fn from(value: PyBytesLike) -> Self { @@ -328,7 +359,11 @@ async def forward(gen, sender): try: try: async for item in gen: - if not await sender.send(item, False): + # Queued at once while the upload budget has room, otherwise awaited. + sent = sender.try_send(item) + if not isinstance(sent, bool): + sent = await sent + if not sent: return finally: close = getattr(gen, 'aclose', None) @@ -353,10 +388,12 @@ async def forward(gen, sender): .getattr("forward") .map(Bound::unbind) })?; - let (tx, rx) = mpsc::channel(1); + let (tx, rx) = mpsc::unbounded_channel(); + let budget = Arc::new(Semaphore::new(UPLOAD_BUDGET)); let finished = Arc::new(AtomicBool::new(false)); let sender = Sender { tx: Mutex::new(Some(tx)), + budget: budget.clone(), finished: finished.clone(), }; let coroutine = forward.bind(py).call1((generator, sender))?; @@ -368,6 +405,7 @@ async def forward(gen, sender): })?; Ok(Self { rx, + budget, finished, task: Some((task.unbind(), event_loop.unbind())), }) @@ -380,7 +418,8 @@ impl Stream for PyAsyncStream { fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { let this = self.get_mut(); match this.rx.poll_recv(cx) { - Poll::Ready(Some(item)) => Poll::Ready(Some(item)), + // Taking the item returns its share of the budget. + Poll::Ready(Some(queued)) => Poll::Ready(Some(queued.item)), // Every sender is gone: the end after `finish`. Otherwise the forwarding task was // cancelled or destroyed, so the body is incomplete. Poll::Ready(None) @@ -399,6 +438,8 @@ impl Stream for PyAsyncStream { impl Drop for PyAsyncStream { fn drop(&mut self) { self.rx.close(); + // A send waiting for room resolves to `False` at once. + self.budget.close(); if let Some((task, event_loop)) = self.task.take() { // Body drop can run on Tokio: cancel from a blocking thread, preferring the current // runtime so a drop on a client's own runtime does not start the shared one. @@ -422,8 +463,34 @@ impl Drop for PyAsyncStream { #[pymethods] impl Sender { - /// Queue a chunk, or `item` as the error that ends the body. Resolves to `False` - /// once the body is dropped or the sender closed, which stops forwarding. + /// Queue a chunk: `True` once queued within the budget, `False` once the body is dropped + /// or the sender closed, and otherwise a coroutine that queues it once there is room. + fn try_send<'py>(&self, item: Bound<'py, PyAny>) -> PyResult> { + let py = item.py(); + let chunk = item.extract::()?; + let tx = self.lock().clone(); + let sent = match &tx { + None => false, + Some(queue) => match self.budget.clone().try_acquire_many_owned(chunk.share()) { + Ok(share) => queue + .send(Queued { + item: Ok(chunk), + _share: Some(share), + }) + .is_ok(), + Err(TryAcquireError::Closed) => false, + // The coroutine keeps the extracted chunk, so it is never copied twice. + Err(TryAcquireError::NoPermits) => { + return self.send_later(py, tx, Ok(chunk)).map(Bound::into_any); + } + }, + }; + sent.into_bound_py_any(py) + } + + /// Queue `item` as the error that ends the body, or a chunk once the budget has room. + /// Resolves to `False` once the body is dropped or the sender closed, which stops + /// forwarding. fn send<'py>( &self, py: Python<'py>, @@ -435,14 +502,7 @@ impl Sender { } else { Ok(item.extract()?) }; - let tx = self.lock().clone(); - // Channel readiness is runtime-independent, so this waits on the Python loop. - coroutine::local(py, "Sender.send", async move { - Ok(match tx { - Some(tx) => tx.send(item).await.is_ok(), - None => false, - }) - }) + self.send_later(py, self.lock().clone(), item) } /// Mark the normal end of the body. The last chunk may still be queued: it is read @@ -462,7 +522,37 @@ impl Sender { impl Sender { #[inline] - fn lock(&self) -> MutexGuard<'_, Option>> { + fn lock(&self) -> MutexGuard<'_, Option>> { self.tx.lock().unwrap_or_else(PoisonError::into_inner) } + + /// A coroutine that queues `item` on `tx` once the budget has room; an error, which ends + /// the body, never waits. + fn send_later<'py>( + &self, + py: Python<'py>, + tx: Option>, + item: Item, + ) -> PyResult> { + let budget = self.budget.clone(); + // Budget readiness is runtime-independent, so this waits on the Python loop. + coroutine::local(py, "Sender.send", async move { + let Some(tx) = tx else { + return Ok(false); + }; + let share = match &item { + Ok(chunk) => match budget.acquire_many_owned(chunk.share()).await { + Ok(share) => Some(share), + Err(_) => return Ok(false), + }, + Err(_) => None, + }; + Ok(tx + .send(Queued { + item, + _share: share, + }) + .is_ok()) + }) + } } diff --git a/src/client/resp.rs b/src/client/resp.rs index 0312082f..edced877 100644 --- a/src/client/resp.rs +++ b/src/client/resp.rs @@ -13,10 +13,15 @@ pub use self::{ /// blocking caller without first releasing the GIL. const READ_ATTACHED: u64 = 64 * 1024; -/// The largest unread body the event loop polls for a response of `version`: `limit` up to -/// HTTP/1.1, and none after it, not even an empty one. An HTTP/2 stream shares its -/// connection's state behind a lock that the connection task holds while it handles frames, -/// so polling one would stall the loop. +/// Whether the event loop may poll an unread body of `version`: up to HTTP/1.1, whose body +/// frames the connection task hands over, and never after it, not even an empty body. An +/// HTTP/2 stream shares its connection's state behind a lock that the connection task holds +/// while it handles frames, so polling one would stall the loop. +fn loop_polls(version: wreq::Version) -> bool { + version <= wreq::Version::HTTP_11 +} + +/// The largest unread body the event loop collects for a response of `version`. fn loop_limit(version: wreq::Version, limit: u64) -> Option { - (version <= wreq::Version::HTTP_11).then_some(limit) + loop_polls(version).then_some(limit) } diff --git a/src/client/resp/stream.rs b/src/client/resp/stream.rs index 46530778..618e3171 100644 --- a/src/client/resp/stream.rs +++ b/src/client/resp/stream.rs @@ -22,7 +22,7 @@ use tokio::sync::{ }; use tokio_util::task::AbortOnDropHandle; -use super::{READ_ATTACHED, loop_limit}; +use super::loop_polls; use crate::{ buffer::PyBuffer, client::nogil, @@ -159,21 +159,15 @@ impl Streamer { } /// Read a frame the body already holds, before any read-ahead task starts, so a reader - /// of a short body never waits on another thread. With `limit`, only a body of known - /// length up to it is read here, so an async reader never reads a large or decompressing - /// body on the event loop; `end` builds the error that ends iteration. - fn ready_frame(&self, limit: Option, end: fn() -> Error) -> Option> { + /// of a short body never waits on another thread; `end` builds the error that ends + /// iteration. On the event loop only an HTTP/1 body of known length is read here, whatever + /// its size: taking a frame it holds is cheap, while a decompressing body is not. + fn ready_frame(&self, on_loop: bool, end: fn() -> Error) -> Option> { let mut state = self.reader.lock(); let State::Idle(resp) = &mut *state else { return None; }; - if let Some(limit) = limit - && !resp - .size_hint() - .exact() - .zip(loop_limit(resp.version(), limit)) - .is_some_and(|(len, limit)| len <= limit) - { + if on_loop && !(loop_polls(resp.version()) && resp.size_hint().exact().is_some()) { return None; } let _runtime = self.runtime.handle().enter(); @@ -290,7 +284,7 @@ impl Streamer { // Refused before any frame, so a nested read never leaves a stream half consumed. runtime::refuse_nested()?; // Frames already received are returned without releasing the GIL. - if let Some(frame) = self.ready_frame(None, || Error::StopIteration) { + if let Some(frame) = self.ready_frame(false, || Error::StopIteration) { return frame; } nogil::run(py, &self.runtime, self.next(|| Error::StopIteration)) @@ -325,9 +319,8 @@ impl Streamer { let slf = slf.unbind(); coroutine::local(py, "Streamer.__anext__", async move { let this = slf.get(); - // A short body already buffered is read without starting the read-ahead task. - if let Some(frame) = this.ready_frame(Some(READ_ATTACHED), || Error::StopAsyncIteration) - { + // A body already buffered is read without starting the read-ahead task. + if let Some(frame) = this.ready_frame(true, || Error::StopAsyncIteration) { return frame; } // Buffered frames complete without suspending; yield to the event loop diff --git a/tests/response_test.py b/tests/response_test.py index 1dfbb405..bd89ba4d 100644 --- a/tests/response_test.py +++ b/tests/response_test.py @@ -348,6 +348,28 @@ async def test_only_http1_bodies_are_read_on_the_event_loop(http2, body, read): assert (suspended > 0) == http2 +@pytest.mark.asyncio +async def test_streams_read_large_http1_bodies_on_the_event_loop(): + # A stream takes frames that an HTTP/1 body of known length already holds on the loop, + # whatever its size; bytes() collects bodies over 64 KiB on the runtime. + body = b"x" * (128 << 10) + async with body_server(False, body) as url: + async with wreq.Client(proxies=[]) as client: + response = await asyncio.wait_for(client.get(url), 5) + await asyncio.sleep(0.05) + result, suspended = await steps(response.stream().__anext__()) + assert body.startswith(bytes(result)) and suspended == 0 + # The runtime may finish one collect before the first poll on a loaded host. + suspensions = [] + for _ in range(3): + response = await asyncio.wait_for(client.get(url), 5) + await asyncio.sleep(0.05) + result, suspended = await steps(response.bytes()) + assert bytes(result) == body + suspensions.append(suspended) + assert any(suspensions) + + @pytest.mark.asyncio async def test_stream_read_ahead_yields_and_closes_waiting_readers(): stalled = threading.Event() diff --git a/tests/shutdown_test.py b/tests/shutdown_test.py index 6a15f5e0..03b2e61d 100644 --- a/tests/shutdown_test.py +++ b/tests/shutdown_test.py @@ -118,7 +118,8 @@ async def chunks(): try: await asyncio.wait_for(full.wait(), 5) await asyncio.sleep(0.05) - assert produced == [0, 1] + # 64 small chunks fill the upload budget at its minimum charge; the next waits. + assert produced == list(range(65)) if retain: return part finally: diff --git a/tests/upload_test.py b/tests/upload_test.py index 8bed63c8..14ccb303 100644 --- a/tests/upload_test.py +++ b/tests/upload_test.py @@ -95,6 +95,24 @@ async def chunks(): await (await asyncio.wait_for(task, 5)).close() +@pytest.mark.asyncio +async def test_async_upload_beyond_the_budget_keeps_order(): + # 1 MiB in small chunks queues within the budget, then waits for room, in order. + parts = [bytes((index % 256,)) * 4096 for index in range(256)] + + async def chunks(): + for part in parts: + yield part + + async with local_server() as (url, connections), wreq.Client(proxies=[]) as client: + task = asyncio.create_task(client.post(url, body=chunks())) + reader, writer = await asyncio.wait_for(connections.get(), 5) + assert await asyncio.wait_for(read_chunked(reader), 10) == b"".join(parts) + writer.write(b"HTTP/1.1 200 OK\r\nContent-Length: 0\r\n\r\n") + await writer.drain() + await (await asyncio.wait_for(task, 5)).close() + + @pytest.mark.asyncio @pytest.mark.parametrize("failure", ["exception", "type", "cancelled"]) async def test_upload_errors(failure): From 6945aa3785070866ca7ee01b612aa407a7683fbc Mon Sep 17 00:00:00 2001 From: 0x676e67 Date: Tue, 6 Oct 2026 15:56:50 +0000 Subject: [PATCH 06/25] perf(upload): read sync iterators ahead within the upload budget --- docs/source/guide/advanced.md | 6 +- src/client/body/stream.rs | 110 +++++++++++++++++++++++----------- tests/sync_upload_test.py | 64 ++++++++++++++++++++ 3 files changed, 143 insertions(+), 37 deletions(-) diff --git a/docs/source/guide/advanced.md b/docs/source/guide/advanced.md index 7f9f9159..e88156d1 100644 --- a/docs/source/guide/advanced.md +++ b/docs/source/guide/advanced.md @@ -15,8 +15,10 @@ Keep the loop running until generator cleanup has finished. This also applies to Construct async-generator `Part` objects inside a running event loop; their producers start at construction. A producer runs at most 256 KiB, or 64 small chunks, ahead of the upload. -Use synchronous iterators for blocking uploads. A blocking call on the producer's -event-loop thread prevents async generators from progressing. +Use synchronous iterators for blocking uploads. They run on the runtime's blocking pool, +also at most 256 KiB or 64 small chunks ahead, except on a current-thread runtime. A +blocking call on the producer's event-loop thread prevents async generators from +progressing. ```python import asyncio diff --git a/src/client/body/stream.rs b/src/client/body/stream.rs index ae4a3060..fa766a02 100644 --- a/src/client/body/stream.rs +++ b/src/client/body/stream.rs @@ -24,10 +24,7 @@ use pyo3::{ }; use tokio::{ runtime::Handle, - sync::{ - OwnedSemaphorePermit, Semaphore, TryAcquireError, - mpsc::{self, error::TrySendError}, - }, + sync::{OwnedSemaphorePermit, Semaphore, TryAcquireError, mpsc}, task::{spawn_blocking, yield_now}, time, }; @@ -81,23 +78,26 @@ enum SyncStream { Done, } -/// An iterator read by pump tasks on the blocking pool. A pump reads an item only once -/// the channel has room, staying at most one item ahead of the upload. When the upload -/// stalls, it parks the iterator and frees its thread; the next item taken restarts it. +/// An iterator read by pump tasks on the blocking pool, up to [`UPLOAD_BUDGET`] ahead of +/// the upload. When the budget stays full, a pump parks the iterator with the chunk it read +/// and frees its thread; the next item taken restarts it. struct Pumped { - rx: mpsc::Receiver, - /// The iterator and sender while no pump runs: before the first poll, or once a pump - /// parked at a full channel. + rx: mpsc::UnboundedReceiver, + budget: Arc, + /// The pump while none runs: before the first poll, or once parked at a full budget. parked: Arc>>, } /// What a pump task needs to read the iterator into the body. struct Pump { iter: Py, - tx: mpsc::Sender, + tx: mpsc::UnboundedSender, + budget: Arc, + /// A chunk read before the pump parked, queued first when it restarts. + pending: Option, } -/// Bytes an async upload may queue ahead of the connection. +/// Bytes an upload may queue ahead of the connection. const UPLOAD_BUDGET: usize = 256 * 1024; /// The least budget a queued chunk holds, so small chunks queue at most 64 items. @@ -257,16 +257,28 @@ impl SyncStream { impl Pumped { fn new(iter: Py) -> Self { - let (tx, rx) = mpsc::channel(1); + let (tx, rx) = mpsc::unbounded_channel(); + let budget = Arc::new(Semaphore::new(UPLOAD_BUDGET)); + let pump = Pump { + iter, + tx, + budget: budget.clone(), + pending: None, + }; Pumped { rx, - parked: Arc::new(Mutex::new(Some(Pump { iter, tx }))), + budget, + parked: Arc::new(Mutex::new(Some(pump))), } } fn poll_next(&mut self, cx: &mut Context<'_>) -> Poll> { - let poll = self.rx.poll_recv(cx); - // A parked pump waits for room, which the first poll or a taken item makes. The + // Taking an item returns its share of the budget before the parked check. + let poll = self + .rx + .poll_recv(cx) + .map(|queued| queued.map(|queued| queued.item)); + // A parked pump waits for budget, which the first poll or a taken item frees. The // check follows the take, so a pump parking concurrently is seen here. let parked = lock(&self.parked).take(); if let Some(pump) = parked { @@ -277,36 +289,65 @@ impl Pumped { } } +impl Drop for Pumped { + fn drop(&mut self) { + // A pump waiting for budget stops at once, without reading another item. + self.budget.close(); + } +} + // ===== impl Pump ===== impl Pump { - /// How long a pump waits for room before parking to free its thread. + /// How long a pump waits for budget before parking to free its thread. const PARK_AFTER: Duration = Duration::from_millis(10); - /// Send items until the iterator ends or raises, the body is dropped, or the channel + /// Queue chunks until the iterator ends or raises, the body is dropped, or the budget /// stays full past [`PARK_AFTER`](Self::PARK_AFTER). - fn run(self, parked: &Mutex>) { + fn run(mut self, parked: &Mutex>) { let handle = Handle::current(); // Once Python is unavailable, stop reading without creating a PyErr that could // require another attachment to format. Python::try_attach(|py| { loop { - let permit = match self.tx.clone().try_reserve_owned() { - Ok(permit) => permit, - Err(TrySendError::Closed(_)) => return, - Err(TrySendError::Full(tx)) => { - let wait = time::timeout(Self::PARK_AFTER, tx.reserve_owned()); + if self.tx.is_closed() { + return; + } + let chunk = match self.pending.take() { + Some(chunk) => chunk, + None => match next_item(py, &self.iter) { + Some(Ok(chunk)) => chunk, + Some(Err(err)) => { + let _ = self.tx.send(Queued { + item: Err(err), + _share: None, + }); + return; + } + None => return, + }, + }; + let share = match self.budget.clone().try_acquire_many_owned(chunk.share()) { + Ok(share) => share, + Err(TryAcquireError::Closed) => return, + Err(TryAcquireError::NoPermits) => { + let budget = self.budget.clone(); + let wait = time::timeout( + Self::PARK_AFTER, + budget.acquire_many_owned(chunk.share()), + ); match py.detach(|| handle.block_on(wait)) { - Ok(Ok(permit)) => permit, + Ok(Ok(share)) => share, Ok(Err(_)) => return, Err(_) => { // Park under the lock the body takes after receiving, so - // either it sees the pump parked or the pump sees room. + // either it sees the pump parked or the pump sees budget. let mut slot = lock(parked); - match self.tx.clone().try_reserve_owned() { - Ok(permit) => permit, - Err(TrySendError::Closed(_)) => return, - Err(TrySendError::Full(_)) => { + match self.budget.clone().try_acquire_many_owned(chunk.share()) { + Ok(share) => share, + Err(TryAcquireError::Closed) => return, + Err(TryAcquireError::NoPermits) => { + self.pending = Some(chunk); *slot = Some(self); return; } @@ -315,12 +356,11 @@ impl Pump { } } }; - let Some(item) = next_item(py, &self.iter) else { - return; + let queued = Queued { + item: Ok(chunk), + _share: Some(share), }; - let failed = item.is_err(); - permit.send(item); - if failed { + if self.tx.send(queued).is_err() { return; } } diff --git a/tests/sync_upload_test.py b/tests/sync_upload_test.py index f6b69bd6..5a1c5c95 100644 --- a/tests/sync_upload_test.py +++ b/tests/sync_upload_test.py @@ -8,6 +8,8 @@ import wreq from wreq.runtime import Runtime, Scheduler +from upload_test import read_chunked + @asynccontextmanager async def upload_server(): @@ -137,6 +139,68 @@ async def send(): client.close() +@pytest.mark.asyncio +@pytest.mark.parametrize("blocking", [False, True]) +async def test_sync_upload_reads_ahead_within_the_budget(blocking): + # A huge first chunk keeps the connection's write buffer full while the peer waits, so + # only the budget bounds the read-ahead: 64 small chunks queue at the minimum charge + # and the next parks with the pump. Once the peer reads, the parked chunk goes first. + release, bodies, calls = threading.Event(), asyncio.Queue(), [] + parts = [b"\0" * (1 << 24)] + [bytes((index,)) * 4096 for index in range(100)] + + def chunks(): + for part in parts: + calls.append(None) + yield part + + async def accept(reader, writer): + try: + await asyncio.wait_for(reader.readuntil(b"\r\n\r\n"), 5) + await asyncio.to_thread(release.wait, 5) + bodies.put_nowait(await asyncio.wait_for(read_chunked(reader), 10)) + writer.write( + b"HTTP/1.1 200 OK\r\nConnection: close\r\nContent-Length: 2\r\n\r\nok" + ) + await writer.drain() + finally: + writer.close() + + server = await asyncio.start_server(accept, "127.0.0.1", 0) + url = f"http://127.0.0.1:{server.sockets[0].getsockname()[1]}/" + client = upload_client(blocking, False) + try: + if blocking: + + def send_blocking(): + with client.post(url, body=chunks()) as response: + return response.bytes() + + task = asyncio.create_task(asyncio.to_thread(send_blocking)) + else: + + async def send(): + response = await client.post(url, body=chunks()) + async with response: + return await response.bytes() + + task = asyncio.create_task(send()) + seen, stable = -1, 0 + for _ in range(250): + await asyncio.sleep(0.02) + seen, stable = len(calls), stable + 1 if len(calls) == seen else 0 + if stable == 5: + break + assert len(calls) == 66 + release.set() + assert await asyncio.wait_for(task, 10) == b"ok" + assert await asyncio.wait_for(bodies.get(), 5) == b"".join(parts) + finally: + release.set() + client.close() + server.close() + await server.wait_closed() + + @pytest.mark.asyncio async def test_blocking_upload_iterator_can_send_requests(): # The iterator runs on a blocking thread, so it may itself block on requests. From 45ec19920c990e1dfa7fd4495f500acb1da10dfa Mon Sep 17 00:00:00 2001 From: 0x676e67 Date: Wed, 7 Oct 2026 07:22:42 +0000 Subject: [PATCH 07/25] refactor(runtime): route every runtime wait through Runtime --- src/client.rs | 37 +++---- src/client/body/multipart.rs | 12 +-- src/client/nogil.rs | 50 ---------- src/client/resp/http.rs | 18 ++-- src/client/resp/stream.rs | 4 +- src/client/resp/ws.rs | 30 +++--- src/coroutine.rs | 24 +---- src/runtime.rs | 180 +++++++++++++++++++++++++++-------- 8 files changed, 186 insertions(+), 169 deletions(-) delete mode 100644 src/client/nogil.rs diff --git a/src/client.rs b/src/client.rs index 67745582..88e9e4d9 100644 --- a/src/client.rs +++ b/src/client.rs @@ -1,5 +1,4 @@ pub mod body; -pub mod nogil; pub mod req; pub mod resp; @@ -279,11 +278,10 @@ impl Client { kwds: Option>, ) -> PyResult { let kwds = Python::attach(|py| kwds.map(|kwds| kwds.bind(py).extract()).transpose())?; - coroutine::run( - self.runtime.clone(), - execute_request(self, method, url, kwds), - ) - .await + self.runtime + .clone() + .run_task(execute_request(self, method, url, kwds)) + .await } /// Open a WebSocket on the client's runtime, extracting options on first await. @@ -293,11 +291,10 @@ impl Client { kwds: Option>, ) -> PyResult { let kwds = Python::attach(|py| kwds.map(|kwds| kwds.bind(py).extract()).transpose())?; - coroutine::run( - self.runtime.clone(), - execute_websocket_request(self, url, kwds), - ) - .await + self.runtime + .clone() + .run_task(execute_websocket_request(self, url, kwds)) + .await } } @@ -834,13 +831,10 @@ impl BlockingClient { url: PyBackedStr, kwds: Option, ) -> PyResult { - py.detach(|| { - nogil::block_on( - &self.0.runtime, - execute_request(self.0.clone(), method, url, kwds), - ) + self.0 + .runtime + .block_on_task(py, execute_request(self.0.clone(), method, url, kwds)) .map(Into::into) - }) } /// Make a WebSocket request to the specified URL. @@ -851,13 +845,10 @@ impl BlockingClient { url: PyBackedStr, kwds: Option, ) -> PyResult { - py.detach(|| { - nogil::block_on( - &self.0.runtime, - execute_websocket_request(self.0.clone(), url, kwds), - ) + self.0 + .runtime + .block_on_task(py, execute_websocket_request(self.0.clone(), url, kwds)) .map(Into::into) - }) } } diff --git a/src/client/body/multipart.rs b/src/client/body/multipart.rs index 1427f52a..3a6d7cab 100644 --- a/src/client/body/multipart.rs +++ b/src/client/body/multipart.rs @@ -1,5 +1,6 @@ use std::path::PathBuf; +use futures_util::TryFutureExt; use pyo3::{prelude::*, types::PyTuple}; use wreq::{Body, multipart}; @@ -135,13 +136,10 @@ impl Part { Value::Text(text) => multipart::Part::stream(text.0), Value::Bytes(bytes) => multipart::Part::stream(bytes.0), // Opening the file blocks, so only that waits detached. - Value::File(path) => py.detach(|| { - runtime::get().block_on(async { - multipart::Part::file(path) - .await - .map_err(|err| Error::from(err).into()) - }) - })?, + Value::File(path) => runtime::get().block_on( + py, + multipart::Part::file(path).map_err(|err| Error::from(err).into()), + )?, Value::Stream(stream) => { let stream = Body::wrap_stream(stream); match self.length { diff --git a/src/client/nogil.rs b/src/client/nogil.rs deleted file mode 100644 index 4aa6c53e..00000000 --- a/src/client/nogil.rs +++ /dev/null @@ -1,50 +0,0 @@ -//! Waits for the blocking API. Network work runs on the client's [`Runtime`] while the -//! caller waits detached from Python, so other threads keep the interpreter. - -use std::{future::Future, pin::pin}; - -use futures_util::FutureExt; -use pyo3::{exceptions::PyRuntimeError, prelude::*}; - -use crate::runtime::{self, Runtime}; - -/// Run network work on its selected worker and only wait on the caller, or drive it here -/// on a current-thread runtime. The caller must be detached from Python and outside an -/// async Tokio context. Its runtime borrow keeps the workers alive until the join completes. -pub fn block_on(runtime: &Runtime, future: F) -> PyResult -where - F: Future> + Send + 'static, - T: Send + 'static, -{ - if runtime.is_current_thread() { - return runtime.block_on(future); - } - runtime.block_on(async { - runtime - .handle() - .spawn(future) - .await - .map_err(|err| PyRuntimeError::new_err(err.to_string()))? - }) -} - -/// Poll `future` on the caller instead of spawning it: return still attached when it is -/// ready at once, as a buffered body usually is, and otherwise wait detached. -/// The caller must be outside an async Tokio context. -pub fn run(py: Python<'_>, runtime: &Runtime, future: F) -> PyResult -where - F: Future> + Send, - T: Send, -{ - // Refused before the first poll, which a later refusal would leave half done. - runtime::refuse_nested()?; - let mut future = pin!(future); - let ready = { - let _runtime = runtime.handle().enter(); - future.as_mut().now_or_never() - }; - match ready { - Some(output) => output, - None => py.detach(|| runtime.block_on(future)), - } -} diff --git a/src/client/resp/http.rs b/src/client/resp/http.rs index b5cadc54..53084b57 100644 --- a/src/client/resp/http.rs +++ b/src/client/resp/http.rs @@ -6,7 +6,7 @@ use std::{ }; use bytes::Bytes; -use futures_util::{FutureExt, TryFutureExt, future::Either}; +use futures_util::{TryFutureExt, future::Either}; use http::response::{Parts, Response as HttpResponse}; use http_body::Body as _; use http_body_util::{BodyExt, Collected, combinators::Collect}; @@ -16,7 +16,7 @@ use wreq::Uri; use super::{READ_ATTACHED, ext::ResponseExt, loop_limit, stream::Streamer}; use crate::{ buffer::PyBuffer, - client::{SocketAddr, body::Json, nogil}, + client::{SocketAddr, body::Json}, cookie::Cookie, coroutine::{self, Coroutine, EntersSelf}, error::Error, @@ -136,9 +136,7 @@ impl Response { .is_some_and(|(len, limit)| len <= limit); let mut collect = body.collect(); let ready = if attached { - // A read timeout starts a timer on its first poll, which needs the runtime. - let _runtime = self.runtime.handle().enter(); - (&mut collect).now_or_never() + self.runtime.poll_now(&mut collect) } else { None }; @@ -224,7 +222,7 @@ impl Response { if inline { read.await } else { - coroutine::run(this.runtime.clone(), read).await + this.runtime.clone().run_task(read).await } }) } @@ -397,7 +395,7 @@ impl Response { .collect_later(collect) .map_ok(|(_, bytes)| PyBuffer::from(bytes)) .map_err(Into::into); - coroutine::run(this.runtime.clone(), read).await + this.runtime.clone().run_task(read).await } } }) @@ -493,9 +491,9 @@ impl BlockingResponse { let runtime = &self.0.runtime; let (read, inline) = self.0.read_body(Some(READ_ATTACHED), read)?; if inline { - nogil::run(py, runtime, read) + runtime.block_on_eager(py, read) } else { - py.detach(|| runtime.block_on(read)) + runtime.block_on(py, read) } } } @@ -594,7 +592,7 @@ impl BlockingResponse { .collect_later(collect) .map_ok(|(_, bytes)| PyBuffer::from(bytes)) .map_err(Into::into); - py.detach(|| response.runtime.block_on(read)) + response.runtime.block_on(py, read) } } } diff --git a/src/client/resp/stream.rs b/src/client/resp/stream.rs index 618e3171..692bf35d 100644 --- a/src/client/resp/stream.rs +++ b/src/client/resp/stream.rs @@ -25,7 +25,6 @@ use tokio_util::task::AbortOnDropHandle; use super::loop_polls; use crate::{ buffer::PyBuffer, - client::nogil, coroutine::{self, Coroutine}, error::Error, header::HeaderMap, @@ -287,7 +286,8 @@ impl Streamer { if let Some(frame) = self.ready_frame(false, || Error::StopIteration) { return frame; } - nogil::run(py, &self.runtime, self.next(|| Error::StopIteration)) + self.runtime + .block_on_eager(py, self.next(|| Error::StopIteration)) } fn __enter__(slf: PyRef) -> PyRef { diff --git a/src/client/resp/ws.rs b/src/client/resp/ws.rs index d83afb43..da2616d7 100644 --- a/src/client/resp/ws.rs +++ b/src/client/resp/ws.rs @@ -11,7 +11,7 @@ use pyo3::{prelude::*, sync::PyOnceLock}; use wreq::{header::HeaderValue, ws::WebSocketResponse}; use crate::{ - client::{SocketAddr, nogil}, + client::SocketAddr, cookie::Cookie, coroutine::{self, Coroutine, EntersSelf}, extractor::Text, @@ -242,33 +242,33 @@ impl BlockingWebSocket { /// Receive a message from the WebSocket. #[pyo3(signature = (timeout=None))] pub fn recv(&self, py: Python, timeout: Option) -> PyResult> { - nogil::run(py, &self.0.runtime, cmd::recv(self.0.cmd.clone(), timeout)) + self.0 + .runtime + .block_on_eager(py, cmd::recv(self.0.cmd.clone(), timeout)) } /// Send a message to the WebSocket. #[pyo3(signature = (message))] pub fn send(&self, py: Python, message: Message) -> PyResult<()> { - nogil::run(py, &self.0.runtime, cmd::send(self.0.cmd.clone(), message)) + self.0 + .runtime + .block_on_eager(py, cmd::send(self.0.cmd.clone(), message)) } /// Send multiple messages to the WebSocket. #[pyo3(signature = (messages))] pub fn send_all(&self, py: Python, messages: Vec) -> PyResult<()> { - nogil::run( - py, - &self.0.runtime, - cmd::send_all(self.0.cmd.clone(), messages), - ) + self.0 + .runtime + .block_on_eager(py, cmd::send_all(self.0.cmd.clone(), messages)) } /// Close the WebSocket connection. #[pyo3(signature = (code=None, reason=None))] pub fn close(&self, py: Python, code: Option, reason: Option) -> PyResult<()> { - nogil::run( - py, - &self.0.runtime, - cmd::close(self.0.cmd.clone(), code, reason), - ) + self.0 + .runtime + .block_on_eager(py, cmd::close(self.0.cmd.clone(), code, reason)) } } @@ -286,7 +286,9 @@ impl BlockingWebSocket { _exc_value: &Bound<'py, PyAny>, _traceback: &Bound<'py, PyAny>, ) -> PyResult<()> { - nogil::run(py, &self.0.runtime, cmd::close_on_exit(self.0.cmd.clone())) + self.0 + .runtime + .block_on_eager(py, cmd::close_on_exit(self.0.cmd.clone())) } } diff --git a/src/coroutine.rs b/src/coroutine.rs index f6131284..0397661a 100644 --- a/src/coroutine.rs +++ b/src/coroutine.rs @@ -21,10 +21,7 @@ use std::{ task::Poll, }; -use pyo3::{ - IntoPyObjectExt, PyTypeInfo, exceptions::PyRuntimeError, intern, prelude::*, sync::PyOnceLock, -}; -use tokio_util::task::AbortOnDropHandle; +use pyo3::{IntoPyObjectExt, PyTypeInfo, intern, prelude::*, sync::PyOnceLock}; pub(crate) use self::asyncio::running_loop; pub use self::awaitable::Coroutine; @@ -43,7 +40,7 @@ where F: Future> + Send + 'static, T: for<'a> IntoPyObject<'a> + Send + 'static, { - local(py, qualname, run(runtime.clone(), fut)) + local(py, qualname, runtime.clone().run_task(fut)) } /// Poll `fut` on the event loop thread, attached to Python, when awaited. @@ -142,20 +139,3 @@ where }) }) } - -/// Spawn `fut` on the runtime and wait for it; dropping the wait aborts the task. -/// The task keeps the runtime alive until it ends. -#[inline] -pub async fn run(runtime: Runtime, fut: F) -> PyResult -where - F: Future> + Send + 'static, - T: Send + 'static, -{ - let task = runtime.handle().clone().spawn(async move { - let _owner = runtime; - fut.await - }); - AbortOnDropHandle::new(task) - .await - .map_err(|err| PyRuntimeError::new_err(err.to_string()))? -} diff --git a/src/runtime.rs b/src/runtime.rs index 81f9be14..8724a410 100644 --- a/src/runtime.rs +++ b/src/runtime.rs @@ -1,17 +1,30 @@ +//! The Tokio runtime behind clients, and how calls wait on it. +//! +//! Requests, and async reads that must wait, run as tasks on the client's selected worker +//! ([`Runtime::block_on_task`], [`Runtime::run_task`]); blocking reads are polled by the +//! caller ([`Runtime::block_on`], [`Runtime::block_on_eager`]). The `block_on*` calls release +//! the GIL while they wait, and drive a current-thread runtime meanwhile. + use std::{ cell::{Cell, RefCell}, future::Future, + pin::pin, sync::{Arc, OnceLock}, thread, time::Duration, }; +use futures_util::FutureExt; use pingora_runtime::{BlockingPoolOpts, Runtime as PingoraRuntime, RuntimeBuilder}; use pyo3::{ exceptions::{PyRuntimeError, PyValueError}, prelude::*, }; -use tokio::runtime::{Builder, Handle}; +use tokio::{ + runtime::{Builder, Handle}, + task::JoinError, +}; +use tokio_util::task::AbortOnDropHandle; /// How a [`Runtime`] schedules client work. #[pyclass(eq, eq_int, frozen, from_py_object)] @@ -41,7 +54,7 @@ pub struct Runtime { /// The runtime behind [`Runtime`]. enum Shared { Workers(PingoraRuntime), - /// Driven only by the threads blocked in [`Runtime::block_on`]. + /// Driven only by threads in a blocking call, through [`Driving::drive`]. CurrentThread(tokio::runtime::Runtime), } @@ -80,20 +93,6 @@ impl Runtime { matches!(self.inner.as_deref(), Some(Shared::CurrentThread(_))) } - /// Poll `future` to completion on the calling thread, which must be detached from - /// Python. A current-thread runtime runs its tasks and IO here meanwhile. - pub fn block_on(&self, future: impl Future>) -> PyResult { - refuse_nested()?; - match self.inner.as_deref() { - Some(Shared::CurrentThread(runtime)) => { - let _driving = Driving::enter(); - let output = runtime.block_on(future); - INTERRUPT.take().map_or(output, Err) - } - _ => self.handle.block_on(future), - } - } - /// Share the runtime and select a worker for a new client. pub fn select(&self) -> PyResult { let inner = self @@ -105,6 +104,91 @@ impl Runtime { inner: Some(inner.clone()), }) } + + /// Poll `future` once inside the runtime, where a read timeout can start its timer; + /// `None` if it must wait. + pub fn poll_now(&self, future: &mut F) -> Option { + let _runtime = self.handle.enter(); + future.now_or_never() + } + + /// Poll `future` to completion on the caller, detached from Python. A current-thread + /// runtime runs its tasks and IO here meanwhile. + pub fn block_on(&self, py: Python<'_>, future: F) -> PyResult + where + F: Future> + Send, + T: Send, + { + py.detach(|| self.wait(future)) + } + + /// Like [`block_on`](Self::block_on), but poll once still attached first and return + /// without releasing the GIL when ready, as a buffered body or frame usually is. + pub fn block_on_eager(&self, py: Python<'_>, future: F) -> PyResult + where + F: Future> + Send, + T: Send, + { + // Refused before the first poll, which a later refusal would leave half done. + refuse_nested()?; + let mut future = pin!(future); + match self.poll_now(&mut future) { + Some(output) => output, + None => self.block_on(py, future), + } + } + + /// Run `future` as a task on the selected worker while the caller waits detached. A + /// current-thread runtime's only worker is its caller, which drives the future itself. + pub fn block_on_task(&self, py: Python<'_>, future: F) -> PyResult + where + F: Future> + Send + 'static, + T: Send + 'static, + { + py.detach(|| self.wait_task(future)) + } + + /// Run `future` as a task on the selected worker and wait for it from any executor. + /// Dropping the wait aborts the task, which keeps the runtime alive until it ends. + #[inline] + pub async fn run_task(self, future: F) -> PyResult + where + F: Future> + Send + 'static, + T: Send + 'static, + { + let task = self.handle.clone().spawn(async move { + let _owner = self; + future.await + }); + AbortOnDropHandle::new(task).await.map_err(join_error)? + } + + /// [`block_on`](Self::block_on) for a caller already detached. It refuses a nested call + /// itself, where Tokio would panic. + fn wait(&self, future: impl Future>) -> PyResult { + refuse_nested()?; + match self.inner.as_deref() { + Some(Shared::CurrentThread(runtime)) => Driving::drive(runtime, future), + _ => self.handle.block_on(future), + } + } + + /// [`block_on_task`](Self::block_on_task) for a caller already detached; refused before + /// the spawn, so a nested request starts nothing. + fn wait_task(&self, future: F) -> PyResult + where + F: Future> + Send + 'static, + T: Send + 'static, + { + refuse_nested()?; + match self.inner.as_deref() { + Some(Shared::CurrentThread(runtime)) => Driving::drive(runtime, future), + _ => self + .handle + .block_on(self.handle.spawn(future)) + .map_err(join_error)?, + } + } } impl From for Runtime { @@ -234,9 +318,16 @@ impl Shared { // ===== impl Driving ===== impl Driving { - fn enter() -> Self { + /// Drive `runtime` on this thread until `future` completes, marked as driving it; an + /// interrupt its upload iterator stashed meanwhile replaces the output. + fn drive( + runtime: &tokio::runtime::Runtime, + future: impl Future>, + ) -> PyResult { DRIVING.set(true); - Driving + let _driving = Driving; + let output = runtime.block_on(future); + INTERRUPT.take().map_or(output, Err) } } @@ -259,6 +350,11 @@ pub fn refuse_nested() -> PyResult<()> { Ok(()) } +/// A task that panicked or was aborted, as a Python error. +fn join_error(err: JoinError) -> PyErr { + PyRuntimeError::new_err(err.to_string()) +} + /// Workers used when none are given: one per available CPU. fn parallelism() -> usize { thread::available_parallelism().map_or(1, usize::from) @@ -298,19 +394,20 @@ mod tests { let caller = std::thread::current().id(); let mut ids = Vec::new(); for worker in [&first, &second, &first.clone(), &second.clone()] { - let id = crate::client::nogil::block_on(worker, async { - let thread = std::thread::current().id(); - for _ in 0..8 { - tokio::task::yield_now().await; - assert_eq!(thread, std::thread::current().id()); - } - let child = tokio::spawn(async { std::thread::current().id() }) - .await - .unwrap(); - assert_eq!(thread, child); - Ok(thread) - }) - .unwrap(); + let id = worker + .wait_task(async { + let thread = std::thread::current().id(); + for _ in 0..8 { + tokio::task::yield_now().await; + assert_eq!(thread, std::thread::current().id()); + } + let child = tokio::spawn(async { std::thread::current().id() }) + .await + .unwrap(); + assert_eq!(thread, child); + Ok(thread) + }) + .unwrap(); assert_ne!(id, caller); ids.push(id); } @@ -328,22 +425,23 @@ mod tests { }; let runtime = current_thread(); // Spawned tasks and timers run on the blocked caller. - let child = crate::client::nogil::block_on(&runtime, async { - Ok(tokio::spawn(async { - tokio::time::sleep(Duration::from_millis(1)).await; - thread::current().id() + let child = runtime + .wait_task(async { + Ok(tokio::spawn(async { + tokio::time::sleep(Duration::from_millis(1)).await; + thread::current().id() + }) + .await + .unwrap()) }) - .await - .unwrap()) - }) - .unwrap(); + .unwrap(); assert_eq!(child, thread::current().id()); // A blocking call from code the runtime runs fails instead of panicking, and the // last owner of another runtime can be released there. let other = current_thread(); let nested = runtime - .block_on(async { - let nested = other.block_on(async { Ok(()) }).is_err(); + .wait(async { + let nested = other.wait(async { Ok(()) }).is_err(); drop(other); Ok(nested) }) From 1fb5fa64b8087fccbdf1866bf9f60cdfb3581592 Mon Sep 17 00:00:00 2001 From: 0x676e67 Date: Wed, 7 Oct 2026 07:22:57 +0000 Subject: [PATCH 08/25] refactor(upload): split upload streams around a shared budget queue --- src/client/body/stream.rs | 538 ++-------------------------- src/client/body/stream/async_gen.rs | 202 +++++++++++ src/client/body/stream/iter.rs | 132 +++++++ src/client/body/stream/pump.rs | 137 +++++++ src/client/body/stream/queue.rs | 149 ++++++++ 5 files changed, 647 insertions(+), 511 deletions(-) create mode 100644 src/client/body/stream/async_gen.rs create mode 100644 src/client/body/stream/iter.rs create mode 100644 src/client/body/stream/pump.rs create mode 100644 src/client/body/stream/queue.rs diff --git a/src/client/body/stream.rs b/src/client/body/stream.rs index fa766a02..5b93b9f6 100644 --- a/src/client/body/stream.rs +++ b/src/client/body/stream.rs @@ -1,39 +1,31 @@ //! Request bodies streamed from Python iterators and async generators. +//! +//! `iter` reads a Python iterator, inline on a current-thread runtime or ahead of the upload +//! through `pump`, and `async_gen` forwards an async generator from its event loop. Both +//! queue chunks through `queue`, which bounds the bytes read ahead of the connection. + +mod async_gen; +mod iter; +mod pump; +mod queue; use std::{ - mem, - pin::{Pin, pin}, - sync::{ - Arc, Mutex, MutexGuard, PoisonError, - atomic::{AtomicBool, Ordering}, - }, + pin::Pin, + sync::{Mutex, MutexGuard, PoisonError}, task::{Context, Poll}, - thread::{self, ThreadId}, - time::Duration, }; use bytes::Bytes; -use futures_util::Stream; +use futures_util::{Stream, StreamExt, future::Either}; use pyo3::{ - IntoPyObjectExt, - exceptions::{PyException, PyKeyboardInterrupt, PyRuntimeError, PyStopIteration}, + exceptions::PyStopIteration, intern, prelude::*, - sync::PyOnceLock, types::{PyIterator, PyString}, }; -use tokio::{ - runtime::Handle, - sync::{OwnedSemaphorePermit, Semaphore, TryAcquireError, mpsc}, - task::{spawn_blocking, yield_now}, - time, -}; -use crate::{ - coroutine::{self, Coroutine}, - extractor::{Binary, Text}, - runtime, -}; +use self::{async_gen::AsyncStream, iter::SyncStream}; +use crate::extractor::{Binary, Text}; type Item = PyResult; @@ -45,89 +37,7 @@ pub enum PyBytesLike { /// A request body read from a Python iterator or async generator, for `body=` and /// multipart parts. -pub struct PyStream(Source); - -/// The source behind a [`PyStream`]. -enum Source { - Sync(SyncStream), - Async(PyAsyncStream), -} - -/// A request body from a Python iterator. -/// -/// Worker runtimes read the iterator on Tokio's blocking pool: neither a worker nor a -/// blocked caller runs it, so a slow `__next__` cannot delay the response, a timeout or -/// cancellation. A current-thread runtime reads it on its driving thread instead, which -/// waits for the request anyway. -enum SyncStream { - /// Not polled yet; the first poll picks where the iterator is read. `owner` passed the - /// iterator to the request. - Idle { - iter: Py, - owner: ThreadId, - }, - /// Read on the thread driving a current-thread runtime; `flushed` once a poll has - /// returned to the connection since the last item. - Inline { - iter: Py, - owner: ThreadId, - flushed: bool, - }, - Pumped(Pumped), - /// The inline iterator ended or raised. - Done, -} - -/// An iterator read by pump tasks on the blocking pool, up to [`UPLOAD_BUDGET`] ahead of -/// the upload. When the budget stays full, a pump parks the iterator with the chunk it read -/// and frees its thread; the next item taken restarts it. -struct Pumped { - rx: mpsc::UnboundedReceiver, - budget: Arc, - /// The pump while none runs: before the first poll, or once parked at a full budget. - parked: Arc>>, -} - -/// What a pump task needs to read the iterator into the body. -struct Pump { - iter: Py, - tx: mpsc::UnboundedSender, - budget: Arc, - /// A chunk read before the pump parked, queued first when it restarts. - pending: Option, -} - -/// Bytes an upload may queue ahead of the connection. -const UPLOAD_BUDGET: usize = 256 * 1024; - -/// The least budget a queued chunk holds, so small chunks queue at most 64 items. -const UPLOAD_CHARGE: usize = 4 * 1024; - -/// A request body from a Python async generator, forwarded by a task on the loop that was -/// running at extraction, up to [`UPLOAD_BUDGET`] ahead. Dropping it cancels that task. -struct PyAsyncStream { - rx: mpsc::UnboundedReceiver, - budget: Arc, - /// Set by [`Sender::finish`], so the channel closing reads as the end of the body. - finished: Arc, - task: Option<(Py, Py)>, -} - -/// A queued item and the share of the upload budget it holds until the body takes it. -struct Queued { - item: Item, - _share: Option, -} - -/// The channel end given to the forwarding coroutine: `try_send` queues a chunk within the -/// budget or waits for room, and `send` also queues the error that ends the body. Closing -/// it without `finish` fails the body. -#[pyclass(frozen)] -struct Sender { - tx: Mutex>>, - budget: Arc, - finished: Arc, -} +pub struct PyStream(Either); // ===== impl PyBytesLike ===== @@ -145,14 +55,12 @@ impl FromPyObject<'_, '_> for PyBytesLike { } impl PyBytesLike { - /// The share of the upload budget this chunk holds while queued; one larger than the - /// budget takes all of it. - fn share(&self) -> u32 { - let len = match self { + /// The chunk's length in bytes. + fn len(&self) -> usize { + match self { PyBytesLike::Bytes(b) => b.0.len(), PyBytesLike::String(s) => s.0.len(), - }; - u32::try_from(len.clamp(UPLOAD_CHARGE, UPLOAD_BUDGET)).unwrap_or(u32::MAX) + } } } @@ -174,12 +82,9 @@ impl FromPyObject<'_, '_> for PyStream { fn extract(ob: Borrowed) -> PyResult { // Iterators are checked first; probing `asend` would raise for each of them. let source = if ob.cast::().is_err() && ob.hasattr(intern!(ob.py(), "asend"))? { - Source::Async(PyAsyncStream::new(ob.to_owned())?) + Either::Right(AsyncStream::new(ob.to_owned())?) } else { - Source::Sync(SyncStream::Idle { - iter: ob.to_owned().unbind(), - owner: thread::current().id(), - }) + Either::Left(SyncStream::new(ob.to_owned().unbind())) }; Ok(PyStream(source)) } @@ -189,190 +94,10 @@ impl Stream for PyStream { type Item = Item; fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - match &mut self.get_mut().0 { - Source::Sync(stream) => stream.poll_next(cx), - Source::Async(stream) => Pin::new(stream).poll_next(cx), - } + self.get_mut().0.poll_next_unpin(cx) } } -// ===== impl SyncStream ===== - -impl SyncStream { - fn poll_next(&mut self, cx: &mut Context<'_>) -> Poll> { - if let SyncStream::Idle { .. } = self - && let SyncStream::Idle { iter, owner } = mem::replace(self, SyncStream::Done) - { - *self = if runtime::driving() { - SyncStream::Inline { - iter, - owner, - flushed: false, - } - } else { - SyncStream::Pumped(Pumped::new(iter)) - }; - } - match self { - SyncStream::Inline { - iter, - owner, - flushed, - } => { - // Return to the connection before each read, so the request head and the - // chunks already yielded go out before `__next__` can block. The deferred - // wake lets the runtime poll its IO and the caller first. - if !mem::replace(flushed, true) { - let _ = pin!(yield_now()).poll(cx); - return Poll::Pending; - } - *flushed = false; - // Like a pump, stop reading once Python is unavailable. - let item = Python::try_attach(|py| { - let item = next_item(py, iter); - // A KeyboardInterrupt or SystemExit reaches the caller as itself: the - // owner's call, or for Ctrl+C the main thread signals reach. - if let Some(Err(err)) = &item - && !err.is_instance_of::(py) - && (*owner == thread::current().id() - || err.is_instance_of::(py)) - { - runtime::interrupt(err.clone_ref(py)); - } - item - }) - .flatten(); - if !matches!(item, Some(Ok(_))) { - *self = SyncStream::Done; - } - Poll::Ready(item) - } - SyncStream::Pumped(pumped) => pumped.poll_next(cx), - SyncStream::Idle { .. } | SyncStream::Done => Poll::Ready(None), - } - } -} - -// ===== impl Pumped ===== - -impl Pumped { - fn new(iter: Py) -> Self { - let (tx, rx) = mpsc::unbounded_channel(); - let budget = Arc::new(Semaphore::new(UPLOAD_BUDGET)); - let pump = Pump { - iter, - tx, - budget: budget.clone(), - pending: None, - }; - Pumped { - rx, - budget, - parked: Arc::new(Mutex::new(Some(pump))), - } - } - - fn poll_next(&mut self, cx: &mut Context<'_>) -> Poll> { - // Taking an item returns its share of the budget before the parked check. - let poll = self - .rx - .poll_recv(cx) - .map(|queued| queued.map(|queued| queued.item)); - // A parked pump waits for budget, which the first poll or a taken item frees. The - // check follows the take, so a pump parking concurrently is seen here. - let parked = lock(&self.parked).take(); - if let Some(pump) = parked { - let parked = self.parked.clone(); - spawn_blocking(move || pump.run(&parked)); - } - poll - } -} - -impl Drop for Pumped { - fn drop(&mut self) { - // A pump waiting for budget stops at once, without reading another item. - self.budget.close(); - } -} - -// ===== impl Pump ===== - -impl Pump { - /// How long a pump waits for budget before parking to free its thread. - const PARK_AFTER: Duration = Duration::from_millis(10); - - /// Queue chunks until the iterator ends or raises, the body is dropped, or the budget - /// stays full past [`PARK_AFTER`](Self::PARK_AFTER). - fn run(mut self, parked: &Mutex>) { - let handle = Handle::current(); - // Once Python is unavailable, stop reading without creating a PyErr that could - // require another attachment to format. - Python::try_attach(|py| { - loop { - if self.tx.is_closed() { - return; - } - let chunk = match self.pending.take() { - Some(chunk) => chunk, - None => match next_item(py, &self.iter) { - Some(Ok(chunk)) => chunk, - Some(Err(err)) => { - let _ = self.tx.send(Queued { - item: Err(err), - _share: None, - }); - return; - } - None => return, - }, - }; - let share = match self.budget.clone().try_acquire_many_owned(chunk.share()) { - Ok(share) => share, - Err(TryAcquireError::Closed) => return, - Err(TryAcquireError::NoPermits) => { - let budget = self.budget.clone(); - let wait = time::timeout( - Self::PARK_AFTER, - budget.acquire_many_owned(chunk.share()), - ); - match py.detach(|| handle.block_on(wait)) { - Ok(Ok(share)) => share, - Ok(Err(_)) => return, - Err(_) => { - // Park under the lock the body takes after receiving, so - // either it sees the pump parked or the pump sees budget. - let mut slot = lock(parked); - match self.budget.clone().try_acquire_many_owned(chunk.share()) { - Ok(share) => share, - Err(TryAcquireError::Closed) => return, - Err(TryAcquireError::NoPermits) => { - self.pending = Some(chunk); - *slot = Some(self); - return; - } - } - } - } - } - }; - let queued = Queued { - item: Ok(chunk), - _share: Some(share), - }; - if self.tx.send(queued).is_err() { - return; - } - } - }); - } -} - -#[inline] -fn lock(parked: &Mutex>) -> MutexGuard<'_, Option> { - parked.lock().unwrap_or_else(PoisonError::into_inner) -} - /// Call `__next__`, ending the stream at StopIteration. fn next_item(py: Python<'_>, iter: &Py) -> Option { match iter.call_method0(py, intern!(py, "__next__")) { @@ -382,217 +107,8 @@ fn next_item(py: Python<'_>, iter: &Py) -> Option { } } -// ===== impl PyAsyncStream ===== - -impl PyAsyncStream { - /// Start forwarding on the running loop, so the body must be extracted in a coroutine. - fn new(generator: Bound<'_, PyAny>) -> PyResult { - static FORWARD: PyOnceLock> = PyOnceLock::new(); - let py = generator.py(); - let event_loop = coroutine::running_loop(py)?; - let forward = FORWARD.get_or_try_init(py, || { - PyModule::from_code( - py, - c"import asyncio - -async def forward(gen, sender): - try: - try: - async for item in gen: - # Queued at once while the upload budget has room, otherwise awaited. - sent = sender.try_send(item) - if not isinstance(sent, bool): - sent = await sent - if not sent: - return - finally: - close = getattr(gen, 'aclose', None) - if close is not None: - await close() - except asyncio.CancelledError as error: - # Task cancellation must not wait for space in a retained body, and must not leave - # the body waiting while a traceback keeps this frame and its sender alive. - if asyncio.current_task().cancelling(): - sender.close() - else: - await sender.send(error, True) - raise - except BaseException as error: - await sender.send(error, True) - else: - sender.finish() -", - c"wreq/_async_stream.py", - c"wreq._async_stream", - )? - .getattr("forward") - .map(Bound::unbind) - })?; - let (tx, rx) = mpsc::unbounded_channel(); - let budget = Arc::new(Semaphore::new(UPLOAD_BUDGET)); - let finished = Arc::new(AtomicBool::new(false)); - let sender = Sender { - tx: Mutex::new(Some(tx)), - budget: budget.clone(), - finished: finished.clone(), - }; - let coroutine = forward.bind(py).call1((generator, sender))?; - // create_task captures the caller's contextvars on the running loop. - let task = event_loop - .call_method1(intern!(py, "create_task"), (&coroutine,)) - .inspect_err(|_| { - let _ = coroutine.call_method0(intern!(py, "close")); - })?; - Ok(Self { - rx, - budget, - finished, - task: Some((task.unbind(), event_loop.unbind())), - }) - } -} - -impl Stream for PyAsyncStream { - type Item = Item; - - fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - let this = self.get_mut(); - match this.rx.poll_recv(cx) { - // Taking the item returns its share of the budget. - Poll::Ready(Some(queued)) => Poll::Ready(Some(queued.item)), - // Every sender is gone: the end after `finish`. Otherwise the forwarding task was - // cancelled or destroyed, so the body is incomplete. - Poll::Ready(None) - if this.task.take().is_some() && !this.finished.load(Ordering::Acquire) => - { - Poll::Ready(Some(Err(PyRuntimeError::new_err( - "async body generator stopped before it finished", - )))) - } - Poll::Ready(None) => Poll::Ready(None), - Poll::Pending => Poll::Pending, - } - } -} - -impl Drop for PyAsyncStream { - fn drop(&mut self) { - self.rx.close(); - // A send waiting for room resolves to `False` at once. - self.budget.close(); - if let Some((task, event_loop)) = self.task.take() { - // Body drop can run on Tokio: cancel from a blocking thread, preferring the current - // runtime so a drop on a client's own runtime does not start the shared one. - let handle = Handle::try_current().unwrap_or_else(|_| runtime::get().handle().clone()); - handle.spawn_blocking(move || { - Python::try_attach(|py| { - if let Ok(cancel) = task.bind(py).getattr(intern!(py, "cancel")) { - let _ = event_loop.call_method1( - py, - intern!(py, "call_soon_threadsafe"), - (cancel,), - ); - } - }); - }); - } - } -} - -// ===== impl Sender ===== - -#[pymethods] -impl Sender { - /// Queue a chunk: `True` once queued within the budget, `False` once the body is dropped - /// or the sender closed, and otherwise a coroutine that queues it once there is room. - fn try_send<'py>(&self, item: Bound<'py, PyAny>) -> PyResult> { - let py = item.py(); - let chunk = item.extract::()?; - let tx = self.lock().clone(); - let sent = match &tx { - None => false, - Some(queue) => match self.budget.clone().try_acquire_many_owned(chunk.share()) { - Ok(share) => queue - .send(Queued { - item: Ok(chunk), - _share: Some(share), - }) - .is_ok(), - Err(TryAcquireError::Closed) => false, - // The coroutine keeps the extracted chunk, so it is never copied twice. - Err(TryAcquireError::NoPermits) => { - return self.send_later(py, tx, Ok(chunk)).map(Bound::into_any); - } - }, - }; - sent.into_bound_py_any(py) - } - - /// Queue `item` as the error that ends the body, or a chunk once the budget has room. - /// Resolves to `False` once the body is dropped or the sender closed, which stops - /// forwarding. - fn send<'py>( - &self, - py: Python<'py>, - item: Bound<'py, PyAny>, - error: bool, - ) -> PyResult> { - let item = if error { - Err(PyErr::from_value(item)) - } else { - Ok(item.extract()?) - }; - self.send_later(py, self.lock().clone(), item) - } - - /// Mark the normal end of the body. The last chunk may still be queued: it is read - /// before the closed channel, so finishing never waits for room. - fn finish(&self) { - self.finished.store(true, Ordering::Release); - // Python may retain the sender after completion, especially on PyPy. - self.close(); - } - - /// Drop the channel end at once, so the body fails instead of waiting for more. - fn close(&self) { - let tx = self.lock().take(); - drop(tx); - } -} - -impl Sender { - #[inline] - fn lock(&self) -> MutexGuard<'_, Option>> { - self.tx.lock().unwrap_or_else(PoisonError::into_inner) - } - - /// A coroutine that queues `item` on `tx` once the budget has room; an error, which ends - /// the body, never waits. - fn send_later<'py>( - &self, - py: Python<'py>, - tx: Option>, - item: Item, - ) -> PyResult> { - let budget = self.budget.clone(); - // Budget readiness is runtime-independent, so this waits on the Python loop. - coroutine::local(py, "Sender.send", async move { - let Some(tx) = tx else { - return Ok(false); - }; - let share = match &item { - Ok(chunk) => match budget.acquire_many_owned(chunk.share()).await { - Ok(share) => Some(share), - Err(_) => return Ok(false), - }, - Err(_) => None, - }; - Ok(tx - .send(Queued { - item, - _share: share, - }) - .is_ok()) - }) - } +/// Lock `mutex`, recovering it from a panicked holder. +#[inline] +fn lock(mutex: &Mutex) -> MutexGuard<'_, T> { + mutex.lock().unwrap_or_else(PoisonError::into_inner) } diff --git a/src/client/body/stream/async_gen.rs b/src/client/body/stream/async_gen.rs new file mode 100644 index 00000000..ce450676 --- /dev/null +++ b/src/client/body/stream/async_gen.rs @@ -0,0 +1,202 @@ +//! Request bodies from Python async generators. + +use std::{ + ffi::CStr, + pin::Pin, + sync::{ + Arc, Mutex, + atomic::{AtomicBool, Ordering}, + }, + task::{Context, Poll, ready}, +}; + +use futures_util::{Stream, StreamExt}; +use pyo3::{IntoPyObjectExt, exceptions::PyRuntimeError, intern, prelude::*, sync::PyOnceLock}; +use tokio::{runtime::Handle, sync::mpsc::error::TrySendError}; + +use super::{Item, PyBytesLike, lock, queue}; +use crate::{coroutine, runtime}; + +/// A request body from a Python async generator, forwarded within the upload budget by a +/// task on the loop that was running at extraction. Dropping it cancels that task. +pub(super) struct AsyncStream { + rx: queue::Rx, + /// Set by [`Sender::finish`], so the queue closing reads as the end of the body. + finished: Arc, + /// The forwarding task and its loop, until forwarding ends. + task: Option<(Py, Py)>, +} + +/// The queue end given to the forwarding coroutine: `try_send` queues a chunk within the +/// budget or returns a coroutine that waits for room, and `fail` queues the error that ends +/// the body. Closing it without `finish` fails the body. +#[pyclass(frozen)] +struct Sender { + tx: Mutex>, + finished: Arc, +} + +// ===== impl AsyncStream ===== + +impl AsyncStream { + /// Start forwarding on the running loop, so the body must be extracted in a coroutine. + pub(super) fn new(generator: Bound<'_, PyAny>) -> PyResult { + static FORWARD: PyOnceLock> = PyOnceLock::new(); + let py = generator.py(); + let event_loop = coroutine::running_loop(py)?; + let forward = FORWARD.get_or_try_init(py, || { + PyModule::from_code( + py, + FORWARD_SOURCE, + c"wreq/_async_stream.py", + c"wreq._async_stream", + )? + .getattr("forward") + .map(Bound::unbind) + })?; + let (tx, rx) = queue::channel(); + let finished = Arc::new(AtomicBool::new(false)); + let sender = Sender { + tx: Mutex::new(Some(tx)), + finished: finished.clone(), + }; + let coroutine = forward.bind(py).call1((generator, sender))?; + // create_task captures the caller's contextvars on the running loop. + let task = event_loop + .call_method1(intern!(py, "create_task"), (&coroutine,)) + .inspect_err(|_| { + let _ = coroutine.call_method0(intern!(py, "close")); + })?; + Ok(Self { + rx, + finished, + task: Some((task.unbind(), event_loop.unbind())), + }) + } +} + +impl Stream for AsyncStream { + type Item = Item; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + Poll::Ready(match ready!(this.rx.poll_next_unpin(cx)) { + Some(item) => Some(item), + // Every sender is gone: the end after `finish`. Otherwise the forwarding task was + // cancelled or destroyed, so the body is incomplete. + None if this.task.take().is_some() && !this.finished.load(Ordering::Acquire) => { + Some(Err(PyRuntimeError::new_err( + "async body generator stopped before it finished", + ))) + } + None => None, + }) + } +} + +impl Drop for AsyncStream { + fn drop(&mut self) { + // Senders see the body gone at once, and a send waiting for room resolves to `False`. + self.rx.close(); + if let Some((task, event_loop)) = self.task.take() { + // Body drop can run on Tokio: cancel from a blocking thread, preferring the current + // runtime so a drop on a client's own runtime does not start the shared one. + let handle = Handle::try_current().unwrap_or_else(|_| runtime::get().handle().clone()); + handle.spawn_blocking(move || { + Python::try_attach(|py| { + if let Ok(cancel) = task.bind(py).getattr(intern!(py, "cancel")) { + let _ = event_loop.call_method1( + py, + intern!(py, "call_soon_threadsafe"), + (cancel,), + ); + } + }); + }); + } + } +} + +// ===== impl Sender ===== + +#[pymethods] +impl Sender { + /// Queue a chunk: `True` once queued within the budget, `False` once the body is dropped + /// or the sender closed, and otherwise a coroutine that queues it once there is room. + fn try_send<'py>(&self, item: Bound<'py, PyAny>) -> PyResult> { + let py = item.py(); + let chunk = item.extract::()?; + // Dropping a chunk runs no Python code, so a failed send may drop it under the lock. + let tx = match &*lock(&self.tx) { + Some(tx) => match tx.try_reserve(&chunk) { + Ok(permit) => return permit.send(chunk).into_bound_py_any(py), + Err(TrySendError::Full(())) => tx.clone(), + Err(TrySendError::Closed(())) => return false.into_bound_py_any(py), + }, + None => return false.into_bound_py_any(py), + }; + // The coroutine keeps the extracted chunk, so it is never copied twice. Budget + // readiness is runtime-independent, so it waits on the Python loop. Its name is the + // pending send that shutdown inspects. + coroutine::local(py, "Sender.send", async move { + let permit = tx.reserve(&chunk).await; + Ok(permit.is_some_and(|permit| permit.send(chunk))) + }) + .map(Bound::into_any) + } + + /// Queue `error` as the end of the body, without waiting for room. + fn fail(&self, error: Bound<'_, PyAny>) { + // Released first: dropping an unsent error can run Python code. + let tx = lock(&self.tx).clone(); + if let Some(tx) = tx { + tx.fail(PyErr::from_value(error)); + } + } + + /// Mark the normal end of the body. The last chunk may still be queued: it is read + /// before the closed channel, so finishing never waits for room. + fn finish(&self) { + self.finished.store(true, Ordering::Release); + // Python may retain the sender after completion, especially on PyPy. + self.close(); + } + + /// Drop the channel end at once, so the body fails instead of waiting for more. + fn close(&self) { + let tx = lock(&self.tx).take(); + drop(tx); + } +} + +/// The coroutine that forwards an async generator into its [`Sender`]. +const FORWARD_SOURCE: &CStr = c"import asyncio + +async def forward(gen, sender): + try: + try: + async for item in gen: + # Queued at once while the upload budget has room, otherwise awaited. + sent = sender.try_send(item) + if not isinstance(sent, bool): + sent = await sent + if not sent: + return + finally: + close = getattr(gen, 'aclose', None) + if close is not None: + await close() + except asyncio.CancelledError as error: + # On task cancellation, close the sender so the body fails as incomplete at once, + # even while a traceback keeps this frame and its sender alive. Any other + # CancelledError came from the generator and ends the body as its error. + if asyncio.current_task().cancelling(): + sender.close() + else: + sender.fail(error) + raise + except BaseException as error: + sender.fail(error) + else: + sender.finish() +"; diff --git a/src/client/body/stream/iter.rs b/src/client/body/stream/iter.rs new file mode 100644 index 00000000..c5ca948e --- /dev/null +++ b/src/client/body/stream/iter.rs @@ -0,0 +1,132 @@ +//! Request bodies from Python iterators. + +use std::{ + mem, + pin::{Pin, pin}, + task::{Context, Poll, ready}, + thread::{self, ThreadId}, +}; + +use futures_util::{Stream, StreamExt}; +use pyo3::{ + exceptions::{PyException, PyKeyboardInterrupt}, + prelude::*, +}; +use tokio::task::yield_now; + +use super::{Item, next_item, pump::Pumped}; +use crate::runtime; + +/// A request body from a Python iterator. +/// +/// Worker runtimes read the iterator on Tokio's blocking pool ([`Pumped`]): neither a worker +/// nor a blocked caller runs it, so a slow `__next__` cannot delay the response, a timeout or +/// cancellation. A current-thread runtime reads it on its driving thread instead +/// ([`Inline`]), which waits for the request anyway. +pub(super) enum SyncStream { + /// Not polled yet; the first poll picks where the iterator is read. `owner` passed the + /// iterator to the request. + Idle { + iter: Py, + owner: ThreadId, + }, + Inline(Inline), + Pumped(Pumped), + /// The inline iterator ended or raised. + Done, +} + +/// An iterator read on the thread driving a current-thread runtime, between its IO polls. +pub(super) struct Inline { + iter: Py, + /// The thread that passed the iterator to the request. + owner: ThreadId, + /// Whether a poll has returned to the connection since the last item. + flushed: bool, +} + +// ===== impl SyncStream ===== + +impl SyncStream { + /// An iterator the calling thread passes to a request. + pub(super) fn new(iter: Py) -> Self { + SyncStream::Idle { + iter, + owner: thread::current().id(), + } + } +} + +impl Stream for SyncStream { + type Item = Item; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + if let SyncStream::Idle { .. } = this + && let SyncStream::Idle { iter, owner } = mem::replace(this, SyncStream::Done) + { + *this = if runtime::driving() { + SyncStream::Inline(Inline { + iter, + owner, + flushed: false, + }) + } else { + SyncStream::Pumped(Pumped::new(iter)) + }; + } + match this { + SyncStream::Inline(inline) => { + let poll = inline.poll_next_unpin(cx); + // Release the iterator once it ends or raises. + if let Poll::Ready(None | Some(Err(_))) = poll { + *this = SyncStream::Done; + } + poll + } + SyncStream::Pumped(pumped) => pumped.poll_next_unpin(cx), + SyncStream::Idle { .. } | SyncStream::Done => Poll::Ready(None), + } + } +} + +// ===== impl Inline ===== + +impl Inline { + /// Return to the connection once before each read, so the request head and the chunks + /// already yielded go out before `__next__` can block. The deferred wake lets the runtime + /// poll its IO and the caller first. + fn poll_flush(&mut self, cx: &mut Context<'_>) -> Poll<()> { + if mem::take(&mut self.flushed) { + return Poll::Ready(()); + } + self.flushed = true; + let _ = pin!(yield_now()).poll(cx); + Poll::Pending + } +} + +impl Stream for Inline { + type Item = Item; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + ready!(this.poll_flush(cx)); + // Like a pump, stop reading once Python is unavailable. + Python::try_attach(|py| { + let item = next_item(py, &this.iter); + // A KeyboardInterrupt or SystemExit reaches the caller as itself: the owner's call, + // or for Ctrl+C the main thread signals reach. + if let Some(Err(err)) = &item + && !err.is_instance_of::(py) + && (this.owner == thread::current().id() + || err.is_instance_of::(py)) + { + runtime::interrupt(err.clone_ref(py)); + } + item + }) + .flatten() + .into() + } +} diff --git a/src/client/body/stream/pump.rs b/src/client/body/stream/pump.rs new file mode 100644 index 00000000..6f790114 --- /dev/null +++ b/src/client/body/stream/pump.rs @@ -0,0 +1,137 @@ +//! Python iterators read ahead of the upload on Tokio's blocking pool. + +use std::{ + pin::Pin, + sync::{Arc, Mutex}, + task::{Context, Poll}, + time::Duration, +}; + +use futures_util::{Stream, StreamExt}; +use pyo3::prelude::*; +use tokio::{runtime::Handle, sync::mpsc::error::TrySendError, task::spawn_blocking, time}; + +use super::{ + Item, PyBytesLike, lock, next_item, + queue::{self, Permit}, +}; + +/// An iterator read ahead by pump tasks on the blocking pool, within the upload budget. When +/// the budget stays full, a pump parks the iterator with the chunk it read and frees its +/// thread; the next item taken restarts it. +pub(super) struct Pumped { + rx: queue::Rx, + /// The pump while none runs: before the first poll, or once parked at a full budget. + parked: Arc>>, +} + +/// What a pump task needs to read the iterator into the body. +struct Pump { + iter: Py, + tx: queue::Tx, + /// A chunk read before the pump parked, queued first when it restarts. + pending: Option, +} + +// ===== impl Pumped ===== + +impl Pumped { + pub(super) fn new(iter: Py) -> Self { + let (tx, rx) = queue::channel(); + let pump = Pump { + iter, + tx, + pending: None, + }; + Pumped { + rx, + parked: Arc::new(Mutex::new(Some(pump))), + } + } +} + +impl Stream for Pumped { + type Item = Item; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + // Taking an item returns its share of the budget before the parked check. + let poll = this.rx.poll_next_unpin(cx); + // A parked pump waits for budget, which the first poll or a taken item frees. The + // check follows the take, so a pump parking concurrently is seen here. + let parked = lock(&this.parked).take(); + if let Some(pump) = parked { + let parked = this.parked.clone(); + spawn_blocking(move || pump.run(&parked)); + } + poll + } +} + +// ===== impl Pump ===== + +impl Pump { + /// How long a pump waits for budget before parking to free its thread. + const PARK_AFTER: Duration = Duration::from_millis(10); + + /// Queue chunks until the iterator ends or raises, the body is dropped, or the budget + /// stays full past [`PARK_AFTER`](Self::PARK_AFTER). + fn run(mut self, parked: &Mutex>) { + let handle = Handle::current(); + // Once Python is unavailable, stop reading without creating a PyErr that could + // require another attachment to format. + Python::try_attach(|py| { + while !self.tx.is_closed() { + let chunk = match self.pending.take() { + Some(chunk) => chunk, + None => match next_item(py, &self.iter) { + Some(Ok(chunk)) => chunk, + Some(Err(err)) => { + self.tx.fail(err); + return; + } + None => return, + }, + }; + let permit = match self.wait_for_room(py, &handle, &chunk) { + Some(permit) => permit, + None => { + // Park under the lock the body takes after receiving, so either it + // sees the pump parked or the pump sees budget. + let mut slot = lock(parked); + match self.tx.try_reserve(&chunk) { + Ok(permit) => permit, + Err(TrySendError::Closed(())) => return, + Err(TrySendError::Full(())) => { + self.pending = Some(chunk); + *slot = Some(self); + return; + } + } + } + }; + if !permit.send(chunk) { + return; + } + } + }); + } + + /// Reserve room for `chunk`, waiting detached up to [`PARK_AFTER`](Self::PARK_AFTER) while + /// the budget is full; `None` if it stays full or the body is dropped. + fn wait_for_room( + &self, + py: Python<'_>, + handle: &Handle, + chunk: &PyBytesLike, + ) -> Option> { + match self.tx.try_reserve(chunk) { + Ok(permit) => Some(permit), + Err(TrySendError::Closed(())) => None, + Err(TrySendError::Full(())) => { + let wait = time::timeout(Self::PARK_AFTER, self.tx.reserve(chunk)); + py.detach(|| handle.block_on(wait)).ok().flatten() + } + } + } +} diff --git a/src/client/body/stream/queue.rs b/src/client/body/stream/queue.rs new file mode 100644 index 00000000..32cc0d51 --- /dev/null +++ b/src/client/body/stream/queue.rs @@ -0,0 +1,149 @@ +//! The queue between a chunk producer and its request body: each queued chunk holds a share +//! of the upload budget until the body takes it, and dropping the body closes the budget. + +use std::{ + pin::Pin, + sync::Arc, + task::{Context, Poll}, +}; + +use futures_util::Stream; +use pyo3::PyErr; +use tokio::sync::{ + OwnedSemaphorePermit, Semaphore, TryAcquireError, + mpsc::{self, error::TrySendError}, +}; + +use super::{Item, PyBytesLike}; + +/// Bytes an upload may queue ahead of the connection. +const UPLOAD_BUDGET: usize = 256 * 1024; + +/// The least budget a queued chunk holds, so small chunks queue at most 64 items. +const UPLOAD_CHARGE: usize = 4 * 1024; + +/// The producer end. Errors skip the budget, so the error that ends a body never waits. +#[derive(Clone)] +pub(super) struct Tx { + chan: mpsc::UnboundedSender, + budget: Arc, +} + +/// A chunk's reserved share of the budget, spent by [`Permit::send`]. +pub(super) struct Permit<'a> { + chan: &'a mpsc::UnboundedSender, + share: OwnedSemaphorePermit, +} + +/// The body end: taking an item returns its share, and dropping it stops waiting producers. +pub(super) struct Rx { + chan: mpsc::UnboundedReceiver, + budget: Arc, +} + +/// A queued item and the share it holds until the body takes it. +struct Queued { + item: Item, + _share: Option, +} + +/// A queue with the whole budget free. +pub(super) fn channel() -> (Tx, Rx) { + let (tx, rx) = mpsc::unbounded_channel(); + let budget = Arc::new(Semaphore::new(UPLOAD_BUDGET)); + let tx = Tx { + chan: tx, + budget: budget.clone(), + }; + (tx, Rx { chan: rx, budget }) +} + +/// The budget a queued `chunk` holds; one larger than the budget takes all of it. +fn charge(chunk: &PyBytesLike) -> u32 { + u32::try_from(chunk.len().clamp(UPLOAD_CHARGE, UPLOAD_BUDGET)).unwrap_or(u32::MAX) +} + +// ===== impl Tx ===== + +impl Tx { + /// Reserve `chunk`'s share if the budget has room; `Closed` once the body is dropped. + pub(super) fn try_reserve(&self, chunk: &PyBytesLike) -> Result, TrySendError<()>> { + match self.budget.clone().try_acquire_many_owned(charge(chunk)) { + Ok(share) => Ok(Permit { + chan: &self.chan, + share, + }), + Err(TryAcquireError::NoPermits) => Err(TrySendError::Full(())), + Err(TryAcquireError::Closed) => Err(TrySendError::Closed(())), + } + } + + /// Wait for room for `chunk`'s share; `None` once the body is dropped. + pub(super) async fn reserve(&self, chunk: &PyBytesLike) -> Option> { + let share = self + .budget + .clone() + .acquire_many_owned(charge(chunk)) + .await + .ok()?; + Some(Permit { + chan: &self.chan, + share, + }) + } + + /// Queue the error that ends the body. + pub(super) fn fail(&self, err: PyErr) { + let _ = self.chan.send(Queued { + item: Err(err), + _share: None, + }); + } + + /// Whether the body is dropped. + pub(super) fn is_closed(&self) -> bool { + self.chan.is_closed() + } +} + +// ===== impl Permit ===== + +impl Permit<'_> { + /// Queue `chunk`, the chunk this share was reserved for; `false` once the body is dropped. + pub(super) fn send(self, chunk: PyBytesLike) -> bool { + self.chan + .send(Queued { + item: Ok(chunk), + _share: Some(self.share), + }) + .is_ok() + } +} + +// ===== impl Rx ===== + +impl Rx { + /// Refuse further items and stop producers waiting for room: a pump stops without reading + /// another item, and a pending send resolves to `False`. + pub(super) fn close(&mut self) { + self.chan.close(); + self.budget.close(); + } +} + +impl Stream for Rx { + type Item = Item; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + // Taking an item returns its share of the budget. + self.chan + .poll_recv(cx) + .map(|queued| queued.map(|queued| queued.item)) + } +} + +impl Drop for Rx { + fn drop(&mut self) { + self.close(); + } +} From 13f636379a8b3f7244a108616578202dbc6c720b Mon Sep 17 00:00:00 2001 From: 0x676e67 Date: Wed, 7 Oct 2026 07:23:08 +0000 Subject: [PATCH 09/25] refactor(stream): name the iteration mode and share the frame read --- src/client/resp/stream.rs | 101 ++++++++++++++++++++++++-------------- 1 file changed, 64 insertions(+), 37 deletions(-) diff --git a/src/client/resp/stream.rs b/src/client/resp/stream.rs index 692bf35d..eae6de0c 100644 --- a/src/client/resp/stream.rs +++ b/src/client/resp/stream.rs @@ -73,6 +73,16 @@ enum State { /// Wakes readers however the read-ahead task ends, including when aborted. struct NotifyOnDrop(Arc); +/// How Python iterates a [`Streamer`], which sets the frames read on the caller and the error +/// that ends iteration. +#[derive(Clone, Copy)] +enum Iteration { + /// `__next__` on a blocking caller. + Sync, + /// `__anext__` on the event loop. + Async, +} + // ===== impl Streamer ===== impl Streamer { @@ -122,10 +132,8 @@ impl Streamer { /// Send frames until the body ends, a frame fails or the reader is gone. async fn pump(mut resp: wreq::Response, tx: &mpsc::Sender>, arrived: &Notify) { let mut unannounced = false; - while let Some(frame) = burst(resp.frame(), arrived, &mut unannounced).await { - let Some(frame) = Self::convert(frame, &mut resp) else { - continue; - }; + while let Some(frame) = burst(Self::next_frame(&mut resp), arrived, &mut unannounced).await + { let failed = frame.is_err(); if burst(tx.send(frame), arrived, &mut unannounced) .await @@ -138,6 +146,16 @@ impl Streamer { } } + /// Read the next frame, skipping unknown kinds; `None` at the end of the body. + async fn next_frame(resp: &mut wreq::Response) -> Option> { + loop { + let frame = resp.frame().await?; + if let Some(frame) = Self::convert(frame, resp) { + return Some(frame); + } + } + } + /// Convert a body frame, skipping unknown kinds. fn convert( frame: wreq::Result>, @@ -158,21 +176,20 @@ impl Streamer { } /// Read a frame the body already holds, before any read-ahead task starts, so a reader - /// of a short body never waits on another thread; `end` builds the error that ends - /// iteration. On the event loop only an HTTP/1 body of known length is read here, whatever - /// its size: taking a frame it holds is cheap, while a decompressing body is not. - fn ready_frame(&self, on_loop: bool, end: fn() -> Error) -> Option> { + /// of a short body never waits on another thread. + fn ready_frame(&self, iteration: Iteration) -> Option> { let mut state = self.reader.lock(); let State::Idle(resp) = &mut *state else { return None; }; - if on_loop && !(loop_polls(resp.version()) && resp.size_hint().exact().is_some()) { + if !iteration.takes_ready(resp) { return None; } + // A read timeout starts a timer on its first poll, which needs the runtime. let _runtime = self.runtime.handle().enter(); - let Some(frame) = Self::poll_ready(resp)? else { + let Some(frame) = Self::next_frame(resp).now_or_never()? else { *state = State::Closed; - return Some(Err(end().into())); + return Some(Err(iteration.end().into())); }; if frame.is_err() { *state = State::Closed; @@ -181,7 +198,7 @@ impl Streamer { // Confirm the end before returning: an unpolled body's read timeout keeps running // while the caller works on this frame. Remaining frames go to the read-ahead task, // so a stalled caller does not stall the network read. - match Self::poll_ready(resp) { + match Self::next_frame(resp).now_or_never() { Some(None) => *state = State::Closed, next => { if let State::Idle(resp) = mem::replace(&mut *state, State::Closed) { @@ -192,18 +209,6 @@ impl Streamer { Some(frame) } - /// Poll the body once: `None` if it must wait, `Some(None)` at its end. - fn poll_ready(resp: &mut wreq::Response) -> Option>> { - loop { - let Some(frame) = resp.frame().now_or_never()? else { - return Some(None); - }; - if let Some(frame) = Self::convert(frame, resp) { - return Some(Some(frame)); - } - } - } - /// Start the read-ahead task after `first`, a frame already read from `resp`. fn read_ahead_from(&self, resp: wreq::Response, first: Option>) -> State { let (tx, rx) = mpsc::channel(Self::READ_AHEAD); @@ -225,9 +230,10 @@ impl Streamer { } } - /// Start the read-ahead task if needed, then return a buffered frame, `end()` once the - /// body is done or closed, or `None` if the read must wait for `arrived`. - fn try_next(&self, end: fn() -> Error) -> Option> { + /// Start the read-ahead task if needed, then return a buffered frame, the end of + /// `iteration` once the body is done or closed, or `None` if the read must wait for + /// `arrived`. + fn try_next(&self, iteration: Iteration) -> Option> { let mut state = self.reader.lock(); if let State::Idle(_) = *state && let State::Idle(resp) = mem::replace(&mut *state, State::Closed) @@ -235,14 +241,14 @@ impl Streamer { *state = self.read_ahead_from(*resp, None); } let State::Reading { rx, .. } = &mut *state else { - return Some(Err(end().into())); + return Some(Err(iteration.end().into())); }; match rx.try_recv() { Ok(frame) => Some(frame), Err(TryRecvError::Empty) => None, Err(TryRecvError::Disconnected) => { *state = State::Closed; - Some(Err(end().into())) + Some(Err(iteration.end().into())) } } } @@ -258,13 +264,13 @@ impl Streamer { Ok(()) } - /// Wait for the next frame; `end` builds the error that ends iteration. - async fn next(&self, end: fn() -> Error) -> PyResult { + /// Wait for the next frame. + async fn next(&self, iteration: Iteration) -> PyResult { loop { // Readers are woken by `notify_waiters`, which reaches a `Notified` from its // creation, so it needs no `enable` and a frame sent after `try_next` still wakes it. let arrived = pin!(self.reader.arrived.notified()); - if let Some(frame) = self.try_next(end) { + if let Some(frame) = self.try_next(iteration) { return frame; } self.reader.since_yield.store(0, Ordering::Relaxed); @@ -283,11 +289,10 @@ impl Streamer { // Refused before any frame, so a nested read never leaves a stream half consumed. runtime::refuse_nested()?; // Frames already received are returned without releasing the GIL. - if let Some(frame) = self.ready_frame(false, || Error::StopIteration) { + if let Some(frame) = self.ready_frame(Iteration::Sync) { return frame; } - self.runtime - .block_on_eager(py, self.next(|| Error::StopIteration)) + self.runtime.block_on_eager(py, self.next(Iteration::Sync)) } fn __enter__(slf: PyRef) -> PyRef { @@ -320,7 +325,7 @@ impl Streamer { coroutine::local(py, "Streamer.__anext__", async move { let this = slf.get(); // A body already buffered is read without starting the read-ahead task. - if let Some(frame) = this.ready_frame(true, || Error::StopAsyncIteration) { + if let Some(frame) = this.ready_frame(Iteration::Async) { return frame; } // Buffered frames complete without suspending; yield to the event loop @@ -329,7 +334,7 @@ impl Streamer { this.reader.since_yield.store(0, Ordering::Relaxed); coroutine::yield_now().await; } - this.next(|| Error::StopAsyncIteration).await + this.next(Iteration::Async).await }) } @@ -392,3 +397,25 @@ impl Drop for NotifyOnDrop { self.0.notify_waiters(); } } + +// ===== impl Iteration ===== + +impl Iteration { + /// Whether the caller takes a frame `resp` already holds. On the event loop only an HTTP/1 + /// body of known length is read, whatever its size: taking a frame it holds is cheap, + /// while a decompressing body is not. + fn takes_ready(self, resp: &wreq::Response) -> bool { + match self { + Iteration::Sync => true, + Iteration::Async => loop_polls(resp.version()) && resp.size_hint().exact().is_some(), + } + } + + /// The error that ends iteration. + fn end(self) -> Error { + match self { + Iteration::Sync => Error::StopIteration, + Iteration::Async => Error::StopAsyncIteration, + } + } +} From 1b449514cd8aaa5ac1d6286a6f7bfb5de13e4bae Mon Sep 17 00:00:00 2001 From: 0x676e67 Date: Wed, 7 Oct 2026 07:23:18 +0000 Subject: [PATCH 10/25] refactor(client): share the close race, map resolved addresses and declare used features --- Cargo.toml | 4 ++-- src/client/req.rs | 28 ++++++++++++++++++---------- src/dns.rs | 18 +----------------- 3 files changed, 21 insertions(+), 29 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index e56a234c..420648e3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -29,7 +29,7 @@ abi3-py314 = ["pyo3/abi3-py314"] [dependencies] pingora-runtime = "0.9.0" -tokio = { version = "1.53.2", features = ["rt-multi-thread", "sync", "time", "net"] } +tokio = { version = "1.53.2", features = ["rt-multi-thread", "sync", "time", "net", "macros"] } tokio-util = { version = "0.7.18", features = ["rt"] } pyo3 = { version = "0.29.3", features = [ "indexmap", @@ -37,7 +37,7 @@ pyo3 = { version = "0.29.3", features = [ "extension-module", "generate-import-lib", ] } -futures-util = { version = "0.3.33", default-features = false, features = ["std"] } +futures-util = { version = "0.3.33", default-features = false, features = ["std", "sink"] } serde = { version = "1.0.228", features = ["derive"] } serde_json = "1" serde_urlencoded = "0.7.1" diff --git a/src/client/req.rs b/src/client/req.rs index 3b079bd2..ed1a0a82 100644 --- a/src/client/req.rs +++ b/src/client/req.rs @@ -1,11 +1,14 @@ use std::{ + future::Future, net::{IpAddr, Ipv4Addr, Ipv6Addr}, + pin::{Pin, pin}, time::Duration, }; use futures_util::TryFutureExt; use http::header::COOKIE; use pyo3::{PyResult, exceptions::asyncio::CancelledError, prelude::*, pybacked::PyBackedStr}; +use tokio_util::sync::CancellationToken; use crate::{ client::{ @@ -268,7 +271,7 @@ pub async fn execute_request( where U: AsRef, { - let future = async { + let future = pin!(async { // Create the request builder. let mut builder = client.inner.request(method.into_ffi(), url.as_ref()); @@ -398,13 +401,9 @@ where .map(|response| Response::new(response, client.runtime.clone())) .map_err(Error::Library) .map_err(Into::into) - }; + }); - tokio::select! { - biased; - _ = client.cancel.cancelled() => Err(CancelledError::new_err("Operation was cancelled: client has been closed")), - result = future => result, - } + until_closed(&client.cancel, future).await } /// Like [`execute_request`], opening a WebSocket. @@ -416,7 +415,7 @@ pub async fn execute_websocket_request( where U: AsRef, { - let future = async { + let future = pin!(async { // Create the WebSocket builder. let mut builder = client.inner.websocket(url.as_ref()); @@ -536,11 +535,20 @@ where .await .map_err(Error::Library) .map_err(Into::into) - }; + }); + + until_closed(&client.cancel, future).await +} +/// Run `future` unless the client is closed first; a closed client never polls it. Taking +/// it pinned keeps the request future stored once, in the caller. +async fn until_closed( + cancel: &CancellationToken, + future: Pin<&mut impl Future>>, +) -> PyResult { tokio::select! { biased; - _ = client.cancel.cancelled() => Err(CancelledError::new_err("Operation was cancelled: client has been closed")), + _ = cancel.cancelled() => Err(CancelledError::new_err("Operation was cancelled: client has been closed")), result = future => result, } } diff --git a/src/dns.rs b/src/dns.rs index 519bcfa1..704357c0 100644 --- a/src/dns.rs +++ b/src/dns.rs @@ -9,7 +9,6 @@ use futures_util::FutureExt; use hickory_resolver::{ TokioResolver, config::{CLOUDFLARE, ResolverConfig}, - lookup_ip::LookupIpIntoIter, net::runtime::TokioRuntimeProvider, }; use pyo3::{prelude::*, pybacked::PyBackedStr}; @@ -101,29 +100,14 @@ impl HickoryResolver { } } -struct SocketAddrs { - iter: LookupIpIntoIter, -} - impl Resolve for HickoryResolver { fn resolve(&self, name: Name) -> Resolving { let resolver = self.clone(); async move { let lookup = resolver.resolver.lookup_ip(name.as_str()).await?; - let addrs: Addrs = Box::new(SocketAddrs { - iter: lookup.into_iter(), - }); + let addrs: Addrs = Box::new(lookup.into_iter().map(|ip| SocketAddr::new(ip, 0))); Ok(addrs) } .boxed() } } - -impl Iterator for SocketAddrs { - type Item = SocketAddr; - - #[inline(always)] - fn next(&mut self) -> Option { - self.iter.next().map(|ip_addr| SocketAddr::new(ip_addr, 0)) - } -} From e3fa389ab8f4d732094637f833c4a812445702c7 Mon Sep 17 00:00:00 2001 From: 0x676e67 Date: Wed, 7 Oct 2026 09:22:04 +0000 Subject: [PATCH 11/25] refactor(runtime): route every runtime operation through Runtime --- clippy.toml | 11 ++ src/client.rs | 4 +- src/client/body/multipart.rs | 4 +- src/client/body/stream/async_gen.rs | 10 +- src/client/body/stream/iter.rs | 11 +- src/client/body/stream/pump.rs | 19 ++-- src/client/resp/http.rs | 6 +- src/client/resp/stream.rs | 17 ++-- src/client/resp/ws.rs | 2 +- src/client/resp/ws/cmd.rs | 10 +- src/redirect.rs | 5 +- src/runtime.rs | 151 ++++++++++++++++++---------- 12 files changed, 148 insertions(+), 102 deletions(-) create mode 100644 clippy.toml diff --git a/clippy.toml b/clippy.toml new file mode 100644 index 00000000..9f706ea0 --- /dev/null +++ b/clippy.toml @@ -0,0 +1,11 @@ +# Runtime operations go through `crate::runtime::Runtime`, which keeps blocking calls, +# current-thread driving and the shared runtime consistent. +disallowed-methods = [ + { path = "tokio::task::spawn", reason = "use Runtime::spawn" }, + { path = "tokio::task::spawn_blocking", reason = "use Runtime::spawn_blocking" }, + { path = "tokio::task::yield_now", reason = "use Runtime::defer_wake" }, + { path = "tokio::task::block_in_place", reason = "use Runtime::block_on_in_pool" }, + { path = "tokio::runtime::Handle::current", reason = "use Runtime" }, + { path = "tokio::runtime::Handle::try_current", reason = "use Runtime" }, +] +disallowed-types = [{ path = "tokio::runtime::Handle", reason = "use Runtime" }] diff --git a/src/client.rs b/src/client.rs index 88e9e4d9..1ece1bad 100644 --- a/src/client.rs +++ b/src/client.rs @@ -300,7 +300,7 @@ impl Client { impl Default for Client { fn default() -> Self { - let runtime = runtime::get(); + let runtime = runtime::Runtime::shared(); Self { inner: wreq::Client::default(), runtime: runtime.clone(), @@ -317,7 +317,7 @@ impl Client { py.detach(|| { let runtime = match kwds.as_ref().and_then(|config| config.runtime.as_ref()) { Some(runtime) => runtime.select()?, - None => runtime::get().clone(), + None => runtime::Runtime::shared().clone(), }; // Create the client builder. let mut builder = wreq::Client::builder(); diff --git a/src/client/body/multipart.rs b/src/client/body/multipart.rs index 3a6d7cab..1cf8b1b1 100644 --- a/src/client/body/multipart.rs +++ b/src/client/body/multipart.rs @@ -9,7 +9,7 @@ use crate::{ error::Error, extractor::{Binary, Text}, header::HeaderMap, - runtime, + runtime::Runtime, }; /// A multipart form for a request. @@ -136,7 +136,7 @@ impl Part { Value::Text(text) => multipart::Part::stream(text.0), Value::Bytes(bytes) => multipart::Part::stream(bytes.0), // Opening the file blocks, so only that waits detached. - Value::File(path) => runtime::get().block_on( + Value::File(path) => Runtime::shared().block_on( py, multipart::Part::file(path).map_err(|err| Error::from(err).into()), )?, diff --git a/src/client/body/stream/async_gen.rs b/src/client/body/stream/async_gen.rs index ce450676..411d3dc8 100644 --- a/src/client/body/stream/async_gen.rs +++ b/src/client/body/stream/async_gen.rs @@ -12,10 +12,10 @@ use std::{ use futures_util::{Stream, StreamExt}; use pyo3::{IntoPyObjectExt, exceptions::PyRuntimeError, intern, prelude::*, sync::PyOnceLock}; -use tokio::{runtime::Handle, sync::mpsc::error::TrySendError}; +use tokio::sync::mpsc::error::TrySendError; use super::{Item, PyBytesLike, lock, queue}; -use crate::{coroutine, runtime}; +use crate::{coroutine, runtime::Runtime}; /// A request body from a Python async generator, forwarded within the upload budget by a /// task on the loop that was running at extraction. Dropping it cancels that task. @@ -99,10 +99,8 @@ impl Drop for AsyncStream { // Senders see the body gone at once, and a send waiting for room resolves to `False`. self.rx.close(); if let Some((task, event_loop)) = self.task.take() { - // Body drop can run on Tokio: cancel from a blocking thread, preferring the current - // runtime so a drop on a client's own runtime does not start the shared one. - let handle = Handle::try_current().unwrap_or_else(|_| runtime::get().handle().clone()); - handle.spawn_blocking(move || { + // Body drop can run on Tokio: cancel from a blocking thread. + Runtime::spawn_blocking(move || { Python::try_attach(|py| { if let Ok(cancel) = task.bind(py).getattr(intern!(py, "cancel")) { let _ = event_loop.call_method1( diff --git a/src/client/body/stream/iter.rs b/src/client/body/stream/iter.rs index c5ca948e..75336ce3 100644 --- a/src/client/body/stream/iter.rs +++ b/src/client/body/stream/iter.rs @@ -2,7 +2,7 @@ use std::{ mem, - pin::{Pin, pin}, + pin::Pin, task::{Context, Poll, ready}, thread::{self, ThreadId}, }; @@ -12,10 +12,9 @@ use pyo3::{ exceptions::{PyException, PyKeyboardInterrupt}, prelude::*, }; -use tokio::task::yield_now; use super::{Item, next_item, pump::Pumped}; -use crate::runtime; +use crate::runtime::Runtime; /// A request body from a Python iterator. /// @@ -65,7 +64,7 @@ impl Stream for SyncStream { if let SyncStream::Idle { .. } = this && let SyncStream::Idle { iter, owner } = mem::replace(this, SyncStream::Done) { - *this = if runtime::driving() { + *this = if Runtime::driving() { SyncStream::Inline(Inline { iter, owner, @@ -101,7 +100,7 @@ impl Inline { return Poll::Ready(()); } self.flushed = true; - let _ = pin!(yield_now()).poll(cx); + Runtime::defer_wake(cx); Poll::Pending } } @@ -122,7 +121,7 @@ impl Stream for Inline { && (this.owner == thread::current().id() || err.is_instance_of::(py)) { - runtime::interrupt(err.clone_ref(py)); + Runtime::interrupt(err.clone_ref(py)); } item }) diff --git a/src/client/body/stream/pump.rs b/src/client/body/stream/pump.rs index 6f790114..f3ea91bd 100644 --- a/src/client/body/stream/pump.rs +++ b/src/client/body/stream/pump.rs @@ -9,12 +9,13 @@ use std::{ use futures_util::{Stream, StreamExt}; use pyo3::prelude::*; -use tokio::{runtime::Handle, sync::mpsc::error::TrySendError, task::spawn_blocking, time}; +use tokio::{sync::mpsc::error::TrySendError, time}; use super::{ Item, PyBytesLike, lock, next_item, queue::{self, Permit}, }; +use crate::runtime::Runtime; /// An iterator read ahead by pump tasks on the blocking pool, within the upload budget. When /// the budget stays full, a pump parks the iterator with the chunk it read and frees its @@ -62,7 +63,7 @@ impl Stream for Pumped { let parked = lock(&this.parked).take(); if let Some(pump) = parked { let parked = this.parked.clone(); - spawn_blocking(move || pump.run(&parked)); + Runtime::spawn_blocking(move || pump.run(&parked)); } poll } @@ -77,7 +78,6 @@ impl Pump { /// Queue chunks until the iterator ends or raises, the body is dropped, or the budget /// stays full past [`PARK_AFTER`](Self::PARK_AFTER). fn run(mut self, parked: &Mutex>) { - let handle = Handle::current(); // Once Python is unavailable, stop reading without creating a PyErr that could // require another attachment to format. Python::try_attach(|py| { @@ -93,7 +93,7 @@ impl Pump { None => return, }, }; - let permit = match self.wait_for_room(py, &handle, &chunk) { + let permit = match self.wait_for_room(py, &chunk) { Some(permit) => permit, None => { // Park under the lock the body takes after receiving, so either it @@ -119,18 +119,13 @@ impl Pump { /// Reserve room for `chunk`, waiting detached up to [`PARK_AFTER`](Self::PARK_AFTER) while /// the budget is full; `None` if it stays full or the body is dropped. - fn wait_for_room( - &self, - py: Python<'_>, - handle: &Handle, - chunk: &PyBytesLike, - ) -> Option> { + fn wait_for_room(&self, py: Python<'_>, chunk: &PyBytesLike) -> Option> { match self.tx.try_reserve(chunk) { Ok(permit) => Some(permit), Err(TrySendError::Closed(())) => None, Err(TrySendError::Full(())) => { - let wait = time::timeout(Self::PARK_AFTER, self.tx.reserve(chunk)); - py.detach(|| handle.block_on(wait)).ok().flatten() + let wait = async { time::timeout(Self::PARK_AFTER, self.tx.reserve(chunk)).await }; + Runtime::block_on_in_pool(py, wait).ok().flatten() } } } diff --git a/src/client/resp/http.rs b/src/client/resp/http.rs index 53084b57..7b5e1cd5 100644 --- a/src/client/resp/http.rs +++ b/src/client/resp/http.rs @@ -23,7 +23,7 @@ use crate::{ header::HeaderMap, http::{StatusCode, Version}, redirect::History, - runtime::{self, Runtime}, + runtime::Runtime, tls::TlsInfo, }; @@ -487,7 +487,7 @@ impl BlockingResponse { T: Send, { // Refused before the body is taken, so it can still be read afterwards. - runtime::refuse_nested()?; + Runtime::refuse_nested()?; let runtime = &self.0.runtime; let (read, inline) = self.0.read_body(Some(READ_ATTACHED), read)?; if inline { @@ -583,7 +583,7 @@ impl BlockingResponse { /// Read the body as a read-only memoryview, retaining its data after the response closes. pub fn bytes(&self, py: Python) -> PyResult { - runtime::refuse_nested()?; + Runtime::refuse_nested()?; let response = &self.0; match response.take_bytes(Some(READ_ATTACHED))? { BodyRead::Ready(bytes) => Ok(PyBuffer::from(bytes)), diff --git a/src/client/resp/stream.rs b/src/client/resp/stream.rs index eae6de0c..13bc2576 100644 --- a/src/client/resp/stream.rs +++ b/src/client/resp/stream.rs @@ -28,7 +28,7 @@ use crate::{ coroutine::{self, Coroutine}, error::Error, header::HeaderMap, - runtime::{self, Runtime}, + runtime::Runtime, }; /// A response frame exposed as a read-only memoryview or a header map. @@ -185,9 +185,7 @@ impl Streamer { if !iteration.takes_ready(resp) { return None; } - // A read timeout starts a timer on its first poll, which needs the runtime. - let _runtime = self.runtime.handle().enter(); - let Some(frame) = Self::next_frame(resp).now_or_never()? else { + let Some(frame) = self.runtime.poll_now(Self::next_frame(resp))? else { *state = State::Closed; return Some(Err(iteration.end().into())); }; @@ -198,7 +196,7 @@ impl Streamer { // Confirm the end before returning: an unpolled body's read timeout keeps running // while the caller works on this frame. Remaining frames go to the read-ahead task, // so a stalled caller does not stall the network read. - match Self::next_frame(resp).now_or_never() { + match self.runtime.poll_now(Self::next_frame(resp)) { Some(None) => *state = State::Closed, next => { if let State::Idle(resp) = mem::replace(&mut *state, State::Closed) { @@ -220,10 +218,9 @@ impl Streamer { return State::Reading { rx, _task: None }; } } - let task = - self.runtime - .handle() - .spawn(Self::read_ahead(resp, tx, self.reader.arrived.clone())); + let task = self + .runtime + .spawn(Self::read_ahead(resp, tx, self.reader.arrived.clone())); State::Reading { rx, _task: Some(AbortOnDropHandle::new(task)), @@ -287,7 +284,7 @@ impl Streamer { fn __next__(&self, py: Python) -> PyResult { // Refused before any frame, so a nested read never leaves a stream half consumed. - runtime::refuse_nested()?; + Runtime::refuse_nested()?; // Frames already received are returned without releasing the GIL. if let Some(frame) = self.ready_frame(Iteration::Sync) { return frame; diff --git a/src/client/resp/ws.rs b/src/client/resp/ws.rs index da2616d7..a8ddd72c 100644 --- a/src/client/resp/ws.rs +++ b/src/client/resp/ws.rs @@ -65,7 +65,7 @@ impl WebSocket { ); let websocket = response.into_websocket().await?; let protocol = websocket.protocol().cloned(); - let cmd = cmd::spawn(websocket); + let cmd = cmd::spawn(&runtime, websocket); Ok(WebSocket { runtime, diff --git a/src/client/resp/ws/cmd.rs b/src/client/resp/ws/cmd.rs index 286667bf..fcfb013d 100644 --- a/src/client/resp/ws/cmd.rs +++ b/src/client/resp/ws/cmd.rs @@ -26,7 +26,7 @@ use wreq::ws::{ }; use super::Message; -use crate::{error::Error, extractor::Text}; +use crate::{error::Error, extractor::Text, runtime::Runtime}; /// The request channels of a WebSocket's read and write tasks. #[derive(Clone)] @@ -45,14 +45,14 @@ enum Write { Close(Option, Option, Sender>), } -/// Start the read and write tasks of `ws` on the current runtime. -pub fn spawn(ws: WebSocket) -> Handle { +/// Start the read and write tasks of `ws` on `runtime`'s selected worker. +pub fn spawn(runtime: &Runtime, ws: WebSocket) -> Handle { let (writer, reader) = ws.split(); let (reads, read_rx) = mpsc::unbounded_channel(); let (writes, write_rx) = mpsc::unbounded_channel(); let closed = CancellationToken::new(); - tokio::spawn(read(reader, read_rx, closed.clone())); - tokio::spawn(write(writer, write_rx, closed)); + runtime.spawn(read(reader, read_rx, closed.clone())); + runtime.spawn(write(writer, write_rx, closed)); Handle { reads, writes } } diff --git a/src/redirect.rs b/src/redirect.rs index 191b0bd0..c8b0ddbb 100644 --- a/src/redirect.rs +++ b/src/redirect.rs @@ -4,9 +4,8 @@ use std::{ }; use pyo3::prelude::*; -use tokio::task::spawn_blocking; -use crate::{header::HeaderMap, http::StatusCode}; +use crate::{header::HeaderMap, http::StatusCode, runtime::Runtime}; /// Represents the redirect policy for HTTP requests. #[derive(Clone)] @@ -106,7 +105,7 @@ impl Policy { let callback = callback.clone(); attempt.pending(|attempt| async move { let args = Attempt::from(&attempt); - let kind = spawn_blocking(move || { + let kind = Runtime::spawn_blocking(move || { Python::try_attach(|py| { callback .call1(py, (args,)) diff --git a/src/runtime.rs b/src/runtime.rs index 8724a410..484d0afc 100644 --- a/src/runtime.rs +++ b/src/runtime.rs @@ -4,12 +4,17 @@ //! ([`Runtime::block_on_task`], [`Runtime::run_task`]); blocking reads are polled by the //! caller ([`Runtime::block_on`], [`Runtime::block_on_eager`]). The `block_on*` calls release //! the GIL while they wait, and drive a current-thread runtime meanwhile. +//! +//! Every other runtime operation goes through [`Runtime`] too; `clippy.toml` rejects direct +//! Tokio handles, spawns and yields elsewhere. +#![allow(clippy::disallowed_methods, clippy::disallowed_types)] use std::{ cell::{Cell, RefCell}, future::Future, pin::pin, sync::{Arc, OnceLock}, + task::Context, thread, time::Duration, }; @@ -22,7 +27,7 @@ use pyo3::{ }; use tokio::{ runtime::{Builder, Handle}, - task::JoinError, + task::{self, JoinError, JoinHandle}, }; use tokio_util::task::AbortOnDropHandle; @@ -69,25 +74,9 @@ thread_local! { static INTERRUPT: RefCell> = const { RefCell::new(None) }; } -/// Whether this thread drives a current-thread runtime, which then runs the Python code -/// its tasks call, such as upload iterators. -pub fn driving() -> bool { - DRIVING.get() -} - -/// Re-raise `err` from the blocking call this thread is driving, once it returns. -pub fn interrupt(err: PyErr) { - INTERRUPT.set(Some(err)); -} - // ===== impl Runtime ===== impl Runtime { - /// Borrow the selected worker's handle without changing the selection. - pub fn handle(&self) -> &Handle { - &self.handle - } - /// Whether blocking calls drive this runtime on their own thread. pub fn is_current_thread(&self) -> bool { matches!(self.inner.as_deref(), Some(Shared::CurrentThread(_))) @@ -106,8 +95,8 @@ impl Runtime { } /// Poll `future` once inside the runtime, where a read timeout can start its timer; - /// `None` if it must wait. - pub fn poll_now(&self, future: &mut F) -> Option { + /// `None` if it must wait. A pending future is dropped, so pass `&mut` to keep it. + pub fn poll_now(&self, future: F) -> Option { let _runtime = self.handle.enter(); future.now_or_never() } @@ -130,9 +119,9 @@ impl Runtime { T: Send, { // Refused before the first poll, which a later refusal would leave half done. - refuse_nested()?; + Self::refuse_nested()?; let mut future = pin!(future); - match self.poll_now(&mut future) { + match self.poll_now(future.as_mut()) { Some(output) => output, None => self.block_on(py, future), } @@ -163,10 +152,20 @@ impl Runtime { AbortOnDropHandle::new(task).await.map_err(join_error)? } + /// Spawn `future` on the selected worker. Dropping the handle detaches the task, which + /// does not keep the runtime alive. + pub fn spawn(&self, future: F) -> JoinHandle + where + F: Future + Send + 'static, + F::Output: Send + 'static, + { + self.handle.spawn(future) + } + /// [`block_on`](Self::block_on) for a caller already detached. It refuses a nested call /// itself, where Tokio would panic. fn wait(&self, future: impl Future>) -> PyResult { - refuse_nested()?; + Self::refuse_nested()?; match self.inner.as_deref() { Some(Shared::CurrentThread(runtime)) => Driving::drive(runtime, future), _ => self.handle.block_on(future), @@ -180,7 +179,7 @@ impl Runtime { F: Future> + Send + 'static, T: Send + 'static, { - refuse_nested()?; + Self::refuse_nested()?; match self.inner.as_deref() { Some(Shared::CurrentThread(runtime)) => Driving::drive(runtime, future), _ => self @@ -191,6 +190,78 @@ impl Runtime { } } +/// Operations on the calling thread: the shared runtime, the runtime the thread is in, and +/// the state of a current-thread runtime it drives. +impl Runtime { + /// The shared runtime, created on first use and retained for the process lifetime. + pub fn shared() -> &'static Runtime { + static RUNTIME: OnceLock = OnceLock::new(); + + fn create() -> Runtime { + let runtime = RuntimeBuilder::new(parallelism(), env!("CARGO_PKG_NAME")).build(); + runtime.into() + } + + if let Some(runtime) = RUNTIME.get() { + return runtime; + } + + // Never wait for another initializer while holding the interpreter. + Python::try_attach(|py| py.detach(|| RUNTIME.get_or_init(create))) + .unwrap_or_else(|| RUNTIME.get_or_init(create)) + } + + /// Run `f` on the blocking pool of the runtime this thread is in, or the shared runtime's + /// outside one, so a drop on a client's runtime does not start the shared one. + pub fn spawn_blocking(f: F) -> JoinHandle + where + F: FnOnce() -> R + Send + 'static, + R: Send + 'static, + { + current().spawn_blocking(f) + } + + /// Wait for `future` detached from Python on a blocking-pool thread, driven by the + /// runtime that owns the pool. Unlike [`block_on`](Self::block_on), it neither refuses + /// nested calls nor needs a `PyResult`; create timers inside `future`. + pub fn block_on_in_pool(py: Python<'_>, future: F) -> F::Output + where + F: Future + Send, + F::Output: Send, + { + let handle = current(); + py.detach(|| handle.block_on(future)) + } + + /// Wake the task after the runtime polls its IO and other tasks, or at once outside a + /// runtime. + pub fn defer_wake(cx: &mut Context<'_>) { + let _ = pin!(task::yield_now()).poll(cx); + } + + /// Whether this thread drives a current-thread runtime, which then runs the Python code + /// its tasks call, such as upload iterators. + pub fn driving() -> bool { + DRIVING.get() + } + + /// Re-raise `err` from the blocking call this thread is driving, once it returns. + pub fn interrupt(err: PyErr) { + INTERRUPT.set(Some(err)); + } + + /// Fail a blocking wait on a thread driving a current-thread runtime, where Tokio would + /// panic, as when an upload iterator it reads sends a blocking request. + pub fn refuse_nested() -> PyResult<()> { + if Self::driving() { + return Err(PyRuntimeError::new_err( + "Cannot make a blocking call from code a CURRENT_THREAD runtime is running", + )); + } + Ok(()) + } +} + impl From for Runtime { fn from(runtime: Shared) -> Self { let handle = runtime.handle(); @@ -339,17 +410,6 @@ impl Drop for Driving { } } -/// Fail a blocking wait on a thread driving a current-thread runtime, where Tokio would -/// panic, as when an upload iterator it reads sends a blocking request. -pub fn refuse_nested() -> PyResult<()> { - if driving() { - return Err(PyRuntimeError::new_err( - "Cannot make a blocking call from code a CURRENT_THREAD runtime is running", - )); - } - Ok(()) -} - /// A task that panicked or was aborted, as a Python error. fn join_error(err: JoinError) -> PyErr { PyRuntimeError::new_err(err.to_string()) @@ -360,22 +420,9 @@ fn parallelism() -> usize { thread::available_parallelism().map_or(1, usize::from) } -/// Create the shared runtime on first use and retain it for the process lifetime. -pub fn get() -> &'static Runtime { - static RUNTIME: OnceLock = OnceLock::new(); - - fn create() -> Runtime { - let runtime = RuntimeBuilder::new(parallelism(), env!("CARGO_PKG_NAME")).build(); - runtime.into() - } - - if let Some(runtime) = RUNTIME.get() { - return runtime; - } - - // Never wait for another initializer while holding the interpreter. - Python::try_attach(|py| py.detach(|| RUNTIME.get_or_init(create))) - .unwrap_or_else(|| RUNTIME.get_or_init(create)) +/// The runtime this thread is in, else the shared one. +fn current() -> Handle { + Handle::try_current().unwrap_or_else(|_| Runtime::shared().handle.clone()) } #[cfg(test)] @@ -447,7 +494,7 @@ mod tests { }) .unwrap(); assert!(nested); - assert!(!driving()); + assert!(!Runtime::driving()); } #[test] @@ -467,7 +514,7 @@ mod tests { .build(), ); let weak = Arc::downgrade(runtime.inner.as_ref().unwrap()); - let handle = runtime.handle().clone(); + let handle = runtime.handle.clone(); let (dropped, released) = std::sync::mpsc::channel(); let guard = NotifyOnDrop(dropped); let background = handle.spawn(async move { From 9008a96176d8841f26aa39b63a10aab081486369 Mon Sep 17 00:00:00 2001 From: 0x676e67 Date: Wed, 7 Oct 2026 09:22:14 +0000 Subject: [PATCH 12/25] refactor(multipart): open file parts when the request is built --- docs/source/guide/advanced.md | 6 ++-- src/client/body/multipart.rs | 58 ++++++++++++++--------------------- src/client/req.rs | 9 ++---- tests/multipart_test.py | 20 ++++++++++++ 4 files changed, 49 insertions(+), 44 deletions(-) diff --git a/docs/source/guide/advanced.md b/docs/source/guide/advanced.md index e88156d1..e5869d4c 100644 --- a/docs/source/guide/advanced.md +++ b/docs/source/guide/advanced.md @@ -151,9 +151,9 @@ bound or request is sent. (512 and 10 seconds). With `Scheduler.PER_WORKER` these limits apply to **each worker's** blocking pool, not the pool as a whole. Python async upload generators still run on the caller's event loop. -Standalone multipart file preparation and upload-task cleanup can use the -shared runtime; a dedicated client runtime does not isolate Python's GIL or -every process resource. DNS resolvers are owned by individual clients so their +Multipart files open on the client's runtime when the request is built; only +upload cleanup that runs outside any runtime uses the shared runtime. A dedicated +client runtime does not isolate Python's GIL or every process resource. DNS resolvers are owned by individual clients so their connections are not shared across runtimes. Closing a client cancels pending requests and rejects new requests with diff --git a/src/client/body/multipart.rs b/src/client/body/multipart.rs index 1cf8b1b1..0a350858 100644 --- a/src/client/body/multipart.rs +++ b/src/client/body/multipart.rs @@ -1,6 +1,5 @@ use std::path::PathBuf; -use futures_util::TryFutureExt; use pyo3::{prelude::*, types::PyTuple}; use wreq::{Body, multipart}; @@ -9,13 +8,12 @@ use crate::{ error::Error, extractor::{Binary, Text}, header::HeaderMap, - runtime::Runtime, }; -/// A multipart form for a request. +/// A multipart form for a request. Its files open when the request is built, on the +/// client's runtime. #[pyclass(subclass)] pub struct Multipart { - pub form: Option, pub parts: Vec, } @@ -50,22 +48,20 @@ impl Multipart { let mut new_parts = Vec::with_capacity(parts.len()); for part in parts { let part = part.cast::()?; - let mut part = part.borrow_mut(); + let mut part = part.try_borrow_mut()?; new_parts.push(part.try_clone()?); } - Ok(Self { - form: None, - parts: new_parts, - }) + Ok(Self { parts: new_parts }) } } impl Multipart { - fn build_form(&mut self, py: Python) -> PyResult { + /// Build the form to send, opening its files. + pub async fn into_form(self) -> PyResult { let mut form = multipart::Form::new(); - for part in &mut self.parts { - let (name, inner) = part.build_part(py)?; + for part in self.parts { + let (name, inner) = part.into_part().await?; form = form.part(name, inner); } Ok(form) @@ -77,13 +73,14 @@ impl FromPyObject<'_, '_> for Multipart { fn extract(ob: Borrowed) -> PyResult { let multipart = ob.cast::()?; - let mut multipart = multipart.borrow_mut(); - let form = multipart.build_form(ob.py())?; - - Ok(Multipart { - form: Some(form), - parts: Vec::new(), - }) + let mut multipart = multipart.try_borrow_mut()?; + // Copied for one request: reusable values stay, a stream moves into it. + let parts = multipart + .parts + .iter_mut() + .map(Part::try_clone) + .collect::>()?; + Ok(Multipart { parts }) } } @@ -124,22 +121,13 @@ impl Part { } } - fn build_part(&mut self, py: Python) -> PyResult<(String, multipart::Part)> { - let value = self - .value - .as_ref() - .and_then(Value::try_clone) - .or_else(|| self.value.take()) - .ok_or_else(|| Error::Memory)?; - + /// Build the part to send, opening its file. + async fn into_part(self) -> PyResult<(String, multipart::Part)> { + let value = self.value.ok_or_else(|| Error::Memory)?; let mut inner = match value { Value::Text(text) => multipart::Part::stream(text.0), Value::Bytes(bytes) => multipart::Part::stream(bytes.0), - // Opening the file blocks, so only that waits detached. - Value::File(path) => Runtime::shared().block_on( - py, - multipart::Part::file(path).map_err(|err| Error::from(err).into()), - )?, + Value::File(path) => multipart::Part::file(path).await.map_err(Error::from)?, Value::Stream(stream) => { let stream = Body::wrap_stream(stream); match self.length { @@ -149,7 +137,7 @@ impl Part { } }; - if let Some(filename) = self.filename.clone() { + if let Some(filename) = self.filename { inner = inner.file_name(filename); } @@ -157,11 +145,11 @@ impl Part { inner = inner.mime_str(mime).map_err(Error::Library)?; } - if let Some(headers) = self.headers.clone() { + if let Some(headers) = self.headers { inner = inner.headers(headers.0); } - Ok((self.name.clone(), inner)) + Ok((self.name, inner)) } /// Copy the part, moving a stream value out since it can be sent only once. diff --git a/src/client/req.rs b/src/client/req.rs index ed1a0a82..a39f3bec 100644 --- a/src/client/req.rs +++ b/src/client/req.rs @@ -372,12 +372,9 @@ where // Body options. apply_option!(set_if_some_ref, builder, request.form, form); apply_option!(set_if_some_ref, builder, request.json, json); - apply_option!( - set_if_some, - builder, - request.multipart.and_then(|form| form.form), - multipart - ); + if let Some(multipart) = request.multipart { + builder = builder.multipart(multipart.into_form().await?); + } apply_option!( set_if_some_map_try, builder, diff --git a/tests/multipart_test.py b/tests/multipart_test.py index f61a51e0..6c901c04 100644 --- a/tests/multipart_test.py +++ b/tests/multipart_test.py @@ -3,6 +3,7 @@ import pytest import wreq from wreq import Multipart, Part +from wreq.runtime import Runtime, Scheduler client = wreq.Client(tls_info=True) @@ -90,3 +91,22 @@ def bytes_stream(): with pytest.raises(RuntimeError): Multipart(part) + + +@pytest.mark.asyncio +async def test_file_parts_open_when_the_request_is_built(tmp_path): + url = "http://localhost:8080/post" + path = tmp_path / "part.txt" + path.write_bytes(b"file part") + blocking = wreq.blocking.Client(runtime=Runtime(scheduler=Scheduler.CURRENT_THREAD)) + + # A current-thread client opens the file on its own runtime, driven by the caller. + with blocking.post(url, multipart=Multipart(Part(name="f", value=path))) as resp: + assert resp.json()["files"]["f"] == "file part" + + # A missing file fails the request rather than building the form. + form = Multipart(Part(name="f", value=tmp_path / "missing.txt")) + with pytest.raises(RuntimeError, match="IO error"): + blocking.post(url, multipart=form) + with pytest.raises(RuntimeError, match="IO error"): + await client.post(url, multipart=form) From 9e3af1df7aab0e0e04c11055afe7e368b41fd0b9 Mon Sep 17 00:00:00 2001 From: 0x676e67 Date: Wed, 7 Oct 2026 09:40:36 +0000 Subject: [PATCH 13/25] test(multipart): assert a closed client never opens file parts --- tests/multipart_test.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/multipart_test.py b/tests/multipart_test.py index 6c901c04..144ba359 100644 --- a/tests/multipart_test.py +++ b/tests/multipart_test.py @@ -1,3 +1,4 @@ +import asyncio from pathlib import Path import pytest @@ -110,3 +111,8 @@ async def test_file_parts_open_when_the_request_is_built(tmp_path): blocking.post(url, multipart=form) with pytest.raises(RuntimeError, match="IO error"): await client.post(url, multipart=form) + + # A closed client rejects the request before building the form, so it never opens. + blocking.close() + with pytest.raises(asyncio.CancelledError): + blocking.post(url, multipart=form) From aca3074d15553a58de11528385ade623d4ab5b4e Mon Sep 17 00:00:00 2001 From: 0x676e67 Date: Wed, 7 Oct 2026 09:40:53 +0000 Subject: [PATCH 14/25] chore(clippy): point the block_in_place lint at spawn_blocking --- clippy.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/clippy.toml b/clippy.toml index 9f706ea0..d0673276 100644 --- a/clippy.toml +++ b/clippy.toml @@ -4,7 +4,7 @@ disallowed-methods = [ { path = "tokio::task::spawn", reason = "use Runtime::spawn" }, { path = "tokio::task::spawn_blocking", reason = "use Runtime::spawn_blocking" }, { path = "tokio::task::yield_now", reason = "use Runtime::defer_wake" }, - { path = "tokio::task::block_in_place", reason = "use Runtime::block_on_in_pool" }, + { path = "tokio::task::block_in_place", reason = "use Runtime::spawn_blocking, and Runtime::block_on_in_pool inside it" }, { path = "tokio::runtime::Handle::current", reason = "use Runtime" }, { path = "tokio::runtime::Handle::try_current", reason = "use Runtime" }, ] From 669054e7cd5472cd5427ec01af79f1ee73b931be Mon Sep 17 00:00:00 2001 From: 0x676e67 Date: Wed, 7 Oct 2026 09:49:22 +0000 Subject: [PATCH 15/25] docs(upload): describe how far uploads read ahead of the budget --- docs/source/guide/advanced.md | 14 ++++++++------ src/client/body/stream.rs | 4 ++-- src/client/body/stream/queue.rs | 2 +- 3 files changed, 11 insertions(+), 9 deletions(-) diff --git a/docs/source/guide/advanced.md b/docs/source/guide/advanced.md index e5869d4c..36db6d05 100644 --- a/docs/source/guide/advanced.md +++ b/docs/source/guide/advanced.md @@ -13,12 +13,14 @@ Their exceptions fail the request. When an upload ends, the generator is closed; cancelling or dropping the upload schedules producer cancellation and cleanup on that loop. Keep the loop running until generator cleanup has finished. This also applies to async multipart parts. Construct async-generator `Part` objects inside a running event loop; their producers -start at construction. A producer runs at most 256 KiB, or 64 small chunks, ahead of the -upload. -Use synchronous iterators for blocking uploads. They run on the runtime's blocking pool, -also at most 256 KiB or 64 small chunks ahead, except on a current-thread runtime. A -blocking call on the producer's event-loop thread prevents async generators from -progressing. +start at construction. A producer reads ahead until its queued chunks fill a 256 KiB +budget. Each chunk counts as at least 4 KiB, so at most 64 small chunks queue, and a chunk +above 256 KiB queues alone at its full size. One more chunk can wait for room outside the +budget, so memory held ahead of the upload also grows with the chunk size. +Use synchronous iterators for blocking uploads. They run on the runtime's blocking pool +with the same budget, except on a current-thread runtime, which reads each chunk only when +the upload needs it. A blocking call on the producer's event-loop thread prevents async +generators from progressing. ```python import asyncio diff --git a/src/client/body/stream.rs b/src/client/body/stream.rs index 5b93b9f6..4282fc1c 100644 --- a/src/client/body/stream.rs +++ b/src/client/body/stream.rs @@ -1,8 +1,8 @@ //! Request bodies streamed from Python iterators and async generators. //! //! `iter` reads a Python iterator, inline on a current-thread runtime or ahead of the upload -//! through `pump`, and `async_gen` forwards an async generator from its event loop. Both -//! queue chunks through `queue`, which bounds the bytes read ahead of the connection. +//! through `pump`, and `async_gen` forwards an async generator from its event loop. The pump +//! and the forwarder queue chunks through `queue`, which bounds how far they read ahead. mod async_gen; mod iter; diff --git a/src/client/body/stream/queue.rs b/src/client/body/stream/queue.rs index 32cc0d51..be3b9840 100644 --- a/src/client/body/stream/queue.rs +++ b/src/client/body/stream/queue.rs @@ -16,7 +16,7 @@ use tokio::sync::{ use super::{Item, PyBytesLike}; -/// Bytes an upload may queue ahead of the connection. +/// Chunk bytes an upload may queue ahead of the connection; a larger chunk queues alone. const UPLOAD_BUDGET: usize = 256 * 1024; /// The least budget a queued chunk holds, so small chunks queue at most 64 items. From c0498018d2670057c9757f39d947fea08b1d5d81 Mon Sep 17 00:00:00 2001 From: 0x676e67 Date: Wed, 7 Oct 2026 13:49:54 +0000 Subject: [PATCH 16/25] fix(upload): end uploads explicitly and give multipart iterators to their sender --- src/client/body/multipart.rs | 11 ++++- src/client/body/stream.rs | 10 ++++ src/client/body/stream/async_gen.rs | 45 ++++++----------- src/client/body/stream/iter.rs | 45 ++++++++--------- src/client/body/stream/pump.rs | 19 +++++--- src/client/body/stream/queue.rs | 75 +++++++++++++++++++---------- 6 files changed, 119 insertions(+), 86 deletions(-) diff --git a/src/client/body/multipart.rs b/src/client/body/multipart.rs index 0a350858..6ff8bd06 100644 --- a/src/client/body/multipart.rs +++ b/src/client/body/multipart.rs @@ -74,11 +74,18 @@ impl FromPyObject<'_, '_> for Multipart { fn extract(ob: Borrowed) -> PyResult { let multipart = ob.cast::()?; let mut multipart = multipart.try_borrow_mut()?; - // Copied for one request: reusable values stay, a stream moves into it. + // Copied for one request: reusable values stay, a stream moves into it and belongs + // to the thread sending it. let parts = multipart .parts .iter_mut() - .map(Part::try_clone) + .map(|part| { + let mut part = part.try_clone()?; + if let Some(Value::Stream(stream)) = &mut part.value { + stream.claim(); + } + Ok(part) + }) .collect::>()?; Ok(Multipart { parts }) } diff --git a/src/client/body/stream.rs b/src/client/body/stream.rs index 4282fc1c..8fff96a6 100644 --- a/src/client/body/stream.rs +++ b/src/client/body/stream.rs @@ -90,6 +90,16 @@ impl FromPyObject<'_, '_> for PyStream { } } +impl PyStream { + /// Hand an unread iterator to a request the calling thread sends; an async generator + /// keeps forwarding on its own loop. + pub fn claim(&mut self) { + if let Either::Left(stream) = &mut self.0 { + stream.claim(); + } + } +} + impl Stream for PyStream { type Item = Item; diff --git a/src/client/body/stream/async_gen.rs b/src/client/body/stream/async_gen.rs index 411d3dc8..7bd7b656 100644 --- a/src/client/body/stream/async_gen.rs +++ b/src/client/body/stream/async_gen.rs @@ -3,16 +3,13 @@ use std::{ ffi::CStr, pin::Pin, - sync::{ - Arc, Mutex, - atomic::{AtomicBool, Ordering}, - }, + sync::Mutex, task::{Context, Poll, ready}, }; use futures_util::{Stream, StreamExt}; -use pyo3::{IntoPyObjectExt, exceptions::PyRuntimeError, intern, prelude::*, sync::PyOnceLock}; -use tokio::sync::mpsc::error::TrySendError; +use pyo3::{IntoPyObjectExt, intern, prelude::*, sync::PyOnceLock}; +use tokio::sync::TryAcquireError; use super::{Item, PyBytesLike, lock, queue}; use crate::{coroutine, runtime::Runtime}; @@ -21,8 +18,6 @@ use crate::{coroutine, runtime::Runtime}; /// task on the loop that was running at extraction. Dropping it cancels that task. pub(super) struct AsyncStream { rx: queue::Rx, - /// Set by [`Sender::finish`], so the queue closing reads as the end of the body. - finished: Arc, /// The forwarding task and its loop, until forwarding ends. task: Option<(Py, Py)>, } @@ -33,7 +28,6 @@ pub(super) struct AsyncStream { #[pyclass(frozen)] struct Sender { tx: Mutex>, - finished: Arc, } // ===== impl AsyncStream ===== @@ -55,10 +49,8 @@ impl AsyncStream { .map(Bound::unbind) })?; let (tx, rx) = queue::channel(); - let finished = Arc::new(AtomicBool::new(false)); let sender = Sender { tx: Mutex::new(Some(tx)), - finished: finished.clone(), }; let coroutine = forward.bind(py).call1((generator, sender))?; // create_task captures the caller's contextvars on the running loop. @@ -69,7 +61,6 @@ impl AsyncStream { })?; Ok(Self { rx, - finished, task: Some((task.unbind(), event_loop.unbind())), }) } @@ -80,17 +71,12 @@ impl Stream for AsyncStream { fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { let this = self.get_mut(); - Poll::Ready(match ready!(this.rx.poll_next_unpin(cx)) { - Some(item) => Some(item), - // Every sender is gone: the end after `finish`. Otherwise the forwarding task was - // cancelled or destroyed, so the body is incomplete. - None if this.task.take().is_some() && !this.finished.load(Ordering::Acquire) => { - Some(Err(PyRuntimeError::new_err( - "async body generator stopped before it finished", - ))) - } - None => None, - }) + let item = ready!(this.rx.poll_next_unpin(cx)); + // Forwarding is over once the body ends or fails, so a drop need not cancel it. + if !matches!(item, Some(Ok(_))) { + this.task = None; + } + Poll::Ready(item) } } @@ -128,8 +114,8 @@ impl Sender { let tx = match &*lock(&self.tx) { Some(tx) => match tx.try_reserve(&chunk) { Ok(permit) => return permit.send(chunk).into_bound_py_any(py), - Err(TrySendError::Full(())) => tx.clone(), - Err(TrySendError::Closed(())) => return false.into_bound_py_any(py), + Err(TryAcquireError::NoPermits) => tx.clone(), + Err(TryAcquireError::Closed) => return false.into_bound_py_any(py), }, None => return false.into_bound_py_any(py), }; @@ -152,12 +138,13 @@ impl Sender { } } - /// Mark the normal end of the body. The last chunk may still be queued: it is read - /// before the closed channel, so finishing never waits for room. + /// Mark the normal end of the body, which skips the budget, so finishing never waits. fn finish(&self) { - self.finished.store(true, Ordering::Release); // Python may retain the sender after completion, especially on PyPy. - self.close(); + let tx = lock(&self.tx).take(); + if let Some(tx) = tx { + tx.finish(); + } } /// Drop the channel end at once, so the body fails instead of waiting for more. diff --git a/src/client/body/stream/iter.rs b/src/client/body/stream/iter.rs index 75336ce3..212f0444 100644 --- a/src/client/body/stream/iter.rs +++ b/src/client/body/stream/iter.rs @@ -13,22 +13,19 @@ use pyo3::{ prelude::*, }; -use super::{Item, next_item, pump::Pumped}; +use super::{Item, next_item, pump::Pumped, queue}; use crate::runtime::Runtime; /// A request body from a Python iterator. /// /// Worker runtimes read the iterator on Tokio's blocking pool ([`Pumped`]): neither a worker /// nor a blocked caller runs it, so a slow `__next__` cannot delay the response, a timeout or -/// cancellation. A current-thread runtime reads it on its driving thread instead -/// ([`Inline`]), which waits for the request anyway. +/// cancellation. A current-thread runtime reads it on whichever thread drives it +/// ([`Inline`]), which is already blocked in a call on that runtime, though not necessarily +/// this request's. pub(super) enum SyncStream { - /// Not polled yet; the first poll picks where the iterator is read. `owner` passed the - /// iterator to the request. - Idle { - iter: Py, - owner: ThreadId, - }, + /// Not polled yet; the first poll picks where the iterator is read. + Idle(Inline), Inline(Inline), Pumped(Pumped), /// The inline iterator ended or raised. @@ -38,7 +35,7 @@ pub(super) enum SyncStream { /// An iterator read on the thread driving a current-thread runtime, between its IO polls. pub(super) struct Inline { iter: Py, - /// The thread that passed the iterator to the request. + /// The thread sending the request, whose call an interrupt from the iterator may end. owner: ThreadId, /// Whether a poll has returned to the connection since the last item. flushed: bool, @@ -49,9 +46,17 @@ pub(super) struct Inline { impl SyncStream { /// An iterator the calling thread passes to a request. pub(super) fn new(iter: Py) -> Self { - SyncStream::Idle { + SyncStream::Idle(Inline { iter, owner: thread::current().id(), + flushed: false, + }) + } + + /// Hand an unread iterator to a request the calling thread sends. + pub(super) fn claim(&mut self) { + if let SyncStream::Idle(inline) = self { + inline.owner = thread::current().id(); } } } @@ -61,17 +66,13 @@ impl Stream for SyncStream { fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { let this = self.get_mut(); - if let SyncStream::Idle { .. } = this - && let SyncStream::Idle { iter, owner } = mem::replace(this, SyncStream::Done) + if let SyncStream::Idle(_) = this + && let SyncStream::Idle(inline) = mem::replace(this, SyncStream::Done) { *this = if Runtime::driving() { - SyncStream::Inline(Inline { - iter, - owner, - flushed: false, - }) + SyncStream::Inline(inline) } else { - SyncStream::Pumped(Pumped::new(iter)) + SyncStream::Pumped(Pumped::new(inline.iter)) }; } match this { @@ -84,7 +85,7 @@ impl Stream for SyncStream { poll } SyncStream::Pumped(pumped) => pumped.poll_next_unpin(cx), - SyncStream::Idle { .. } | SyncStream::Done => Poll::Ready(None), + SyncStream::Idle(_) | SyncStream::Done => Poll::Ready(None), } } } @@ -111,7 +112,7 @@ impl Stream for Inline { fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { let this = self.get_mut(); ready!(this.poll_flush(cx)); - // Like a pump, stop reading once Python is unavailable. + // Like a pump, fail the body once Python is unavailable instead of ending it short. Python::try_attach(|py| { let item = next_item(py, &this.iter); // A KeyboardInterrupt or SystemExit reaches the caller as itself: the owner's call, @@ -125,7 +126,7 @@ impl Stream for Inline { } item }) - .flatten() + .unwrap_or_else(|| Some(Err(queue::unfinished()))) .into() } } diff --git a/src/client/body/stream/pump.rs b/src/client/body/stream/pump.rs index f3ea91bd..9e23acfb 100644 --- a/src/client/body/stream/pump.rs +++ b/src/client/body/stream/pump.rs @@ -9,7 +9,7 @@ use std::{ use futures_util::{Stream, StreamExt}; use pyo3::prelude::*; -use tokio::{sync::mpsc::error::TrySendError, time}; +use tokio::{sync::TryAcquireError, time}; use super::{ Item, PyBytesLike, lock, next_item, @@ -78,8 +78,8 @@ impl Pump { /// Queue chunks until the iterator ends or raises, the body is dropped, or the budget /// stays full past [`PARK_AFTER`](Self::PARK_AFTER). fn run(mut self, parked: &Mutex>) { - // Once Python is unavailable, stop reading without creating a PyErr that could - // require another attachment to format. + // Once Python is unavailable, stop without ending the body, which then fails as + // incomplete instead of sending a truncated one. Python::try_attach(|py| { while !self.tx.is_closed() { let chunk = match self.pending.take() { @@ -90,7 +90,10 @@ impl Pump { self.tx.fail(err); return; } - None => return, + None => { + self.tx.finish(); + return; + } }, }; let permit = match self.wait_for_room(py, &chunk) { @@ -101,8 +104,8 @@ impl Pump { let mut slot = lock(parked); match self.tx.try_reserve(&chunk) { Ok(permit) => permit, - Err(TrySendError::Closed(())) => return, - Err(TrySendError::Full(())) => { + Err(TryAcquireError::Closed) => return, + Err(TryAcquireError::NoPermits) => { self.pending = Some(chunk); *slot = Some(self); return; @@ -122,8 +125,8 @@ impl Pump { fn wait_for_room(&self, py: Python<'_>, chunk: &PyBytesLike) -> Option> { match self.tx.try_reserve(chunk) { Ok(permit) => Some(permit), - Err(TrySendError::Closed(())) => None, - Err(TrySendError::Full(())) => { + Err(TryAcquireError::Closed) => None, + Err(TryAcquireError::NoPermits) => { let wait = async { time::timeout(Self::PARK_AFTER, self.tx.reserve(chunk)).await }; Runtime::block_on_in_pool(py, wait).ok().flatten() } diff --git a/src/client/body/stream/queue.rs b/src/client/body/stream/queue.rs index be3b9840..8128e1cb 100644 --- a/src/client/body/stream/queue.rs +++ b/src/client/body/stream/queue.rs @@ -1,18 +1,16 @@ //! The queue between a chunk producer and its request body: each queued chunk holds a share //! of the upload budget until the body takes it, and dropping the body closes the budget. +//! The producer marks the end of the body, so a producer that stops early fails it. use std::{ pin::Pin, sync::Arc, - task::{Context, Poll}, + task::{Context, Poll, ready}, }; use futures_util::Stream; -use pyo3::PyErr; -use tokio::sync::{ - OwnedSemaphorePermit, Semaphore, TryAcquireError, - mpsc::{self, error::TrySendError}, -}; +use pyo3::{PyErr, exceptions::PyRuntimeError}; +use tokio::sync::{OwnedSemaphorePermit, Semaphore, TryAcquireError, mpsc}; use super::{Item, PyBytesLike}; @@ -22,7 +20,7 @@ const UPLOAD_BUDGET: usize = 256 * 1024; /// The least budget a queued chunk holds, so small chunks queue at most 64 items. const UPLOAD_CHARGE: usize = 4 * 1024; -/// The producer end. Errors skip the budget, so the error that ends a body never waits. +/// The producer end. The end and errors skip the budget, so finishing never waits. #[derive(Clone)] pub(super) struct Tx { chan: mpsc::UnboundedSender, @@ -39,11 +37,14 @@ pub(super) struct Permit<'a> { pub(super) struct Rx { chan: mpsc::UnboundedReceiver, budget: Arc, + /// Set once the body ended or failed; later polls end it again. + ended: bool, } -/// A queued item and the share it holds until the body takes it. +/// A queued chunk, the error that ends the body, or `None` for its normal end, with the +/// share a chunk holds until the body takes it. struct Queued { - item: Item, + item: Option, _share: Option, } @@ -55,7 +56,12 @@ pub(super) fn channel() -> (Tx, Rx) { chan: tx, budget: budget.clone(), }; - (tx, Rx { chan: rx, budget }) + let rx = Rx { + chan: rx, + budget, + ended: false, + }; + (tx, rx) } /// The budget a queued `chunk` holds; one larger than the budget takes all of it. @@ -63,19 +69,24 @@ fn charge(chunk: &PyBytesLike) -> u32 { u32::try_from(chunk.len().clamp(UPLOAD_CHARGE, UPLOAD_BUDGET)).unwrap_or(u32::MAX) } +/// The error for a body whose producer stopped before ending it. +pub(super) fn unfinished() -> PyErr { + PyRuntimeError::new_err("request body stopped before it finished") +} + // ===== impl Tx ===== impl Tx { - /// Reserve `chunk`'s share if the budget has room; `Closed` once the body is dropped. - pub(super) fn try_reserve(&self, chunk: &PyBytesLike) -> Result, TrySendError<()>> { - match self.budget.clone().try_acquire_many_owned(charge(chunk)) { - Ok(share) => Ok(Permit { + /// Reserve `chunk`'s share if the budget has room: `NoPermits` while it is full, + /// `Closed` once the body is dropped. + pub(super) fn try_reserve(&self, chunk: &PyBytesLike) -> Result, TryAcquireError> { + self.budget + .clone() + .try_acquire_many_owned(charge(chunk)) + .map(|share| Permit { chan: &self.chan, share, - }), - Err(TryAcquireError::NoPermits) => Err(TrySendError::Full(())), - Err(TryAcquireError::Closed) => Err(TrySendError::Closed(())), - } + }) } /// Wait for room for `chunk`'s share; `None` once the body is dropped. @@ -92,10 +103,18 @@ impl Tx { }) } + /// Queue the normal end of the body. + pub(super) fn finish(&self) { + let _ = self.chan.send(Queued { + item: None, + _share: None, + }); + } + /// Queue the error that ends the body. pub(super) fn fail(&self, err: PyErr) { let _ = self.chan.send(Queued { - item: Err(err), + item: Some(Err(err)), _share: None, }); } @@ -113,7 +132,7 @@ impl Permit<'_> { pub(super) fn send(self, chunk: PyBytesLike) -> bool { self.chan .send(Queued { - item: Ok(chunk), + item: Some(Ok(chunk)), _share: Some(self.share), }) .is_ok() @@ -134,11 +153,17 @@ impl Rx { impl Stream for Rx { type Item = Item; - fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - // Taking an item returns its share of the budget. - self.chan - .poll_recv(cx) - .map(|queued| queued.map(|queued| queued.item)) + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + if this.ended { + return Poll::Ready(None); + } + // Taking an item returns its share of the budget. Every producer gone without the + // end means one stopped early, so the body is incomplete. + let item = ready!(this.chan.poll_recv(cx)) + .map_or_else(|| Some(Err(unfinished())), |queued| queued.item); + this.ended = !matches!(item, Some(Ok(_))); + Poll::Ready(item) } } From 30486a3d602fb21f4e309a35bbcd26e0503b2715 Mon Sep 17 00:00:00 2001 From: 0x676e67 Date: Wed, 7 Oct 2026 13:49:58 +0000 Subject: [PATCH 17/25] refactor(runtime): name the inner runtime and read ready frames inside Streamer::next --- src/client/resp/stream.rs | 17 +++++------- src/runtime.rs | 57 ++++++++++++++++++++------------------- 2 files changed, 35 insertions(+), 39 deletions(-) diff --git a/src/client/resp/stream.rs b/src/client/resp/stream.rs index 13bc2576..43f8cd60 100644 --- a/src/client/resp/stream.rs +++ b/src/client/resp/stream.rs @@ -261,8 +261,11 @@ impl Streamer { Ok(()) } - /// Wait for the next frame. + /// Return a frame the body already holds, else wait for the next one. async fn next(&self, iteration: Iteration) -> PyResult { + if let Some(frame) = self.ready_frame(iteration) { + return frame; + } loop { // Readers are woken by `notify_waiters`, which reaches a `Notified` from its // creation, so it needs no `enable` and a frame sent after `try_next` still wakes it. @@ -283,12 +286,8 @@ impl Streamer { } fn __next__(&self, py: Python) -> PyResult { - // Refused before any frame, so a nested read never leaves a stream half consumed. - Runtime::refuse_nested()?; - // Frames already received are returned without releasing the GIL. - if let Some(frame) = self.ready_frame(Iteration::Sync) { - return frame; - } + // Refused before the first poll, which returns frames already received without + // releasing the GIL, so a nested read never leaves the stream half consumed. self.runtime.block_on_eager(py, self.next(Iteration::Sync)) } @@ -321,10 +320,6 @@ impl Streamer { let slf = slf.unbind(); coroutine::local(py, "Streamer.__anext__", async move { let this = slf.get(); - // A body already buffered is read without starting the read-ahead task. - if let Some(frame) = this.ready_frame(Iteration::Async) { - return frame; - } // Buffered frames complete without suspending; yield to the event loop // periodically so timeouts, cancellation and other tasks run. if this.reader.since_yield.fetch_add(1, Ordering::Relaxed) >= Self::YIELD_EVERY { diff --git a/src/runtime.rs b/src/runtime.rs index 484d0afc..2754b8ee 100644 --- a/src/runtime.rs +++ b/src/runtime.rs @@ -2,8 +2,9 @@ //! //! Requests, and async reads that must wait, run as tasks on the client's selected worker //! ([`Runtime::block_on_task`], [`Runtime::run_task`]); blocking reads are polled by the -//! caller ([`Runtime::block_on`], [`Runtime::block_on_eager`]). The `block_on*` calls release -//! the GIL while they wait, and drive a current-thread runtime meanwhile. +//! caller ([`Runtime::block_on`], [`Runtime::block_on_eager`]). These blocking calls release +//! the GIL while they wait; on a current-thread runtime they drive it and run the request +//! themselves instead of spawning it. //! //! Every other runtime operation goes through [`Runtime`] too; `clippy.toml` rejects direct //! Tokio handles, spawns and yields elsewhere. @@ -51,13 +52,13 @@ pub enum Scheduler { #[pyclass(frozen, skip_from_py_object)] pub struct Runtime { /// Shared owner; `None` only once `Drop` has taken it to shut the runtime down. - inner: Option>, + inner: Option>, /// The worker this copy runs on; per-worker clients keep their tasks there. handle: Handle, } /// The runtime behind [`Runtime`]. -enum Shared { +enum Inner { Workers(PingoraRuntime), /// Driven only by threads in a blocking call, through [`Driving::drive`]. CurrentThread(tokio::runtime::Runtime), @@ -68,9 +69,8 @@ struct Driving; thread_local! { static DRIVING: Cell = const { Cell::new(false) }; - /// A KeyboardInterrupt or other non-`Exception` error that this thread's upload - /// iterator raised while it drove the runtime, re-raised by its call instead of the - /// wrapped request error. + /// An error an upload iterator run on this thread chose to re-raise from the blocking + /// call this thread is driving, in place of that call's output. static INTERRUPT: RefCell> = const { RefCell::new(None) }; } @@ -79,7 +79,7 @@ thread_local! { impl Runtime { /// Whether blocking calls drive this runtime on their own thread. pub fn is_current_thread(&self) -> bool { - matches!(self.inner.as_deref(), Some(Shared::CurrentThread(_))) + matches!(self.inner.as_deref(), Some(Inner::CurrentThread(_))) } /// Share the runtime and select a worker for a new client. @@ -167,7 +167,7 @@ impl Runtime { fn wait(&self, future: impl Future>) -> PyResult { Self::refuse_nested()?; match self.inner.as_deref() { - Some(Shared::CurrentThread(runtime)) => Driving::drive(runtime, future), + Some(Inner::CurrentThread(runtime)) => Driving::drive(runtime, future), _ => self.handle.block_on(future), } } @@ -181,7 +181,7 @@ impl Runtime { { Self::refuse_nested()?; match self.inner.as_deref() { - Some(Shared::CurrentThread(runtime)) => Driving::drive(runtime, future), + Some(Inner::CurrentThread(runtime)) => Driving::drive(runtime, future), _ => self .handle .block_on(self.handle.spawn(future)) @@ -221,9 +221,10 @@ impl Runtime { current().spawn_blocking(f) } - /// Wait for `future` detached from Python on a blocking-pool thread, driven by the - /// runtime that owns the pool. Unlike [`block_on`](Self::block_on), it neither refuses - /// nested calls nor needs a `PyResult`; create timers inside `future`. + /// Wait for `future` detached from Python on a blocking-pool thread of a runtime whose + /// workers drive its IO and timers; a current-thread runtime's pool has none. Unlike + /// [`block_on`](Self::block_on), it neither refuses nested calls nor needs a `PyResult`; + /// create timers inside `future`. pub fn block_on_in_pool(py: Python<'_>, future: F) -> F::Output where F: Future + Send, @@ -262,8 +263,8 @@ impl Runtime { } } -impl From for Runtime { - fn from(runtime: Shared) -> Self { +impl From for Runtime { + fn from(runtime: Inner) -> Self { let handle = runtime.handle(); Self { inner: Some(Arc::new(runtime)), @@ -274,7 +275,7 @@ impl From for Runtime { impl From for Runtime { fn from(runtime: PingoraRuntime) -> Self { - Shared::Workers(runtime).into() + Inner::Workers(runtime).into() } } @@ -288,8 +289,8 @@ impl FromPyObject<'_, '_> for Runtime { #[pymethods] impl Runtime { - /// Create the runtime and start its workers. - /// Workers default to CPU parallelism and do not apply to CURRENT_THREAD. + /// Create the runtime and start its workers, if any. + /// Workers default to CPU parallelism and must be None with CURRENT_THREAD. /// Thread names default to the package name. /// Thread counts must be positive; thread_keep_alive is a nonnegative timedelta. #[new] @@ -335,7 +336,7 @@ impl Runtime { let runtime = builder .build() .map_err(|err| PyRuntimeError::new_err(err.to_string()))?; - Shared::CurrentThread(runtime).into() + Inner::CurrentThread(runtime).into() } _ => { let workers = workers.unwrap_or_else(parallelism); @@ -366,22 +367,22 @@ impl Drop for Runtime { // The final owner may be released on a worker or while holding the GIL. if let Some(runtime) = self.inner.take().and_then(Arc::into_inner) { match runtime { - Shared::Workers(PingoraRuntime::Steal { runtime, .. }) - | Shared::CurrentThread(runtime) => runtime.shutdown_background(), - Shared::Workers(PingoraRuntime::NoSteal(runtime)) => drop(runtime), + Inner::Workers(PingoraRuntime::Steal { runtime, .. }) + | Inner::CurrentThread(runtime) => runtime.shutdown_background(), + Inner::Workers(PingoraRuntime::NoSteal(runtime)) => drop(runtime), } } } } -// ===== impl Shared ===== +// ===== impl Inner ===== -impl Shared { +impl Inner { fn handle(&self) -> Handle { match self { // No-steal workers are lazy upstream; this starts them. - Shared::Workers(runtime) => runtime.get_handle().clone(), - Shared::CurrentThread(runtime) => runtime.handle().clone(), + Inner::Workers(runtime) => runtime.get_handle().clone(), + Inner::CurrentThread(runtime) => runtime.handle().clone(), } } } @@ -390,7 +391,7 @@ impl Shared { impl Driving { /// Drive `runtime` on this thread until `future` completes, marked as driving it; an - /// interrupt its upload iterator stashed meanwhile replaces the output. + /// interrupt an upload iterator stashed meanwhile replaces the output. fn drive( runtime: &tokio::runtime::Runtime, future: impl Future>, @@ -466,7 +467,7 @@ mod tests { #[test] fn current_thread_runs_on_its_callers() { let current_thread = || { - Runtime::from(Shared::CurrentThread( + Runtime::from(Inner::CurrentThread( Builder::new_current_thread().enable_all().build().unwrap(), )) }; From a3b8c7ff38658116577ed50c67ae9f602a7c82d0 Mon Sep 17 00:00:00 2001 From: 0x676e67 Date: Wed, 7 Oct 2026 13:50:03 +0000 Subject: [PATCH 18/25] docs: correct runtime, upload and event-loop read claims --- bench/README.md | 8 ++++---- docs/source/guide/advanced.md | 24 +++++++++++++----------- docs/source/guide/blocking.md | 17 +++++++++++------ python/wreq/runtime.py | 7 ++++--- 4 files changed, 32 insertions(+), 24 deletions(-) diff --git a/bench/README.md b/bench/README.md index f8096c67..b089ac46 100644 --- a/bench/README.md +++ b/bench/README.md @@ -15,7 +15,7 @@ are reported separately. - Async clients: wreq default runtime, wreq single worker, pyreqwest single-threaded and multithreaded runtimes, ry, httpx, aiohttp, niquests, and curl_cffi. -- Blocking variants: wreq MT/ST, ry, requests, httpx, niquests, curl_cffi, and pycurl. +- Blocking variants: wreq MT/ST/CT, ry, requests, httpx, niquests, curl_cffi, and pycurl. - Three rounds, each with one warmup and one timed batch of 300 requests. Case order is shuffled with a recorded seed. - Status, negotiated HTTP version, and total response length are checked for @@ -42,7 +42,7 @@ worker clients have separate connection pools. | 4 MiB | 256 KiB | Both Full and Stream uploads are tested for every payload and concurrency. -The complete matrix contains 1,680 supported cells across 16 client variants; +The complete matrix contains 1,904 supported cells across 18 client variants; each cell has three measured rounds. The Rust benchmark uses Criterion and 600 requests per iteration; this suite uses fixed batches of 300 requests. Custom payload sizes use upload chunks of at most 64 KiB. @@ -70,7 +70,7 @@ Cargo uses a custom target directory, pass the actual executable path to `--serv starts, it prints the number of cases and batches. It saves stdout/stderr logs beside the raw JSON and generates an English `.report.md` once validation passes. -The default suite has 1,680 supported cases, 5,040 timed batches and the same +The default suite has 1,904 supported cases, 5,712 timed batches and the same number of warm-up batches. Allow several hours for a full run. The setup commands above prepare the dependencies and native binaries; the runner won't install or build them for you. Publishing the results to the docs is optional. @@ -163,7 +163,7 @@ To update only the blocking results without rerunning the async clients: uv run --no-sync python bench/run.py --input bench/data/BLOCKING.json --publish-blocking ``` -This requires all eight blocking variants and the same complete workload and +This requires all nine blocking variants and the same complete workload and batch minimums. It freezes the candidate and existing `latest.json` for the docs build, then selects `bench/data/latest-blocking.json` only if the build succeeds. The async source and all historical files stay unchanged. The two publication diff --git a/docs/source/guide/advanced.md b/docs/source/guide/advanced.md index 36db6d05..21770153 100644 --- a/docs/source/guide/advanced.md +++ b/docs/source/guide/advanced.md @@ -103,8 +103,8 @@ if __name__ == "__main__": ### Custom runtimes Clients share a global multi-thread runtime when `runtime` is omitted or `None`. -It starts on first use. Construct a `Runtime` to -start a separate worker pool for an async or blocking client: +It starts on first use. Construct a `Runtime` to give an async or blocking +client its own runtime: ```python from datetime import timedelta @@ -133,18 +133,20 @@ client = Client(runtime=runtime) With `Scheduler.PER_WORKER`, each client is assigned one worker for its lifetime; requests, response reads, streams and WebSocket operations use that -worker. Async reads of HTTP/1 data that has already arrived, and is not -content-encoded, finish on the event loop thread instead: `bytes()` and `text()` -up to 64 KiB, `json()` up to 8 KiB, and the first frame of a stream whose length -is known; later frames are read on the worker. With multiple workers, newly -created clients select a worker randomly and keep that selection. This is not CPU -pinning. Sharing the same `Runtime` between clients is supported, and -`client.runtime` returns the shared runtime object. +worker. With multiple workers, newly created clients select a worker randomly and +keep that selection. This is not CPU pinning. Sharing the same `Runtime` between +clients is supported, and `client.runtime` returns the shared runtime object. + +With either worker scheduler, async reads of an HTTP/1 body of known length that +the client does not decompress finish on the event loop thread once the data has +arrived: `bytes()` and `text()` up to 64 KiB, `json()` up to 8 KiB, and a stream's +first frame of any size, plus the next one if it is already buffered. Other reads +and later stream frames run on the runtime. `workers=None` uses the available CPU parallelism, or 1 if it cannot be determined; `Scheduler.CURRENT_THREAD` requires it. -Custom runtimes start their threads during construction, before any client is -bound or request is sent. +Worker schedulers start their threads during construction, before any client is +bound or request is sent; `Scheduler.CURRENT_THREAD` starts none. `thread_name=None` uses the package name, `wreq-python`, as the thread name. diff --git a/docs/source/guide/blocking.md b/docs/source/guide/blocking.md index 46f13a28..b2af2f3f 100644 --- a/docs/source/guide/blocking.md +++ b/docs/source/guide/blocking.md @@ -80,7 +80,9 @@ with Client(runtime=runtime) as client: print(response.text()) ``` -Network work runs on the selected worker while the calling thread waits. +With a worker scheduler, network work runs on the selected worker while the +calling thread waits; with `Scheduler.CURRENT_THREAD` the calling thread runs it +(below). `client.runtime` is read-only. Closing the client does not shut down a shared runtime; it cancels pending requests and rejects new ones with `asyncio.CancelledError`. See [custom runtimes](advanced.md#custom-runtimes) for @@ -88,9 +90,11 @@ configuration and lifetime details. ### Current-thread runtime -A `Scheduler.CURRENT_THREAD` runtime has no worker threads: the thread in a -blocking call drives its IO, with no handoff to another thread. It is fastest -with a client and runtime per thread, like one curl handle per thread: +A `Scheduler.CURRENT_THREAD` runtime has no worker threads: a blocking call drives +the runtime on the calling thread instead of waiting on a worker. Redirect +callbacks, multipart file reads and `system_dns` lookups still use its blocking +pool. It is fastest with a client and runtime per thread, like one curl handle per +thread: ```python import threading @@ -107,8 +111,9 @@ def client() -> Client: return local.client ``` -- Nothing runs between calls. HTTP/2 pings, idle-connection cleanup, closing - cancelled connections and WebSocket keepalive progress only during a call. +- Nothing runs between calls. HTTP/2 keep-alive pings, idle-connection cleanup, + and closing cancelled connections or dropped WebSockets progress only during a + call. - Threads sharing the runtime take turns driving it, and the driving thread also runs the other threads' connections and upload iterators. A client shared by many threads is faster with `Scheduler.PER_WORKER`. diff --git a/python/wreq/runtime.py b/python/wreq/runtime.py index 88c072ad..564e4f91 100644 --- a/python/wreq/runtime.py +++ b/python/wreq/runtime.py @@ -20,14 +20,15 @@ class Scheduler(Enum): share it, and nothing runs between calls. Best with a client and runtime per thread; only blocking clients can use it. - Upload iterators run on the driving thread and must not wait for, or make, - other calls on this runtime. + Upload iterators run on whichever thread drives it: blocking wreq calls from + them raise RuntimeError, and one that waits for another call on this runtime + never returns. """ @final class Runtime: - """Shared Tokio runtime whose workers start when constructed. + """Shared Tokio runtime; worker schedulers start their threads when constructed. Clients, responses and active work keep it alive; the last owner releases it automatically. From c62c59308d046103328a91ce6d43d7aa7df2ab3b Mon Sep 17 00:00:00 2001 From: 0x676e67 Date: Wed, 7 Oct 2026 13:50:08 +0000 Subject: [PATCH 19/25] test: cover upload budget edges, multipart iterator owners and cross-thread interrupts --- tests/runtime_test.py | 46 +++++++++++++-- tests/shutdown_test.py | 32 +++++++---- tests/sync_upload_test.py | 118 ++++++++++++++++++-------------------- tests/upload_test.py | 18 ------ 4 files changed, 118 insertions(+), 96 deletions(-) diff --git a/tests/runtime_test.py b/tests/runtime_test.py index afbde27d..1aa5267f 100644 --- a/tests/runtime_test.py +++ b/tests/runtime_test.py @@ -240,12 +240,12 @@ def fetch(path): with client.get(url + path) as response: return response.text() - def interrupted(): + def interrupted(error): yield b"x" - raise Interrupt + raise error - def upload_interrupted(): - client.post(url + "/interrupted", body=interrupted()) + def upload_interrupted(error=Interrupt): + client.post(url + "/interrupted", body=interrupted(error)) def run(): caller, readers = threading.get_ident(), [] @@ -304,6 +304,44 @@ def chunks(): release.set() assert await held == "/hold" + # Ctrl+C lands on whichever thread runs Python, so it ends the driving thread's + # own call, while the uploading call gets the wrapped error. + def fetch_interrupted(): + # Caught here: a KeyboardInterrupt reaching an asyncio task stops the loop. + with pytest.raises(KeyboardInterrupt): + fetch("/hold") + + release.clear() + held = asyncio.create_task(run_blocking(fetch_interrupted)) + while await asyncio.wait_for(arrived.get(), 5) != "/hold": + pass + with pytest.raises(wreq.exceptions.RequestError, match="KeyboardInterrupt"): + await run_blocking(upload_interrupted, KeyboardInterrupt) + assert not held.done() + release.set() + await held + + # A multipart iterator belongs to the thread sending it, not the one that built + # its part, so the held thread that built it keeps its own result. + built = [] + + def build_and_fetch(): + built.append(wreq.Part(name="file", value=interrupted(Interrupt))) + return fetch("/hold") + + def upload_part(): + client.post(url + "/interrupted", multipart=wreq.Multipart(built[0])) + + release.clear() + held = asyncio.create_task(run_blocking(build_and_fetch)) + while await asyncio.wait_for(arrived.get(), 5) != "/hold": + pass + with pytest.raises(wreq.exceptions.RequestError, match="Interrupt"): + await run_blocking(upload_part) + assert not held.done() + release.set() + assert await held == "/hold" + # Closing the client from another thread ends a call driving the runtime. release.clear() held = asyncio.create_task(run_blocking(fetch, "/hold")) diff --git a/tests/shutdown_test.py b/tests/shutdown_test.py index 03b2e61d..01e6cf68 100644 --- a/tests/shutdown_test.py +++ b/tests/shutdown_test.py @@ -5,7 +5,7 @@ import pytest # Keep a pending request alive into CPython module teardown, then fail its I/O. -# Upload mode also wakes the Python producer waiting for channel capacity. +# Upload mode also wakes the Python producer waiting for upload budget. SCRIPT = """ import asyncio import os @@ -64,7 +64,7 @@ async def chunks(): assert any( getattr(getattr(t.get_coro(), "cr_await", None), "__name__", None) == "send" for t in asyncio.all_tasks(loop) - ), "upload producer must be waiting for channel capacity" + ), "upload producer must be waiting for upload budget" hold = HoldTeardown(peer, task) del peer, task raise RuntimeError("uncaught error while a request is in flight") @@ -97,7 +97,7 @@ def test_unconsumed_upload_cleanup(): import gc import wreq -async def main(retain): +async def main(retain, first, second, expected): produced = [] full = asyncio.Event() closed = asyncio.Event() @@ -106,9 +106,10 @@ async def chunks(): try: for index in range(100): produced.append(index) - if index == 1: + # The last expected chunk is the first that waits for room. + if index == len(expected) - 1: full.set() - yield b"chunk" + yield first if index == 0 else second if index == 1 else b"chunk" finally: closed.set() print("generator closed", flush=True) @@ -118,8 +119,7 @@ async def chunks(): try: await asyncio.wait_for(full.wait(), 5) await asyncio.sleep(0.05) - # 64 small chunks fill the upload budget at its minimum charge; the next waits. - assert produced == list(range(65)) + assert produced == expected, produced if retain: return part finally: @@ -130,9 +130,19 @@ async def chunks(): gc.collect() await asyncio.wait_for(closed.wait(), 5) -for retain in (False, True): - part = asyncio.run(main(retain)) - print("runner closed", flush=True) +big = b"x" * (1 << 20) +cases = [ + # 64 small chunks fill the budget at its minimum charge, and the next waits. + (b"chunk", b"chunk", list(range(65))), + # A chunk above the budget queues alone, and the next waits. + (big, b"chunk", [0, 1]), + # A chunk above the budget waits for the whole budget. + (b"chunk", big, [0, 1]), +] +for first, second, expected in cases: + for retain in (False, True): + part = asyncio.run(main(retain, first, second, expected)) + print("runner closed", flush=True) """ proc = subprocess.run( [sys.executable, "-c", script], @@ -141,4 +151,4 @@ async def chunks(): timeout=15, ) assert proc.returncode == 0, proc.stderr - assert proc.stdout.splitlines() == ["generator closed", "runner closed"] * 2 + assert proc.stdout.splitlines() == ["generator closed", "runner closed"] * 6 diff --git a/tests/sync_upload_test.py b/tests/sync_upload_test.py index 5a1c5c95..da288534 100644 --- a/tests/sync_upload_test.py +++ b/tests/sync_upload_test.py @@ -68,24 +68,38 @@ async def accept(reader, writer): assert not errors, errors -def upload_client(blocking, current_thread): - """A client whose iterators run on the blocking pool, or inline on a current-thread - runtime's driving thread.""" - if current_thread: +def upload_client(kind): + """A client whose iterators run on the async or blocking client's blocking pool, or + inline on a current-thread runtime's driving thread.""" + if kind == "current_thread": return wreq.blocking.Client( proxies=[], runtime=Runtime(scheduler=Scheduler.CURRENT_THREAD) ) - return (wreq.blocking.Client if blocking else wreq.Client)(proxies=[]) + return (wreq.blocking.Client if kind == "blocking" else wreq.Client)(proxies=[]) -UPLOAD_CLIENTS = [(False, False), (True, False), (True, True)] +UPLOAD_KINDS = ["async", "blocking", "current_thread"] + + +async def post_bytes(client, url, **kwargs): + """Send a request with either client kind and return its body.""" + if isinstance(client, wreq.blocking.Client): + + def send(): + with client.post(url, **kwargs) as response: + return response.bytes() + + return await asyncio.to_thread(send) + response = await client.post(url, **kwargs) + async with response: + return await response.bytes() @pytest.mark.asyncio -@pytest.mark.parametrize(("blocking", "current_thread"), UPLOAD_CLIENTS) +@pytest.mark.parametrize("kind", UPLOAD_KINDS) @pytest.mark.parametrize("multipart", [False, True]) -async def test_sync_upload_iterator_errors(blocking, current_thread, multipart): - client = upload_client(blocking, current_thread) +async def test_sync_upload_iterator_errors(kind, multipart): + client = upload_client(kind) try: async with upload_server() as (url, bodies): for failure in (False, True): @@ -110,24 +124,16 @@ def chunks(): else {"body": chunks()} ) - def send_blocking(): - with client.post(url, **kwargs) as response: - return response.bytes() - - async def send(): - if blocking: - return await asyncio.to_thread(send_blocking) - response = await client.post(url, **kwargs) - async with response: - return await response.bytes() - if failure: with pytest.raises(wreq.exceptions.RequestError) as caught: - await asyncio.wait_for(send(), 5) + await asyncio.wait_for(post_bytes(client, url, **kwargs), 5) assert "ValueError" in str(caught.value) assert "upload exploded" in str(caught.value) else: - assert await asyncio.wait_for(send(), 5) == b"ok" + assert ( + await asyncio.wait_for(post_bytes(client, url, **kwargs), 5) + == b"ok" + ) body = await asyncio.wait_for(bodies.get(), 5) assert ( (b"hello world" in body) @@ -140,19 +146,26 @@ async def send(): @pytest.mark.asyncio -@pytest.mark.parametrize("blocking", [False, True]) -async def test_sync_upload_reads_ahead_within_the_budget(blocking): - # A huge first chunk keeps the connection's write buffer full while the peer waits, so - # only the budget bounds the read-ahead: 64 small chunks queue at the minimum charge - # and the next parks with the pump. Once the peer reads, the parked chunk goes first. +@pytest.mark.parametrize( + ("source", "kind"), + [("async_generator", "async"), ("iterator", "async"), ("iterator", "blocking")], +) +async def test_upload_reads_one_chunk_past_the_budget(source, kind): + # A huge first chunk fills the connection's write buffer while the peer waits. 64 small + # chunks then fill the budget, and the producer waits holding one more chunk it already + # read. Once the peer reads, the waiting chunk goes out first. release, bodies, calls = threading.Event(), asyncio.Queue(), [] parts = [b"\0" * (1 << 24)] + [bytes((index,)) * 4096 for index in range(100)] - def chunks(): + def iterator(): for part in parts: calls.append(None) yield part + async def async_generator(): + for part in iterator(): + yield part + async def accept(reader, writer): try: await asyncio.wait_for(reader.readuntil(b"\r\n\r\n"), 5) @@ -167,29 +180,16 @@ async def accept(reader, writer): server = await asyncio.start_server(accept, "127.0.0.1", 0) url = f"http://127.0.0.1:{server.sockets[0].getsockname()[1]}/" - client = upload_client(blocking, False) + client = upload_client(kind) + body = async_generator() if source == "async_generator" else iterator() try: - if blocking: - - def send_blocking(): - with client.post(url, body=chunks()) as response: - return response.bytes() - - task = asyncio.create_task(asyncio.to_thread(send_blocking)) - else: - - async def send(): - response = await client.post(url, body=chunks()) - async with response: - return await response.bytes() - - task = asyncio.create_task(send()) - seen, stable = -1, 0 + task = asyncio.create_task(post_bytes(client, url, body=body)) for _ in range(250): - await asyncio.sleep(0.02) - seen, stable = len(calls), stable + 1 if len(calls) == seen else 0 - if stable == 5: + if len(calls) >= 66: break + await asyncio.sleep(0.02) + # A producer reading past the budget would continue at once; give it time to show. + await asyncio.sleep(0.1) assert len(calls) == 66 release.set() assert await asyncio.wait_for(task, 10) == b"ok" @@ -228,8 +228,8 @@ def send(): @pytest.mark.asyncio -@pytest.mark.parametrize(("blocking", "current_thread"), UPLOAD_CLIENTS) -async def test_sync_upload_sends_each_chunk_as_yielded(blocking, current_thread): +@pytest.mark.parametrize("kind", UPLOAD_KINDS) +async def test_sync_upload_sends_each_chunk_as_yielded(kind): # A chunk goes out once yielded, before the iterator produces the next one. first_seen = threading.Event() bodies = asyncio.Queue() @@ -259,19 +259,11 @@ def chunks(): server = await asyncio.start_server(accept, "127.0.0.1", 0) url = f"http://127.0.0.1:{server.sockets[0].getsockname()[1]}/" - client = upload_client(blocking, current_thread) + client = upload_client(kind) try: - if blocking: - - def send(): - with client.post(url, body=chunks()) as response: - return response.bytes() - - assert await asyncio.wait_for(asyncio.to_thread(send), 10) == b"ok" - else: - response = await asyncio.wait_for(client.post(url, body=chunks()), 10) - async with response: - assert await response.bytes() == b"ok" + assert ( + await asyncio.wait_for(post_bytes(client, url, body=chunks()), 10) == b"ok" + ) assert await asyncio.wait_for(bodies.get(), 5) == b"firstsecond" finally: client.close() @@ -356,7 +348,7 @@ def chunks(): server = await asyncio.start_server(accept, "127.0.0.1", 0) url = f"http://127.0.0.1:{server.sockets[0].getsockname()[1]}/" - client = upload_client(True, True) + client = upload_client("current_thread") try: def send(): diff --git a/tests/upload_test.py b/tests/upload_test.py index 14ccb303..8bed63c8 100644 --- a/tests/upload_test.py +++ b/tests/upload_test.py @@ -95,24 +95,6 @@ async def chunks(): await (await asyncio.wait_for(task, 5)).close() -@pytest.mark.asyncio -async def test_async_upload_beyond_the_budget_keeps_order(): - # 1 MiB in small chunks queues within the budget, then waits for room, in order. - parts = [bytes((index % 256,)) * 4096 for index in range(256)] - - async def chunks(): - for part in parts: - yield part - - async with local_server() as (url, connections), wreq.Client(proxies=[]) as client: - task = asyncio.create_task(client.post(url, body=chunks())) - reader, writer = await asyncio.wait_for(connections.get(), 5) - assert await asyncio.wait_for(read_chunked(reader), 10) == b"".join(parts) - writer.write(b"HTTP/1.1 200 OK\r\nContent-Length: 0\r\n\r\n") - await writer.drain() - await (await asyncio.wait_for(task, 5)).close() - - @pytest.mark.asyncio @pytest.mark.parametrize("failure", ["exception", "type", "cancelled"]) async def test_upload_errors(failure): From 68f9fd2b0c643de5f904093c0995bb539a9e76ad Mon Sep 17 00:00:00 2001 From: 0x676e67 Date: Wed, 7 Oct 2026 14:53:31 +0000 Subject: [PATCH 20/25] fix(runtime): collect idle IO before a current-thread request picks a connection --- src/runtime.rs | 9 ++++++++- tests/runtime_test.py | 34 ++++++++++++++++++++++++++++++++++ 2 files changed, 42 insertions(+), 1 deletion(-) diff --git a/src/runtime.rs b/src/runtime.rs index 2754b8ee..96d95994 100644 --- a/src/runtime.rs +++ b/src/runtime.rs @@ -181,7 +181,14 @@ impl Runtime { { Self::refuse_nested()?; match self.inner.as_deref() { - Some(Inner::CurrentThread(runtime)) => Driving::drive(runtime, future), + Some(Inner::CurrentThread(runtime)) => Driving::drive(runtime, async { + // Nothing ran since the last call. The first yield lets the driver collect IO + // that arrived meanwhile, such as a server closing an idle connection, and the + // second lets its task act on it before the request checks out a connection. + task::yield_now().await; + task::yield_now().await; + future.await + }), _ => self .handle .block_on(self.handle.spawn(future)) diff --git a/tests/runtime_test.py b/tests/runtime_test.py index 1aa5267f..5a6bcc7b 100644 --- a/tests/runtime_test.py +++ b/tests/runtime_test.py @@ -2,6 +2,7 @@ import base64 import hashlib import threading +import time from contextlib import asynccontextmanager from datetime import timedelta @@ -354,6 +355,39 @@ def upload_part(): client.close() +@pytest.mark.asyncio +async def test_current_thread_skips_connections_closed_while_idle(): + # Nothing drives a current-thread runtime between calls, so a request must first see + # that the server closed the pooled connection meanwhile instead of sending on it. + async def accept(reader, writer): + try: + await reader.readuntil(b"\r\n\r\n") + writer.write(b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\n\r\nok") + await writer.drain() + # A short keep-alive closes the idle connection. + await asyncio.sleep(0.02) + finally: + writer.close() + + def fetch_after_idle(): + for _ in range(3): + with client.get(url) as response: + assert response.text() == "ok" + time.sleep(0.2) + + server = await asyncio.start_server(accept, "127.0.0.1", 0) + url = f"http://127.0.0.1:{server.sockets[0].getsockname()[1]}/" + client = wreq.blocking.Client( + runtime=Runtime(scheduler=Scheduler.CURRENT_THREAD), proxies=[] + ) + try: + await run_blocking(fetch_after_idle) + finally: + client.close() + server.close() + await server.wait_closed() + + @pytest.mark.asyncio @pytest.mark.parametrize( ("blocking", "operation"), From 14fb2f9f2fb13a9ee000788686538920aa786a9d Mon Sep 17 00:00:00 2001 From: 0x676e67 Date: Wed, 7 Oct 2026 14:53:35 +0000 Subject: [PATCH 21/25] docs: note the work_steal migration and name the pending upload send after try_send --- docs/source/guide/advanced.md | 4 +++- src/client/body/stream/async_gen.rs | 5 ++--- src/client/resp/stream.rs | 3 ++- tests/shutdown_test.py | 3 ++- 4 files changed, 9 insertions(+), 6 deletions(-) diff --git a/docs/source/guide/advanced.md b/docs/source/guide/advanced.md index 21770153..838071fa 100644 --- a/docs/source/guide/advanced.md +++ b/docs/source/guide/advanced.md @@ -122,7 +122,9 @@ runtime = Runtime( client = Client(runtime=runtime) ``` -`scheduler` selects how the runtime runs client work: +`scheduler` selects how the runtime runs client work. It replaces `work_steal`: +`work_steal=False` is `Scheduler.PER_WORKER`, and the old default is +`Scheduler.WORK_STEALING`. - `Scheduler.WORK_STEALING` (the default) uses one multi-thread pool whose workers steal work from each other. diff --git a/src/client/body/stream/async_gen.rs b/src/client/body/stream/async_gen.rs index 7bd7b656..0c122e47 100644 --- a/src/client/body/stream/async_gen.rs +++ b/src/client/body/stream/async_gen.rs @@ -120,9 +120,8 @@ impl Sender { None => return false.into_bound_py_any(py), }; // The coroutine keeps the extracted chunk, so it is never copied twice. Budget - // readiness is runtime-independent, so it waits on the Python loop. Its name is the - // pending send that shutdown inspects. - coroutine::local(py, "Sender.send", async move { + // readiness is runtime-independent, so it waits on the Python loop. + coroutine::local(py, "Sender.try_send", async move { let permit = tx.reserve(&chunk).await; Ok(permit.is_some_and(|permit| permit.send(chunk))) }) diff --git a/src/client/resp/stream.rs b/src/client/resp/stream.rs index 43f8cd60..a0d1e75b 100644 --- a/src/client/resp/stream.rs +++ b/src/client/resp/stream.rs @@ -40,7 +40,8 @@ pub enum Frame { /// A response stream yielding read-only memoryviews and any trailing headers. /// -/// Sync and async iteration share [`Streamer::next`]; only the wait differs. +/// Sync and async iteration share [`Streamer::next`]; [`Iteration`] picks the ready frames +/// the caller reads and the error that ends iteration. #[pyclass(subclass, frozen, skip_from_py_object)] pub struct Streamer { reader: Arc, diff --git a/tests/shutdown_test.py b/tests/shutdown_test.py index 01e6cf68..54efdc1a 100644 --- a/tests/shutdown_test.py +++ b/tests/shutdown_test.py @@ -62,7 +62,8 @@ async def chunks(): assert not task.done(), "request must remain pending" if upload: assert any( - getattr(getattr(t.get_coro(), "cr_await", None), "__name__", None) == "send" + getattr(getattr(t.get_coro(), "cr_await", None), "__name__", None) + == "try_send" for t in asyncio.all_tasks(loop) ), "upload producer must be waiting for upload budget" hold = HoldTeardown(peer, task) From 98b70cf8f1dad1f50a9e2af6e9cd8fc8ad1f32bc Mon Sep 17 00:00:00 2001 From: 0x676e67 Date: Wed, 7 Oct 2026 15:34:52 +0000 Subject: [PATCH 22/25] test(cancellation): throw in a helper and collect until PyPy releases the exception --- tests/cancellation_test.py | 21 ++++++++++++++------- 1 file changed, 14 insertions(+), 7 deletions(-) diff --git a/tests/cancellation_test.py b/tests/cancellation_test.py index 63fca185..6ff06d51 100644 --- a/tests/cancellation_test.py +++ b/tests/cancellation_test.py @@ -150,6 +150,16 @@ async def accept(reader, writer): await server.wait_closed() +def throw_cancellation(coroutine): + """Throw a cancellation into `coroutine` and return a weak reference to it, so no + caller frame still holds the exception.""" + error = Cancellation("cancelled after Rust completion") + with pytest.raises(asyncio.CancelledError) as caught: + coroutine.throw(error) + assert caught.value is error + return weakref.ref(error) + + @pytest.mark.asyncio @pytest.mark.parametrize("operation", ["request", "request_error", "stream"]) async def test_cancellation_after_rust_completion(operation): @@ -181,15 +191,12 @@ async def test_cancellation_after_rust_completion(operation): done, _ = await asyncio.wait({waiter}, timeout=5) assert waiter in done, "Rust operation did not finish" - error = Cancellation("cancelled after Rust completion") - error_ref = weakref.ref(error) - with pytest.raises(asyncio.CancelledError) as caught: - coroutine.throw(error) - assert caught.value is error + error_ref = throw_cancellation(coroutine) # Keeping the finished coroutine alive must not retain its exception. - del caught, error # PyPy may need several passes to release it. - for _ in range(3): + for _ in range(10): + if error_ref() is None: + break gc.collect() assert error_ref() is None finally: From 3ac0aa1a8538ccb8e19a49c4f6d760e0dda211e2 Mon Sep 17 00:00:00 2001 From: 0x676e67 Date: Wed, 7 Oct 2026 16:01:45 +0000 Subject: [PATCH 23/25] refactor(coroutine): accept the legacy throw form in Rust and drop the PyPy shim --- python/wreq/__init__.py | 15 ------ python/wreq/_compat.py | 49 ------------------ src/coroutine/awaitable.rs | 103 ++++++++++++++++++++++++++++++++++++- tests/cancellation_test.py | 9 ---- 4 files changed, 101 insertions(+), 75 deletions(-) delete mode 100644 python/wreq/_compat.py diff --git a/python/wreq/__init__.py b/python/wreq/__init__.py index a3955b41..9dc0c8d4 100644 --- a/python/wreq/__init__.py +++ b/python/wreq/__init__.py @@ -14,21 +14,6 @@ from .proxy import * from .runtime import * -import sys as _sys - -if _sys.implementation.name == "pypy": - from ._compat import _install - - # Creating and closing an unpolled coroutine does not start a request or runtime. - _coroutine = get("") - try: - _install(type(_coroutine)) - finally: - _coroutine.close() - del _coroutine, _install - -del _sys - __all__ = ( header.__all__ + cookie.__all__ diff --git a/python/wreq/_compat.py b/python/wreq/_compat.py deleted file mode 100644 index d1cc1f8d..00000000 --- a/python/wreq/_compat.py +++ /dev/null @@ -1,49 +0,0 @@ -"""Compatibility for PyPy's legacy exception delegation to wreq coroutines.""" - -from types import TracebackType - - -def _install(coroutine_type): - original = vars(coroutine_type)["throw"] - if getattr(original, "_wreq_pypy_throw_compat", False): - return - - # Coroutine.throw takes one exception, as asyncio passes; PyPy delegates - # throws from awaiting generators as the legacy (type, value, traceback). - def throw(self, exc, value=None, traceback=None): - if traceback is not None and type(traceback) is not TracebackType: - raise TypeError("throw() third argument must be a traceback object") - if issubclass(type(exc), BaseException): - if value is not None: - raise TypeError("instance exception may not have a separate value") - instance = exc - if traceback is not None: - BaseException.with_traceback(instance, traceback) - elif issubclass(type(exc), type) and issubclass(exc, BaseException): - try: - if issubclass(type(value), BaseException) and issubclass( - type(value), exc - ): - instance = value - elif value is None: - instance = exc() - elif issubclass(type(value), tuple): - instance = exc(*value) - else: - instance = exc(value) - except BaseException as error: - instance = error - own_traceback = BaseException.__traceback__.__get__(error).tb_next - BaseException.with_traceback(instance, own_traceback or traceback) - else: - if not issubclass(type(instance), BaseException): - instance = TypeError( - "exception constructor did not return an exception" - ) - BaseException.with_traceback(instance, traceback) - else: - raise TypeError("exceptions must derive from BaseException") - return original(self, instance) - - throw._wreq_pypy_throw_compat = True - coroutine_type.throw = throw diff --git a/src/coroutine/awaitable.rs b/src/coroutine/awaitable.rs index c0e3ee43..b09c6d65 100644 --- a/src/coroutine/awaitable.rs +++ b/src/coroutine/awaitable.rs @@ -10,9 +10,10 @@ use std::{ use futures_util::{FutureExt, future::BoxFuture}; use pyo3::{ PyTraverseError, PyVisit, - exceptions::{PyRuntimeError, PyStopIteration}, + exceptions::{PyBaseException, PyRuntimeError, PyStopIteration, PyTypeError}, intern, prelude::*, + types::{PyTraceback, PyTuple, PyType}, }; use super::{Port, scope::Scope}; @@ -152,7 +153,17 @@ impl Coroutine { self.step(py, Some(value)).and_then(Step::into_result) } - fn throw(&mut self, py: Python<'_>, exc: Bound<'_, PyAny>) -> PyResult> { + /// Raise `exc` at the await point. Like a generator's, it also takes the legacy + /// `(type, value, traceback)` form, in which PyPy delegates throws. + #[pyo3(signature = (exc, value = None, traceback = None))] + fn throw( + &mut self, + py: Python<'_>, + exc: Bound<'_, PyAny>, + value: Option>, + traceback: Option>, + ) -> PyResult> { + let exc = thrown(exc, value, traceback)?; self.throw_step(py, exc).and_then(Step::into_result) } @@ -204,6 +215,94 @@ impl Step { } } +/// The exception `throw(exc, value, traceback)` raises, built as a generator's `throw` +/// builds it: a constructor that fails raises its own error instead. +fn thrown<'py>( + exc: Bound<'py, PyAny>, + value: Option>, + traceback: Option>, +) -> PyResult> { + let py = exc.py(); + if traceback + .as_ref() + .is_some_and(|tb| !tb.is_exact_instance_of::()) + { + return Err(PyTypeError::new_err( + "throw() third argument must be a traceback object", + )); + } + if exc.is_instance_of::() { + if value.is_some() { + return Err(PyTypeError::new_err( + "instance exception may not have a separate value", + )); + } + if traceback.is_some() { + set_traceback(&exc, traceback.as_ref())?; + } + return Ok(exc); + } + let ty = match exc.cast_into::() { + Ok(ty) if ty.is_subclass_of::()? => ty, + _ => { + return Err(PyTypeError::new_err( + "exceptions must derive from BaseException", + )); + } + }; + let (instance, traceback) = match instantiate(&ty, value) { + Ok(instance) if instance.is_instance_of::() => (instance, traceback), + Ok(_) => ( + PyTypeError::new_err("exception constructor did not return an exception") + .into_value(py) + .into_bound(py) + .into_any(), + traceback, + ), + Err(err) => { + let own = err.traceback(py).map(Bound::into_any); + ( + err.into_value(py).into_bound(py).into_any(), + own.or(traceback), + ) + } + }; + set_traceback(&instance, traceback.as_ref())?; + Ok(instance) +} + +/// `ty`'s exception for `value`: `value` itself if it already is one, else `ty` called with +/// no arguments, the tuple's items or `value`. +fn instantiate<'py>( + ty: &Bound<'py, PyType>, + value: Option>, +) -> PyResult> { + match value { + Some(value) + if value.is_instance_of::() && value.get_type().is_subclass(ty)? => + { + Ok(value) + } + None => ty.call0(), + Some(value) => match value.cast_into::() { + Ok(args) => ty.call1(args), + Err(err) => ty.call1((err.into_inner(),)), + }, + } +} + +/// Set or clear `exception`'s traceback through `BaseException`, never an override. +fn set_traceback<'py>( + exception: &Bound<'py, PyAny>, + traceback: Option<&Bound<'py, PyAny>>, +) -> PyResult<()> { + let py = exception.py(); + py.get_type::() + .getattr(intern!(py, "with_traceback"))? + .call1((exception, traceback)) + .map(drop) +} + /// Fail if `waiter`, the future a task waits on for this coroutine, is still pending. pub(super) fn ensure_done(waiter: &Bound<'_, PyAny>) -> PyResult<()> { if waiter diff --git a/tests/cancellation_test.py b/tests/cancellation_test.py index 6ff06d51..d3c1b5cc 100644 --- a/tests/cancellation_test.py +++ b/tests/cancellation_test.py @@ -1,6 +1,5 @@ import asyncio import gc -import sys import weakref from contextlib import asynccontextmanager @@ -13,12 +12,7 @@ class Cancellation(asyncio.CancelledError): pass -@pytest.mark.skipif( - sys.implementation.name != "pypy", reason="PyPy legacy throw protocol" -) def test_legacy_coroutine_throw(): - from wreq._compat import _install - try: raise RuntimeError("traceback origin") except RuntimeError as error: @@ -105,9 +99,6 @@ def __new__(cls): coroutine = wreq.get("") try: - installed = vars(type(coroutine))["throw"] - _install(type(coroutine)) - assert vars(type(coroutine))["throw"] is installed for args in ( (object(),), (ValueError, None, object()), From 857123362649b520bb3880c124808919d5512bf4 Mon Sep 17 00:00:00 2001 From: 0x676e67 Date: Wed, 7 Oct 2026 16:59:29 +0000 Subject: [PATCH 24/25] fix(coroutine): build thrown exceptions before borrowing and type them as BaseException --- src/coroutine/awaitable.rs | 111 +++++++++++++++++++------------------ src/coroutine/scope.rs | 14 +++-- tests/cancellation_test.py | 43 +++++++++++++- 3 files changed, 108 insertions(+), 60 deletions(-) diff --git a/src/coroutine/awaitable.rs b/src/coroutine/awaitable.rs index b09c6d65..463ac4e7 100644 --- a/src/coroutine/awaitable.rs +++ b/src/coroutine/awaitable.rs @@ -21,8 +21,8 @@ use super::{Port, scope::Scope}; /// An awaitable driving a Rust future on the asyncio event loop thread. /// /// Each suspension hands the task a real asyncio future, so task cancellation and -/// the C task fast path work unchanged. A throw or close drops the Rust future, -/// which aborts any spawned work. PyO3's borrow flag rejects reentrant polls. +/// the C task fast path work unchanged. A throw it accepts, or a close, drops the Rust +/// future, which aborts any spawned work. PyO3's borrow flag rejects reentrant polls. /// /// A coroutine built with [`managed`](super::managed) also works as `async with`, as /// `async with await` would; the `scope` module holds that state. @@ -153,18 +153,20 @@ impl Coroutine { self.step(py, Some(value)).and_then(Step::into_result) } - /// Raise `exc` at the await point. Like a generator's, it also takes the legacy - /// `(type, value, traceback)` form, in which PyPy delegates throws. + // Like a generator's, this also takes the legacy `(type, value, traceback)` form, in + // which PyPy delegates throws. The exception is built before borrowing the coroutine, + // as its constructor may use it. #[pyo3(signature = (exc, value = None, traceback = None))] fn throw( - &mut self, - py: Python<'_>, + slf: &Bound<'_, Self>, exc: Bound<'_, PyAny>, value: Option>, traceback: Option>, ) -> PyResult> { let exc = thrown(exc, value, traceback)?; - self.throw_step(py, exc).and_then(Step::into_result) + slf.try_borrow_mut()? + .throw_step(slf.py(), exc) + .and_then(Step::into_result) } fn close(&mut self, py: Python<'_>) -> PyResult<()> { @@ -216,13 +218,12 @@ impl Step { } /// The exception `throw(exc, value, traceback)` raises, built as a generator's `throw` -/// builds it: a constructor that fails raises its own error instead. +/// builds it. fn thrown<'py>( exc: Bound<'py, PyAny>, value: Option>, traceback: Option>, -) -> PyResult> { - let py = exc.py(); +) -> PyResult> { if traceback .as_ref() .is_some_and(|tb| !tb.is_exact_instance_of::()) @@ -231,75 +232,77 @@ fn thrown<'py>( "throw() third argument must be a traceback object", )); } - if exc.is_instance_of::() { - if value.is_some() { + let ty = match exc.cast_into::() { + Ok(_) if value.is_some() => { return Err(PyTypeError::new_err( "instance exception may not have a separate value", )); } - if traceback.is_some() { - set_traceback(&exc, traceback.as_ref())?; - } - return Ok(exc); - } - let ty = match exc.cast_into::() { - Ok(ty) if ty.is_subclass_of::()? => ty, - _ => { - return Err(PyTypeError::new_err( - "exceptions must derive from BaseException", - )); + Ok(exc) => { + if traceback.is_some() { + set_traceback(&exc, traceback.as_ref())?; + } + return Ok(exc); } + Err(err) => match err.into_inner().cast_into::() { + Ok(ty) if ty.is_subclass_of::()? => ty, + _ => { + return Err(PyTypeError::new_err( + "exceptions must derive from BaseException", + )); + } + }, }; - let (instance, traceback) = match instantiate(&ty, value) { - Ok(instance) if instance.is_instance_of::() => (instance, traceback), - Ok(_) => ( - PyTypeError::new_err("exception constructor did not return an exception") - .into_value(py) - .into_bound(py) - .into_any(), - traceback, - ), - Err(err) => { - let own = err.traceback(py).map(Bound::into_any); - ( - err.into_value(py).into_bound(py).into_any(), - own.or(traceback), - ) + let py = ty.py(); + let exc = match instantiate(&ty, value) { + Ok(exc) => exc, + // A failed constructor raises its own error, keeping that error's traceback if it has + // one. PyPy reports a missing traceback as `None`. + Err(err) if err.traceback(py).is_some_and(|tb| !tb.is_none()) => { + return Ok(err.into_value(py).into_bound(py)); } + Err(err) => err.into_value(py).into_bound(py), }; - set_traceback(&instance, traceback.as_ref())?; - Ok(instance) + set_traceback(&exc, traceback.as_ref())?; + Ok(exc) } -/// `ty`'s exception for `value`: `value` itself if it already is one, else `ty` called with -/// no arguments, the tuple's items or `value`. +/// `value` itself if it already is a `ty` instance, else `ty()`, `ty(*value)` for a tuple or +/// `ty(value)`. A result that is not an exception is a `TypeError`. fn instantiate<'py>( ty: &Bound<'py, PyType>, value: Option>, -) -> PyResult> { - match value { +) -> PyResult> { + let exc = match value { Some(value) if value.is_instance_of::() && value.get_type().is_subclass(ty)? => { - Ok(value) + value } - None => ty.call0(), + None => ty.call0()?, Some(value) => match value.cast_into::() { - Ok(args) => ty.call1(args), - Err(err) => ty.call1((err.into_inner(),)), + Ok(args) => ty.call1(args)?, + Err(err) => ty.call1((err.into_inner(),))?, }, + }; + match exc.cast_into::() { + Ok(exc) => Ok(exc), + Err(err) => Err(PyTypeError::new_err(format!( + "calling {ty} should have returned an instance of BaseException, not {}", + err.into_inner().get_type().name()? + ))), } } -/// Set or clear `exception`'s traceback through `BaseException`, never an override. -fn set_traceback<'py>( - exception: &Bound<'py, PyAny>, - traceback: Option<&Bound<'py, PyAny>>, +/// Set or clear `exc`'s traceback through `BaseException`, never an override. +fn set_traceback( + exc: &Bound<'_, PyBaseException>, + traceback: Option<&Bound<'_, PyAny>>, ) -> PyResult<()> { - let py = exception.py(); + let py = exc.py(); py.get_type::() .getattr(intern!(py, "with_traceback"))? - .call1((exception, traceback)) + .call1((exc, traceback)) .map(drop) } diff --git a/src/coroutine/scope.rs b/src/coroutine/scope.rs index 419bb5ed..6a2ef446 100644 --- a/src/coroutine/scope.rs +++ b/src/coroutine/scope.rs @@ -5,7 +5,7 @@ use std::mem; use pyo3::{ PyTraverseError, PyVisit, - exceptions::{PyRuntimeError, PyStopIteration, PyTypeError}, + exceptions::{PyBaseException, PyRuntimeError, PyStopIteration, PyTypeError}, intern, prelude::*, }; @@ -171,13 +171,17 @@ impl Coroutine { /// Raise `exc` at the await point: inside the `__aenter__` awaitable being driven, if /// any, otherwise ending the coroutine. - pub(super) fn throw_step(&mut self, py: Python<'_>, exc: Bound<'_, PyAny>) -> PyResult { + pub(super) fn throw_step( + &mut self, + py: Python<'_>, + exc: Bound<'_, PyBaseException>, + ) -> PyResult { if let Scope::Opening { delegate, .. } = &self.scope { let step = delegate.clone_ref(py).throw(py, exc); return self.resume(py, step); } self.abandon(); - Err(PyErr::from_value(exc)) + Err(PyErr::from_value(exc.into_any())) } /// Finish without entering, dropping any context manager being entered. @@ -305,12 +309,12 @@ impl Delegate { } } - fn throw(&self, py: Python<'_>, exc: Bound<'_, PyAny>) -> PyResult { + fn throw(&self, py: Python<'_>, exc: Bound<'_, PyBaseException>) -> PyResult { match self { Delegate::Native(coroutine) => coroutine.bind(py).try_borrow_mut()?.throw_step(py, exc), Delegate::Foreign(iter) => match iter.bind(py).getattr(intern!(py, "throw")) { Ok(throw) => iter_step(py, throw.call1((exc,))), - Err(_) => Err(PyErr::from_value(exc)), + Err(_) => Err(PyErr::from_value(exc.into_any())), }, } } diff --git a/tests/cancellation_test.py b/tests/cancellation_test.py index d3c1b5cc..8aa82481 100644 --- a/tests/cancellation_test.py +++ b/tests/cancellation_test.py @@ -95,9 +95,43 @@ def __new__(cls): assert type(caught) is HiddenTraceback assert caught.args == ("constructor failure",) assert not contains(BaseException.__traceback__.__get__(caught)) - assert type(invoke(NotAnException)) is TypeError + # A failure without a traceback of its own takes the given one. + for args in ((NotAnException, None), (UnicodeDecodeError, ("a",))): + caught = invoke(*args, origin) + assert type(caught) is TypeError + assert contains(caught.__traceback__) + + unrelated = KeyError("unrelated") + caught = invoke(ValueError, unrelated, origin) + assert type(caught) is ValueError + assert caught.args == (unrelated,) + assert contains(caught.__traceback__) + assert invoke(ValueError, value=("keyword",), traceback=origin).args == ("keyword",) + + class RefusingMeta(type): + def __subclasscheck__(cls, subclass): + raise ZeroDivisionError + + class Refusing(Exception, metaclass=RefusingMeta): + pass + + assert type(invoke(Refusing, ValueError())) is ZeroDivisionError + + constructed = [] + + class NotRaisable: + def __init__(self): + constructed.append(self) + + assert type(invoke(NotRaisable)) is TypeError + assert not constructed coroutine = wreq.get("") + + class UsesCoroutine(Exception): + def __init__(self): + super().__init__(coroutine.__qualname__) + try: for args in ( (object(),), @@ -107,6 +141,13 @@ def __new__(cls): ): with pytest.raises(TypeError): coroutine.throw(*args) + # A rejected throw leaves the coroutine running, as a generator's does. + with pytest.raises(Exception) as caught: + coroutine.send(None) + assert "cannot reuse" not in str(caught.value) + # The exception is built before the throw borrows the coroutine. + with pytest.raises(UsesCoroutine): + coroutine.throw(UsesCoroutine) with pytest.raises(ValueError) as caught: coroutine.throw(error) assert caught.value is error From 278dda4ddbb85a2eae0342c7fe17dfd1b11c8d84 Mon Sep 17 00:00:00 2001 From: 0x676e67 Date: Wed, 7 Oct 2026 17:36:45 +0000 Subject: [PATCH 25/25] docs(coroutine): keep the Python docstring free of Rust links and cover reentrant subclass checks --- src/coroutine/awaitable.rs | 8 +++++--- tests/cancellation_test.py | 11 ++++++++++- 2 files changed, 15 insertions(+), 4 deletions(-) diff --git a/src/coroutine/awaitable.rs b/src/coroutine/awaitable.rs index 463ac4e7..a609707e 100644 --- a/src/coroutine/awaitable.rs +++ b/src/coroutine/awaitable.rs @@ -22,10 +22,12 @@ use super::{Port, scope::Scope}; /// /// Each suspension hands the task a real asyncio future, so task cancellation and /// the C task fast path work unchanged. A throw it accepts, or a close, drops the Rust -/// future, which aborts any spawned work. PyO3's borrow flag rejects reentrant polls. +/// future, which aborts any spawned work. /// -/// A coroutine built with [`managed`](super::managed) also works as `async with`, as -/// `async with await` would; the `scope` module holds that state. +/// A request coroutine also works as `async with`, as `async with await` would. +// +// The doc above is the Python docstring. PyO3's borrow flag rejects reentrant polls; +// coroutines built by `managed` support `async with`, and the `scope` module holds that state. #[pyclass(module = "wreq")] pub struct Coroutine { qualname: &'static str, diff --git a/tests/cancellation_test.py b/tests/cancellation_test.py index 8aa82481..d6c70b34 100644 --- a/tests/cancellation_test.py +++ b/tests/cancellation_test.py @@ -128,7 +128,12 @@ def __init__(self): coroutine = wreq.get("") - class UsesCoroutine(Exception): + class ChecksCoroutine(type): + def __subclasscheck__(cls, subclass): + assert coroutine.__qualname__ + return super().__subclasscheck__(subclass) + + class UsesCoroutine(Exception, metaclass=ChecksCoroutine): def __init__(self): super().__init__(coroutine.__qualname__) @@ -148,6 +153,10 @@ def __init__(self): # The exception is built before the throw borrows the coroutine. with pytest.raises(UsesCoroutine): coroutine.throw(UsesCoroutine) + uses = UsesCoroutine() + with pytest.raises(UsesCoroutine) as caught: + coroutine.throw(UsesCoroutine, uses) + assert caught.value is uses with pytest.raises(ValueError) as caught: coroutine.throw(error) assert caught.value is error