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/bench/README.md b/bench/README.md index e587298c..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 @@ -205,13 +205,17 @@ 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 -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. +`wreq_blocking_st` shares one +`Runtime(scheduler=Scheduler.PER_WORKER, workers=1)` across all logical workers +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/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..3d433b50 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,11 +117,19 @@ def settings(self): runtime={ "kind": "custom", "workers": 1, - "work_steal": False, + "scheduler": "PER_WORKER", "scope": "shared per case", }, 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 01858289..4caea682 100644 --- a/bench/report.py +++ b/bench/report.py @@ -42,8 +42,13 @@ 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: + 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/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_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/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/clippy.toml b/clippy.toml new file mode 100644 index 00000000..d0673276 --- /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::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" }, +] +disallowed-types = [{ path = "tokio::runtime::Handle", reason = "use Runtime" }] 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..838071fa 100644 --- a/docs/source/guide/advanced.md +++ b/docs/source/guide/advanced.md @@ -13,9 +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, with bounded buffering before the request consumes them. -Use synchronous iterators for blocking uploads. 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 @@ -98,18 +103,18 @@ 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 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,28 +122,44 @@ 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 -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. -Custom runtimes start their threads during construction, before any client is -bound or request is sent. +`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. +- `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. 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. +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. `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. -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 +(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. +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/docs/source/guide/blocking.md b/docs/source/guide/blocking.md index e731f43d..b2af2f3f 100644 --- a/docs/source/guide/blocking.md +++ b/docs/source/guide/blocking.md @@ -72,20 +72,58 @@ 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()) ``` -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 configuration and lifetime details. +### Current-thread runtime + +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 + +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 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`. +- 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..5de10013 100644 --- a/docs/templates/benchmark.md +++ b/docs/templates/benchmark.md @@ -60,13 +60,15 @@ 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 -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 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/python/wreq/runtime.py b/python/wreq/runtime.py index 1a3c6ab7..564e4f91 100644 --- a/python/wreq/runtime.py +++ b/python/wreq/runtime.py @@ -1,35 +1,56 @@ 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 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. - 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..1ece1bad 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; @@ -11,7 +10,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; @@ -277,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. @@ -291,17 +291,16 @@ 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 } } impl Default for Client { fn default() -> Self { - let runtime = runtime::get(); + let runtime = runtime::Runtime::shared(); Self { inner: wreq::Client::default(), runtime: runtime.clone(), @@ -312,16 +311,13 @@ 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()?, - None => runtime::get().clone(), + None => runtime::Runtime::shared().clone(), }; // Create the client builder. let mut builder = wreq::Client::builder(); @@ -541,6 +537,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 +719,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. @@ -815,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. @@ -832,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 854897dd..6ff8bd06 100644 --- a/src/client/body/multipart.rs +++ b/src/client/body/multipart.rs @@ -8,13 +8,12 @@ use crate::{ error::Error, extractor::{Binary, Text}, header::HeaderMap, - 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, } @@ -49,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) @@ -76,13 +73,21 @@ 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 and belongs + // to the thread sending it. + let parts = multipart + .parts + .iter_mut() + .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 }) } } @@ -123,25 +128,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) => py - .detach(|| { - runtime::get() - .handle() - .block_on(multipart::Part::file(path)) - }) - .map_err(Error::from)?, + Value::File(path) => multipart::Part::file(path).await.map_err(Error::from)?, Value::Stream(stream) => { let stream = Body::wrap_stream(stream); match self.length { @@ -151,7 +144,7 @@ impl Part { } }; - if let Some(filename) = self.filename.clone() { + if let Some(filename) = self.filename { inner = inner.file_name(filename); } @@ -159,11 +152,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/body/stream.rs b/src/client/body/stream.rs index bb051763..8fff96a6 100644 --- a/src/client/body/stream.rs +++ b/src/client/body/stream.rs @@ -1,36 +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. The pump +//! and the forwarder queue chunks through `queue`, which bounds how far they read ahead. + +mod async_gen; +mod iter; +mod pump; +mod queue; use std::{ pin::Pin, - sync::{ - Arc, Mutex, MutexGuard, PoisonError, - atomic::{AtomicBool, Ordering}, - }, + sync::{Mutex, MutexGuard, PoisonError}, task::{Context, Poll}, - time::Duration, }; use bytes::Bytes; -use futures_util::Stream; +use futures_util::{Stream, StreamExt, future::Either}; use pyo3::{ - exceptions::{PyRuntimeError, PyStopIteration}, + exceptions::PyStopIteration, intern, prelude::*, - sync::PyOnceLock, types::{PyIterator, PyString}, }; -use tokio::{ - runtime::Handle, - sync::mpsc::{self, error::TrySendError}, - task::spawn_blocking, - 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; @@ -42,49 +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, read on Tokio's blocking pool. -/// -/// 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 { - rx: mpsc::Receiver, - /// The iterator and sender while no pump runs: before the first poll, or once a pump - /// parked at a full channel. - parked: Arc>>, -} - -/// What a pump task needs to read the iterator into the body. -struct Pump { - iter: Py, - 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. -struct PyAsyncStream { - rx: mpsc::Receiver, - /// 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. -#[pyclass(frozen)] -struct Sender { - tx: Mutex>>, - finished: Arc, -} +pub struct PyStream(Either); // ===== impl PyBytesLike ===== @@ -101,6 +54,16 @@ impl FromPyObject<'_, '_> for PyBytesLike { } } +impl PyBytesLike { + /// The chunk's length in bytes. + fn len(&self) -> usize { + match self { + PyBytesLike::Bytes(b) => b.0.len(), + PyBytesLike::String(s) => s.0.len(), + } + } +} + impl From for Bytes { #[inline] fn from(value: PyBytesLike) -> Self { @@ -119,105 +82,32 @@ 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::new(ob.to_owned().unbind())) + Either::Left(SyncStream::new(ob.to_owned().unbind())) }; Ok(PyStream(source)) } } -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), - } - } -} - -// ===== impl SyncStream ===== - -impl SyncStream { - fn new(iter: Py) -> Self { - let (tx, rx) = mpsc::channel(1); - SyncStream { - rx, - parked: Arc::new(Mutex::new(Some(Pump { iter, tx }))), - } - } - - 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 - // 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)); +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(); } - poll } } -// ===== impl Pump ===== - -impl Pump { - /// How long a pump waits for room before parking to free its thread. - const PARK_AFTER: Duration = Duration::from_millis(10); +impl Stream for PyStream { + type Item = Item; - /// Send items until the iterator ends or raises, the body is dropped, or the channel - /// stays full past [`PARK_AFTER`](Self::PARK_AFTER). - fn run(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()); - match py.detach(|| handle.block_on(wait)) { - Ok(Ok(permit)) => permit, - 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. - let mut slot = lock(parked); - match self.tx.clone().try_reserve_owned() { - Ok(permit) => permit, - Err(TrySendError::Closed(_)) => return, - Err(TrySendError::Full(_)) => { - *slot = Some(self); - return; - } - } - } - } - } - }; - let Some(item) = next_item(py, &self.iter) else { - return; - }; - let failed = item.is_err(); - permit.send(item); - if failed { - return; - } - } - }); + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + self.get_mut().0.poll_next_unpin(cx) } } -#[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__")) { @@ -227,158 +117,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: - if not await sender.send(item, False): - 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::channel(1); - 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 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) { - Poll::Ready(Some(item)) => Poll::Ready(Some(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(); - 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, or `item` as the error that ends the body. 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()?) - }; - 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, - }) - }) - } - - /// 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) - } +/// 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..0c122e47 --- /dev/null +++ b/src/client/body/stream/async_gen.rs @@ -0,0 +1,186 @@ +//! Request bodies from Python async generators. + +use std::{ + ffi::CStr, + pin::Pin, + sync::Mutex, + task::{Context, Poll, ready}, +}; + +use futures_util::{Stream, StreamExt}; +use pyo3::{IntoPyObjectExt, intern, prelude::*, sync::PyOnceLock}; +use tokio::sync::TryAcquireError; + +use super::{Item, PyBytesLike, lock, queue}; +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. +pub(super) struct AsyncStream { + rx: queue::Rx, + /// 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>, +} + +// ===== 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 sender = Sender { + tx: Mutex::new(Some(tx)), + }; + 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, + 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(); + 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) + } +} + +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. + Runtime::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(TryAcquireError::NoPermits) => tx.clone(), + Err(TryAcquireError::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. + coroutine::local(py, "Sender.try_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, which skips the budget, so finishing never waits. + fn finish(&self) { + // Python may retain the sender after completion, especially on PyPy. + 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. + 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..212f0444 --- /dev/null +++ b/src/client/body/stream/iter.rs @@ -0,0 +1,132 @@ +//! Request bodies from Python iterators. + +use std::{ + mem, + pin::Pin, + task::{Context, Poll, ready}, + thread::{self, ThreadId}, +}; + +use futures_util::{Stream, StreamExt}; +use pyo3::{ + exceptions::{PyException, PyKeyboardInterrupt}, + prelude::*, +}; + +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 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. + Idle(Inline), + 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 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, +} + +// ===== impl SyncStream ===== + +impl SyncStream { + /// An iterator the calling thread passes to a request. + pub(super) fn new(iter: Py) -> Self { + 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(); + } + } +} + +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(inline) = mem::replace(this, SyncStream::Done) + { + *this = if Runtime::driving() { + SyncStream::Inline(inline) + } else { + SyncStream::Pumped(Pumped::new(inline.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; + Runtime::defer_wake(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, 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, + // 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 + }) + .unwrap_or_else(|| Some(Err(queue::unfinished()))) + .into() + } +} diff --git a/src/client/body/stream/pump.rs b/src/client/body/stream/pump.rs new file mode 100644 index 00000000..9e23acfb --- /dev/null +++ b/src/client/body/stream/pump.rs @@ -0,0 +1,135 @@ +//! 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::{sync::TryAcquireError, 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 +/// 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(); + Runtime::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>) { + // 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() { + Some(chunk) => chunk, + None => match next_item(py, &self.iter) { + Some(Ok(chunk)) => chunk, + Some(Err(err)) => { + self.tx.fail(err); + return; + } + None => { + self.tx.finish(); + return; + } + }, + }; + 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 + // sees the pump parked or the pump sees budget. + let mut slot = lock(parked); + match self.tx.try_reserve(&chunk) { + Ok(permit) => permit, + Err(TryAcquireError::Closed) => return, + Err(TryAcquireError::NoPermits) => { + 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<'_>, chunk: &PyBytesLike) -> Option> { + match self.tx.try_reserve(chunk) { + Ok(permit) => Some(permit), + 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 new file mode 100644 index 00000000..8128e1cb --- /dev/null +++ b/src/client/body/stream/queue.rs @@ -0,0 +1,174 @@ +//! 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, ready}, +}; + +use futures_util::Stream; +use pyo3::{PyErr, exceptions::PyRuntimeError}; +use tokio::sync::{OwnedSemaphorePermit, Semaphore, TryAcquireError, mpsc}; + +use super::{Item, PyBytesLike}; + +/// 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. +const UPLOAD_CHARGE: usize = 4 * 1024; + +/// The producer end. The end and errors skip the budget, so finishing 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, + /// Set once the body ended or failed; later polls end it again. + ended: bool, +} + +/// 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: Option, + _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(), + }; + 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. +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: `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, + }) + } + + /// 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 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: Some(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: Some(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(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) + } +} + +impl Drop for Rx { + fn drop(&mut self) { + self.close(); + } +} diff --git a/src/client/nogil.rs b/src/client/nogil.rs deleted file mode 100644 index bb834ff6..00000000 --- a/src/client/nogil.rs +++ /dev/null @@ -1,43 +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::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. -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()))? -} - -/// 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, -{ - let handle = runtime.handle(); - let mut future = pin!(future); - let ready = { - let _runtime = handle.enter(); - future.as_mut().now_or_never() - }; - match ready { - Some(output) => output, - None => py.detach(|| handle.block_on(future)), - } -} diff --git a/src/client/req.rs b/src/client/req.rs index 3b079bd2..a39f3bec 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()); @@ -369,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, @@ -398,13 +398,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 +412,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 +532,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/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/http.rs b/src/client/resp/http.rs index 0080d1ad..7b5e1cd5 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 } } }) @@ -488,12 +486,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) + runtime.block_on_eager(py, read) } else { - py.detach(|| runtime.handle().block_on(read)) + runtime.block_on(py, read) } } } @@ -583,6 +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()?; let response = &self.0; match response.take_bytes(Some(READ_ATTACHED))? { BodyRead::Ready(bytes) => Ok(PyBuffer::from(bytes)), @@ -591,7 +592,7 @@ impl BlockingResponse { .collect_later(collect) .map_ok(|(_, bytes)| PyBuffer::from(bytes)) .map_err(Into::into); - py.detach(|| response.runtime.handle().block_on(read)) + response.runtime.block_on(py, read) } } } diff --git a/src/client/resp/stream.rs b/src/client/resp/stream.rs index b4f9dfb6..a0d1e75b 100644 --- a/src/client/resp/stream.rs +++ b/src/client/resp/stream.rs @@ -22,10 +22,9 @@ use tokio::sync::{ }; use tokio_util::task::AbortOnDropHandle; -use super::{READ_ATTACHED, loop_limit}; +use super::loop_polls; use crate::{ buffer::PyBuffer, - client::nogil, coroutine::{self, Coroutine}, error::Error, header::HeaderMap, @@ -41,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, @@ -74,6 +74,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 { @@ -123,10 +133,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 @@ -139,6 +147,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>, @@ -159,27 +177,18 @@ 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. + fn ready_frame(&self, iteration: Iteration) -> 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 !iteration.takes_ready(resp) { return None; } - let _runtime = self.runtime.handle().enter(); - let Some(frame) = Self::poll_ready(resp)? else { + let Some(frame) = self.runtime.poll_now(Self::next_frame(resp))? else { *state = State::Closed; - return Some(Err(end().into())); + return Some(Err(iteration.end().into())); }; if frame.is_err() { *state = State::Closed; @@ -188,7 +197,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.runtime.poll_now(Self::next_frame(resp)) { Some(None) => *state = State::Closed, next => { if let State::Idle(resp) = mem::replace(&mut *state, State::Closed) { @@ -199,18 +208,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); @@ -222,19 +219,19 @@ 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)), } } - /// 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) @@ -242,25 +239,39 @@ 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())) } } } - /// Wait for the next frame; `end` builds the error that ends iteration. - async fn next(&self, end: fn() -> Error) -> PyResult { + /// 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(()) + } + + /// 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. 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); @@ -276,11 +287,9 @@ impl Streamer { } fn __next__(&self, py: Python) -> PyResult { - // Frames already received are returned without releasing the GIL. - if let Some(frame) = self.ready_frame(None, || Error::StopIteration) { - return frame; - } - nogil::run(py, &self.runtime, self.next(|| Error::StopIteration)) + // 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)) } fn __enter__(slf: PyRef) -> PyRef { @@ -300,32 +309,30 @@ 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 { 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) - { - 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 { this.reader.since_yield.store(0, Ordering::Relaxed); coroutine::yield_now().await; } - this.next(|| Error::StopAsyncIteration).await + this.next(Iteration::Async).await }) } fn __aenter__(slf: Bound<'_, Self>) -> PyResult> { + slf.get().refuse_async()?; coroutine::ready("Streamer.__aenter__", slf) } @@ -383,3 +390,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, + } + } +} diff --git a/src/client/resp/ws.rs b/src/client/resp/ws.rs index d83afb43..a8ddd72c 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, @@ -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, @@ -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/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/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/coroutine/awaitable.rs b/src/coroutine/awaitable.rs index c0e3ee43..a609707e 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}; @@ -20,11 +21,13 @@ 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. /// -/// 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, @@ -152,8 +155,20 @@ impl Coroutine { self.step(py, Some(value)).and_then(Step::into_result) } - fn throw(&mut self, py: Python<'_>, exc: Bound<'_, PyAny>) -> PyResult> { - self.throw_step(py, exc).and_then(Step::into_result) + // 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( + slf: &Bound<'_, Self>, + exc: Bound<'_, PyAny>, + value: Option>, + traceback: Option>, + ) -> PyResult> { + let exc = thrown(exc, value, traceback)?; + slf.try_borrow_mut()? + .throw_step(slf.py(), exc) + .and_then(Step::into_result) } fn close(&mut self, py: Python<'_>) -> PyResult<()> { @@ -204,6 +219,95 @@ impl Step { } } +/// The exception `throw(exc, value, traceback)` raises, built as a generator's `throw` +/// builds it. +fn thrown<'py>( + exc: Bound<'py, PyAny>, + value: Option>, + traceback: Option>, +) -> PyResult> { + 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", + )); + } + let ty = match exc.cast_into::() { + Ok(_) if value.is_some() => { + return Err(PyTypeError::new_err( + "instance exception may not have a separate value", + )); + } + 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 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(&exc, traceback.as_ref())?; + Ok(exc) +} + +/// `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> { + let exc = match value { + Some(value) + if value.is_instance_of::() && value.get_type().is_subclass(ty)? => + { + value + } + None => ty.call0()?, + Some(value) => match value.cast_into::() { + 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 `exc`'s traceback through `BaseException`, never an override. +fn set_traceback( + exc: &Bound<'_, PyBaseException>, + traceback: Option<&Bound<'_, PyAny>>, +) -> PyResult<()> { + let py = exc.py(); + py.get_type::() + .getattr(intern!(py, "with_traceback"))? + .call1((exc, 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/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/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)) - } -} 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/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 37cb05e8..96d95994 100644 --- a/src/runtime.rs +++ b/src/runtime.rs @@ -1,30 +1,85 @@ +//! 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`]). 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. +#![allow(clippy::disallowed_methods, clippy::disallowed_types)] + use std::{ + cell::{Cell, RefCell}, + future::Future, + pin::pin, sync::{Arc, OnceLock}, + task::Context, thread, time::Duration, }; +use futures_util::FutureExt; use pingora_runtime::{BlockingPoolOpts, Runtime as PingoraRuntime, RuntimeBuilder}; use pyo3::{ exceptions::{PyRuntimeError, PyValueError}, prelude::*, }; -use tokio::runtime::Handle; +use tokio::{ + runtime::{Builder, Handle}, + task::{self, JoinError, JoinHandle}, +}; +use tokio_util::task::AbortOnDropHandle; + +/// 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 Inner { + Workers(PingoraRuntime), + /// Driven only by threads in a blocking call, through [`Driving::drive`]. + 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) }; + /// 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) }; +} + +// ===== 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(Inner::CurrentThread(_))) } /// Share the runtime and select a worker for a new client. @@ -34,16 +89,190 @@ 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(), }) } + + /// Poll `future` once inside the runtime, where a read timeout can start its timer; + /// `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() + } + + /// 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. + Self::refuse_nested()?; + let mut future = pin!(future); + match self.poll_now(future.as_mut()) { + 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)? + } + + /// 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 { + Self::refuse_nested()?; + match self.inner.as_deref() { + Some(Inner::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, + { + Self::refuse_nested()?; + match self.inner.as_deref() { + 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)) + .map_err(join_error)?, + } + } } -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(); +/// 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 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, + 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: Inner) -> Self { + let handle = runtime.handle(); Self { inner: Some(Arc::new(runtime)), handle, @@ -51,6 +280,12 @@ impl From for Runtime { } } +impl From for Runtime { + fn from(runtime: PingoraRuntime) -> Self { + Inner::Workers(runtime).into() + } +} + impl FromPyObject<'_, '_> for Runtime { type Error = PyErr; @@ -61,52 +296,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, 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] #[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()))?; + Inner::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,34 +374,63 @@ 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), + Inner::Workers(PingoraRuntime::Steal { runtime, .. }) + | Inner::CurrentThread(runtime) => runtime.shutdown_background(), + Inner::Workers(PingoraRuntime::NoSteal(runtime)) => drop(runtime), } } } } -/// Workers used when none are given: one per available CPU. -fn parallelism() -> usize { - thread::available_parallelism().map_or(1, usize::from) +// ===== impl Inner ===== + +impl Inner { + fn handle(&self) -> Handle { + match self { + // No-steal workers are lazy upstream; this starts them. + Inner::Workers(runtime) => runtime.get_handle().clone(), + Inner::CurrentThread(runtime) => runtime.handle().clone(), + } + } } -/// Create the shared runtime on first use and retain it for the process lifetime. -pub fn get() -> &'static Runtime { - static RUNTIME: OnceLock = OnceLock::new(); +// ===== impl Driving ===== - fn create() -> Runtime { - let runtime = RuntimeBuilder::new(parallelism(), env!("CARGO_PKG_NAME")).build(); - runtime.into() +impl Driving { + /// Drive `runtime` on this thread until `future` completes, marked as driving it; an + /// interrupt an upload iterator stashed meanwhile replaces the output. + fn drive( + runtime: &tokio::runtime::Runtime, + future: impl Future>, + ) -> PyResult { + DRIVING.set(true); + let _driving = Driving; + let output = runtime.block_on(future); + INTERRUPT.take().map_or(output, Err) } +} - if let Some(runtime) = RUNTIME.get() { - return runtime; +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(); } +} + +/// A task that panicked or was aborted, as a Python error. +fn join_error(err: JoinError) -> PyErr { + PyRuntimeError::new_err(err.to_string()) +} - // 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)) +/// Workers used when none are given: one per available CPU. +fn parallelism() -> usize { + thread::available_parallelism().map_or(1, usize::from) +} + +/// 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)] @@ -161,19 +449,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); } @@ -182,6 +471,40 @@ 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(Inner::CurrentThread( + Builder::new_current_thread().enable_all().build().unwrap(), + )) + }; + let runtime = current_thread(); + // Spawned tasks and timers run on the blocked caller. + let child = runtime + .wait_task(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 + .wait(async { + let nested = other.wait(async { Ok(()) }).is_err(); + drop(other); + Ok(nested) + }) + .unwrap(); + assert!(nested); + assert!(!Runtime::driving()); + } + #[test] fn last_owner_can_be_released_on_a_worker() { struct NotifyOnDrop(std::sync::mpsc::Sender<()>); @@ -199,7 +522,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 { diff --git a/tests/cancellation_test.py b/tests/cancellation_test.py index 63fca185..d6c70b34 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: @@ -101,13 +95,49 @@ 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 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__) + try: - installed = vars(type(coroutine))["throw"] - _install(type(coroutine)) - assert vars(type(coroutine))["throw"] is installed for args in ( (object(),), (ValueError, None, object()), @@ -116,6 +146,17 @@ 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) + 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 @@ -150,6 +191,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 +232,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: diff --git a/tests/multipart_test.py b/tests/multipart_test.py index f61a51e0..144ba359 100644 --- a/tests/multipart_test.py +++ b/tests/multipart_test.py @@ -1,8 +1,10 @@ +import asyncio from pathlib import Path import pytest import wreq from wreq import Multipart, Part +from wreq.runtime import Runtime, Scheduler client = wreq.Client(tls_info=True) @@ -90,3 +92,27 @@ 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) + + # 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) 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/runtime_test.py b/tests/runtime_test.py index 7e889a06..5a6bcc7b 100644 --- a/tests/runtime_test.py +++ b/tests/runtime_test.py @@ -1,25 +1,37 @@ import asyncio import base64 import hashlib +import threading +import time +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 +41,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 +64,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 +89,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 +126,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 +160,234 @@ 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(error): + yield b"x" + raise error + + def upload_interrupted(error=Interrupt): + client.post(url + "/interrupted", body=interrupted(error)) + + 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" + + # 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")) + 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 +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"), @@ -155,7 +396,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 +437,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 +466,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 +474,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 +483,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 +497,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 +510,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 +555,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/shutdown_test.py b/tests/shutdown_test.py index 6a15f5e0..54efdc1a 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 @@ -62,9 +62,10 @@ 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 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 +98,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 +107,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,7 +120,7 @@ async def chunks(): try: await asyncio.wait_for(full.wait(), 5) await asyncio.sleep(0.05) - assert produced == [0, 1] + assert produced == expected, produced if retain: return part finally: @@ -129,9 +131,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], @@ -140,4 +152,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 3d35043d..da288534 100644 --- a/tests/sync_upload_test.py +++ b/tests/sync_upload_test.py @@ -6,6 +6,9 @@ import pytest import wreq +from wreq.runtime import Runtime, Scheduler + +from upload_test import read_chunked @asynccontextmanager @@ -65,12 +68,38 @@ async def accept(reader, writer): assert not errors, errors +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 kind == "blocking" else wreq.Client)(proxies=[]) + + +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", [False, True]) +@pytest.mark.parametrize("kind", UPLOAD_KINDS) @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(kind, multipart): + client = upload_client(kind) try: async with upload_server() as (url, bodies): for failure in (False, True): @@ -95,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) @@ -124,6 +145,62 @@ async def send(): client.close() +@pytest.mark.asyncio +@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 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) + 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(kind) + body = async_generator() if source == "async_generator" else iterator() + try: + task = asyncio.create_task(post_bytes(client, url, body=body)) + for _ in range(250): + 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" + 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. @@ -151,8 +228,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("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() @@ -182,20 +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]}/" - factory = wreq.blocking.Client if blocking else wreq.Client - client = factory(proxies=[]) + 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() @@ -254,6 +322,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("current_thread") + 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.