diff --git a/Cargo.lock b/Cargo.lock index bd382ab8..8aedd18f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -802,8 +802,7 @@ dependencies = [ [[package]] name = "http2" version = "0.5.20" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "92d3114be2f413b2e491e686b93a28cda30c355cffc8d091a57f8be4b1342896" +source = "git+https://github.com/0x676e67/http2?rev=fb5bf8bdbf3e9d9aad7073759a6f06f47373342e#fb5bf8bdbf3e9d9aad7073759a6f06f47373342e" dependencies = [ "atomic-waker", "bytes", diff --git a/Cargo.toml b/Cargo.toml index ea83fda9..17fd2967 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -84,6 +84,7 @@ btls = { git = "https://github.com/0x676e67/btls", branch = "main" } btls-sys = { git = "https://github.com/0x676e67/btls", branch = "main" } tokio-btls = { git = "https://github.com/0x676e67/btls", branch = "main" } netty = { git = "https://github.com/0x676e67/netty", branch = "main" } +http2 = { git = "https://github.com/0x676e67/http2", rev = "fb5bf8bdbf3e9d9aad7073759a6f06f47373342e" } [profile.release] codegen-units = 1 diff --git a/docs/source/guide/redirect-errors.md b/docs/source/guide/redirect-errors.md index 1ae03843..921c282f 100644 --- a/docs/source/guide/redirect-errors.md +++ b/docs/source/guide/redirect-errors.md @@ -107,6 +107,9 @@ reading its response body. `connect_timeout` limits connection establishment; `read_timeout` limits the wait for a read and resets after successful reads. Requests can override `timeout` and `read_timeout`. +Waiting before the first response body read does not use up its `read_timeout`. +An overall `timeout` still runs during that wait. + ```python import asyncio from datetime import timedelta diff --git a/docs/source/guide/runtime.md b/docs/source/guide/runtime.md index 5d545843..c889f113 100644 --- a/docs/source/guide/runtime.md +++ b/docs/source/guide/runtime.md @@ -20,6 +20,10 @@ Both worker schedulers support async and blocking clients. `CURRENT_THREAD` supports only `wreq.blocking.Client`; passing it to `wreq.Client` raises `ValueError`. +With `WORK_STEALING`, blocking requests run on the calling thread while +background workers handle connection I/O. `PER_WORKER` submits requests to the +client's selected worker. Both keep processing connections between calls. + Runtime worker count and request concurrency are different settings. For example, one async client can have many requests in flight with `workers=1`. With `PER_WORKER`, increasing the worker count distributes newly created clients @@ -74,7 +78,8 @@ or all process resources. ## Use the calling thread for blocking I/O {#current-thread} -`CURRENT_THREAD` avoids dispatching each request to a background I/O worker. +`CURRENT_THREAD` runs connection I/O on the calling thread, without background +I/O workers. Create the client and runtime together and reuse them for subsequent calls: ```python diff --git a/src/client.rs b/src/client.rs index cc1e2c3e..a5672992 100644 --- a/src/client.rs +++ b/src/client.rs @@ -1,12 +1,14 @@ pub mod body; pub mod req; pub mod resp; +pub mod ws; mod param; mod query; use std::{ net::{IpAddr, Ipv4Addr, Ipv6Addr}, + sync::Arc, time::Duration, }; @@ -20,7 +22,8 @@ use wreq::tls::trust::CertStore; use self::{ req::{execute_request, execute_websocket_request}, - resp::{BlockingResponse, BlockingWebSocket, Response, WebSocket}, + resp::{BlockingResponse, Response}, + ws::{BlockingWebSocket, WebSocket}, }; use crate::{ cookie::Jar, @@ -211,7 +214,11 @@ struct Builder { /// Unknown keyword arguments raise `TypeError` at call time. #[derive(Clone)] #[pyclass(subclass, frozen)] -pub struct Client { +pub struct Client(Arc); + +/// The state a [`Client`] shares with its requests. Each request task owns one clone, which +/// keeps the runtime alive, so starting a request clones nothing else. +pub struct ClientInner { inner: wreq::Client, /// Runs this client's requests and response reads, keeping the runtime alive. runtime: runtime::Runtime, @@ -219,9 +226,8 @@ pub struct Client { cancel: CancellationToken, /// Turn error statuses into exceptions for every request. raise_for_status: bool, - - /// Get the cookie jar of the client. - #[pyo3(get)] + /// Prefetching must not start a configured read timeout before the caller reads. + has_read_timeout: bool, cookie_jar: Option, } @@ -256,10 +262,13 @@ impl Client { kwds: Option>, ) -> PyResult { let kwds = Python::attach(|py| kwds.map(|kwds| kwds.bind(py).extract()).transpose())?; - self.runtime - .clone() - .run_task(execute_request(self, method, url, kwds)) - .await + let Self(client) = self; + let task = client + .runtime + .spawn(execute_request(client.clone(), method, url, kwds)); + // Release ours before waiting, leaving the task's clone as the only extra count. + drop(client); + runtime::Runtime::join(task).await } /// Open a WebSocket on the client's runtime, extracting options on first await. @@ -269,23 +278,27 @@ impl Client { kwds: Option>, ) -> PyResult { let kwds = Python::attach(|py| kwds.map(|kwds| kwds.bind(py).extract()).transpose())?; - self.runtime - .clone() - .run_task(execute_websocket_request(self, url, kwds)) - .await + let Self(client) = self; + let task = client + .runtime + .spawn(execute_websocket_request(client.clone(), url, kwds)); + // Released before waiting, as in `execute`. + drop(client); + runtime::Runtime::join(task).await } } impl Default for Client { fn default() -> Self { let runtime = runtime::Runtime::shared(); - Self { + Self(Arc::new(ClientInner { inner: wreq::Client::default(), runtime: runtime.clone(), cancel: CancellationToken::new(), raise_for_status: false, + has_read_timeout: false, cookie_jar: None, - } + })) } } @@ -297,6 +310,9 @@ impl Client { Some(runtime) => runtime.select()?, None => runtime::Runtime::shared().clone(), }; + let has_read_timeout = kwds + .as_ref() + .is_some_and(|config| config.read_timeout.is_some()); // Create the client builder. let mut builder = wreq::Client::builder(); let mut cookie_jar: Option = None; @@ -504,12 +520,15 @@ impl Client { builder .build() - .map(|inner| Client { - inner, - runtime, - cancel: CancellationToken::new(), - cookie_jar, - raise_for_status, + .map(|inner| { + Client(Arc::new(ClientInner { + inner, + runtime, + cancel: CancellationToken::new(), + cookie_jar, + raise_for_status, + has_read_timeout, + })) }) .map_err(Error::Library) .map_err(Into::into) @@ -539,13 +558,19 @@ impl Client { /// Cancel pending requests and reject new ones with asyncio.CancelledError. /// Existing responses, WebSockets and the shared runtime remain usable. pub fn close(&self) { - self.cancel.cancel(); + self.0.cancel.cancel(); } /// The runtime used by this client and its responses. #[getter] pub fn runtime(&self) -> runtime::Runtime { - self.runtime.clone() + self.0.runtime.clone() + } + + /// Get the cookie jar of the client. + #[getter] + pub fn cookie_jar(&self) -> Option { + self.0.cookie_jar.clone() } /// Make a GET request to the given URL. @@ -677,7 +702,7 @@ impl Client { _exc_val: Py, _traceback: Py, ) -> PyResult> { - let cancel = self.cancel.clone(); + let cancel = self.0.cancel.clone(); coroutine::local(py, "Client.__aexit__", async move { cancel.cancel(); Ok(()) @@ -705,7 +730,7 @@ impl BlockingClient { /// Get the cookie jar of the client. #[getter] pub fn cookie_jar(&self) -> Option { - self.0.cookie_jar.clone() + self.0.cookie_jar() } /// Cancel pending requests and reject new ones with asyncio.CancelledError. @@ -812,9 +837,10 @@ impl BlockingClient { url: PyBackedStr, kwds: Option, ) -> PyResult { - self.0 + let Client(client) = &self.0; + client .runtime - .block_on_task(py, execute_request(self.0.clone(), method, url, kwds)) + .block_on_task(py, execute_request(client.clone(), method, url, kwds)) .map(Into::into) } @@ -827,9 +853,10 @@ impl BlockingClient { url: PyBackedStr, kwds: Option, ) -> PyResult { - self.0 + let Client(client) = &self.0; + client .runtime - .block_on_task(py, execute_websocket_request(self.0.clone(), url, kwds)) + .block_on_task(py, execute_websocket_request(client.clone(), url, kwds)) .map(Into::into) } } diff --git a/src/client/body/stream/async_gen.rs b/src/client/body/stream/async_gen.rs index 0c122e47..a8e4b082 100644 --- a/src/client/body/stream/async_gen.rs +++ b/src/client/body/stream/async_gen.rs @@ -3,12 +3,17 @@ use std::{ ffi::CStr, pin::Pin, - sync::Mutex, + sync::{Mutex, OnceLock}, task::{Context, Poll, ready}, }; use futures_util::{Stream, StreamExt}; -use pyo3::{IntoPyObjectExt, intern, prelude::*, sync::PyOnceLock}; +use pyo3::{ + IntoPyObjectExt, intern, + prelude::*, + sync::PyOnceLock, + types::{PyCode, PyDict, PyType}, +}; use tokio::sync::TryAcquireError; use super::{Item, PyBytesLike, lock, queue}; @@ -53,12 +58,9 @@ impl AsyncStream { 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")); - })?; + let task = start(&event_loop, &coroutine).inspect_err(|_| { + let _ = coroutine.call_method0(intern!(py, "close")); + })?; Ok(Self { rx, task: Some((task.unbind(), event_loop.unbind())), @@ -153,6 +155,81 @@ impl Sender { } } +/// Run `coroutine` as a task on `event_loop`, capturing the caller's contextvars. +/// +/// A standard loop whose `create_task` passes extra keywords to `Task`, as on Python 3.14 and +/// later 3.13 releases, receives an eager-start hint so chunks can be queued before the +/// request is sent. Task creation always goes through the loop, including wrappers that +/// capture context or instrument tasks. +fn start<'py>( + event_loop: &Bound<'py, PyAny>, + coroutine: &Bound<'py, PyAny>, +) -> PyResult> { + let py = event_loop.py(); + let create_task = event_loop.getattr(intern!(py, "create_task"))?; + if eager(event_loop, &create_task) { + let kwargs = PyDict::new(py); + kwargs.set_item(intern!(py, "eager_start"), true)?; + return create_task.call((coroutine,), Some(&kwargs)); + } + create_task.call1((coroutine,)) +} + +/// Only pass the hint where the standard loop API supports it. Older wrappers and custom +/// factories keep their usual arguments; the actual loop method still controls scheduling. +fn eager(event_loop: &Bound<'_, PyAny>, create_task: &Bound<'_, PyAny>) -> bool { + static BASE: PyOnceLock> = PyOnceLock::new(); + static VERSION: OnceLock<(bool, bool)> = OnceLock::new(); + let py = event_loop.py(); + let (supported, forwards_keywords) = *VERSION.get_or_init(|| { + let version = py.version_info(); + (version >= (3, 12), version >= (3, 14)) + }); + if !supported { + return false; + } + let standard = || -> PyResult { + let base = BASE.import(py, "asyncio", "BaseEventLoop")?; + let name = intern!(py, "create_task"); + let method = base.getattr(name)?; + // Python 3.14 always passes extra keywords to `Task`, 3.12 never does, and 3.13 only + // in later releases; there only the unwrapped standard method shows whether it does. + let forwards = forwards_keywords + || (method + .getattr(intern!(py, "__module__"))? + .eq("asyncio.base_events")? + && takes_keywords(&method)?); + Ok(forwards + && event_loop.is_instance(base)? + && event_loop.get_type().getattr(name)?.is(&method) + && event_loop + .call_method0(intern!(py, "get_task_factory"))? + .is_none() + // Don't send a new keyword to an old-signature wrapper. + && takes_keywords(create_task)?) + }; + standard().unwrap_or(false) +} + +/// Whether the Python function `func` accepts `**kwargs` (`inspect.CO_VARKEYWORDS`). +fn takes_keywords(func: &Bound<'_, PyAny>) -> PyResult { + static CODE: PyOnceLock<(Py, bool)> = PyOnceLock::new(); + let py = func.py(); + let code = func.getattr(intern!(py, "__code__"))?; + if let Some((cached, keywords)) = CODE.get(py) + && code.is(cached) + { + return Ok(*keywords); + } + let keywords = code.getattr(intern!(py, "co_flags"))?.extract::()? & 0x08 != 0; + // Code objects are immutable; replacing a function or its code misses this cache. + // A custom callable can expose a mutable __code__ proxy, so never cache that. + if code.is_exact_instance_of::() { + let _ = CODE.set(py, (code.unbind(), keywords)); + } + Ok(keywords) +} + /// The coroutine that forwards an async generator into its [`Sender`]. const FORWARD_SOURCE: &CStr = c"import asyncio diff --git a/src/client/req.rs b/src/client/req.rs index 6ff23d19..4381e488 100644 --- a/src/client/req.rs +++ b/src/client/req.rs @@ -1,7 +1,8 @@ use std::{ future::Future, net::{IpAddr, Ipv4Addr, Ipv6Addr}, - pin::{Pin, pin}, + pin::Pin, + sync::Arc, time::Duration, }; @@ -13,10 +14,11 @@ use tokio_util::sync::CancellationToken; use crate::{ client::{ - Client, + ClientInner, body::{Body, Form, Json, multipart::Multipart}, query::Query, - resp::{Response, WebSocket}, + resp::Response, + ws::WebSocket, }, cookie::{Cookies, Jar}, emulate::EmulationLike, @@ -28,12 +30,21 @@ use crate::{ redirect, }; -/// The parameters for a request. +/// Extracted request parameters. A lone raw body needs no full options object. #[derive(Default, Deftly)] -#[derive_deftly(ExtractOptions)] -#[deftly(deny_unknown)] +#[derive_deftly(ExtractOptionsWrapper)] #[non_exhaustive] pub struct Request { + body: Option, + #[deftly(flatten = "Options")] + options: Option>, +} + +/// Full request options, used for combinations and non-body parameters. +#[derive(Default, Deftly)] +#[derive_deftly(ExtractOptions)] +#[deftly(deny_unknown, single_keyword)] +struct Options { /// The JSON body to use for the request. json: Option, @@ -192,7 +203,7 @@ pub struct WebSocketRequest { /// Build and send a request with `client`, failing with `CancelledError` once the client is /// closed. Callers run it on the client's runtime. pub async fn execute_request( - client: Client, + client: Arc, method: Method, url: U, request: Option, @@ -200,121 +211,133 @@ pub async fn execute_request( where U: AsRef, { - let future = pin!(async { + // Boxed so this future stays small: it is moved into a task on the event loop, while + // the send state below is large and only allocated once the task first runs. + let mut future = Box::pin(async { // Create the request builder. let mut builder = client.inner.request(method.into_ffi(), url.as_ref()); - - if let Some(mut request) = request { - // Emulation options. - apply_option!(set_if_some, builder, request.emulation, emulation); - - // Version options. - apply_option!( - set_if_some_map, - builder, - request.version, - version, - Version::into_ffi - ); - - // Timeout options. - apply_option!(set_if_some, builder, request.timeout, timeout); - apply_option!(set_if_some, builder, request.read_timeout, read_timeout); - - // Network options. - apply_option!(set_if_some_inner, builder, request.proxy, proxy); - apply_option!(set_if_some, builder, request.local_address, local_address); - apply_option!( - set_if_some_tuple_inner, - builder, - request.local_addresses, - local_addresses - ); - - #[cfg(any( - target_os = "android", - target_os = "fuchsia", - target_os = "illumos", - target_os = "ios", - target_os = "linux", - target_os = "macos", - target_os = "solaris", - target_os = "tvos", - target_os = "visionos", - target_os = "watchos", - ))] - apply_option!(set_if_some, builder, request.interface, interface); - - // Headers options. - apply_option!(set_if_some_inner, builder, request.headers, headers); - apply_option!( - set_if_some_inner, - builder, - request.orig_headers, - orig_headers - ); - apply_option!( - set_if_some, - builder, - request.default_headers, - default_headers - ); - - // Cookies options. - apply_option!( - set_if_some_iter_inner_with_key, - builder, - request.cookies, - header, - COOKIE - ); - apply_option!( - set_if_some_inner, - builder, - request.cookie_provider, - cookie_provider - ); - - // Authentication options. - apply_option!( - set_if_some_map_ref, - builder, - request.auth, - auth, - AsRef::::as_ref - ); - apply_option!(set_if_some, builder, request.bearer_auth, bearer_auth); - apply_option!(set_if_some_tuple, builder, request.basic_auth, basic_auth); - - // Allow redirects options. - apply_option!(set_if_some_inner, builder, request.redirect, redirect); - - // Compression options. - apply_option!(set_if_some, builder, request.gzip, gzip); - apply_option!(set_if_some, builder, request.brotli, brotli); - apply_option!(set_if_some, builder, request.deflate, deflate); - apply_option!(set_if_some, builder, request.zstd, zstd); - - // Query options. - apply_option!(set_if_some_ref, builder, request.query, query); - - // Body options. - apply_option!(set_if_some_ref, builder, request.form, form); - apply_option!(set_if_some_ref, builder, request.json, json); - if let Some(multipart) = request.multipart { - builder = builder.multipart(multipart.into_form().await?); + let prefetch_body = !client.has_read_timeout + && request + .as_ref() + .and_then(|request| request.options.as_deref()) + .is_none_or(|options| options.read_timeout.is_none()); + + if let Some(request) = request { + if let Some(body) = request.body { + builder = builder.body(wreq::Body::try_from(body)?); + } + if let Some(mut request) = request.options { + // Emulation options. + apply_option!(set_if_some, builder, request.emulation, emulation); + + // Version options. + apply_option!( + set_if_some_map, + builder, + request.version, + version, + Version::into_ffi + ); + + // Timeout options. + apply_option!(set_if_some, builder, request.timeout, timeout); + apply_option!(set_if_some, builder, request.read_timeout, read_timeout); + + // Network options. + apply_option!(set_if_some_inner, builder, request.proxy, proxy); + apply_option!(set_if_some, builder, request.local_address, local_address); + apply_option!( + set_if_some_tuple_inner, + builder, + request.local_addresses, + local_addresses + ); + + #[cfg(any( + target_os = "android", + target_os = "fuchsia", + target_os = "illumos", + target_os = "ios", + target_os = "linux", + target_os = "macos", + target_os = "solaris", + target_os = "tvos", + target_os = "visionos", + target_os = "watchos", + ))] + apply_option!(set_if_some, builder, request.interface, interface); + + // Headers options. + apply_option!(set_if_some_inner, builder, request.headers, headers); + apply_option!( + set_if_some_inner, + builder, + request.orig_headers, + orig_headers + ); + apply_option!( + set_if_some, + builder, + request.default_headers, + default_headers + ); + + // Cookies options. + apply_option!( + set_if_some_iter_inner_with_key, + builder, + request.cookies, + header, + COOKIE + ); + apply_option!( + set_if_some_inner, + builder, + request.cookie_provider, + cookie_provider + ); + + // Authentication options. + apply_option!( + set_if_some_map_ref, + builder, + request.auth, + auth, + AsRef::::as_ref + ); + apply_option!(set_if_some, builder, request.bearer_auth, bearer_auth); + apply_option!(set_if_some_tuple, builder, request.basic_auth, basic_auth); + + // Allow redirects options. + apply_option!(set_if_some_inner, builder, request.redirect, redirect); + + // Compression options. + apply_option!(set_if_some, builder, request.gzip, gzip); + apply_option!(set_if_some, builder, request.brotli, brotli); + apply_option!(set_if_some, builder, request.deflate, deflate); + apply_option!(set_if_some, builder, request.zstd, zstd); + + // Query options. + apply_option!(set_if_some_ref, builder, request.query, query); + + // Body options. + apply_option!(set_if_some_ref, builder, request.form, form); + apply_option!(set_if_some_ref, builder, request.json, json); + if let Some(multipart) = request.multipart { + builder = builder.multipart(multipart.into_form().await?); + } + apply_option!( + set_if_some_map_try, + builder, + request.body, + body, + wreq::Body::try_from + ); } - apply_option!( - set_if_some_map_try, - builder, - request.body, - body, - wreq::Body::try_from - ); } // Send request. - builder + let response = builder .send() .await .and_then(|r| { @@ -324,24 +347,28 @@ where Ok(r) } }) - .map(|response| Response::new(response, client.runtime.clone())) - .map_err(Error::Library) - .map_err(Into::into) + .map_err(Error::Library)?; + Ok(Response::new( + response, + client.runtime.clone(), + prefetch_body, + )) }); - until_closed(&client.cancel, future).await + until_closed(&client.cancel, future.as_mut()).await } /// Like [`execute_request`], opening a WebSocket. pub async fn execute_websocket_request( - client: Client, + client: Arc, url: U, request: Option, ) -> PyResult where U: AsRef, { - let future = pin!(async { + // Boxed for the same reason as in `execute_request`. + let mut future = Box::pin(async { // Create the WebSocket builder. let mut builder = client.inner.websocket(url.as_ref()); @@ -463,7 +490,7 @@ where .map_err(Into::into) }); - until_closed(&client.cancel, future).await + until_closed(&client.cancel, future.as_mut()).await } /// Run `future` unless the client is closed first; a closed client never polls it. Taking diff --git a/src/client/resp.rs b/src/client/resp.rs index edced877..3012f174 100644 --- a/src/client/resp.rs +++ b/src/client/resp.rs @@ -1,12 +1,34 @@ -mod ext; -mod http; +mod prefetch; mod stream; -mod ws; -pub use self::{ - http::{BlockingResponse, Response}, - stream::Streamer, - ws::{BlockingWebSocket, WebSocket, msg::Message}, +use std::{ + fmt::{self, Display}, + future::Future, + mem, + sync::{Arc, Mutex, MutexGuard, PoisonError}, +}; + +use bytes::Bytes; +use futures_util::{FutureExt, TryFutureExt, future::Either}; +use http::response::{Parts, Response as HttpResponse}; +use http_body::Body as _; +use http_body_util::{BodyExt, Collected, combinators::Collect}; +use pyo3::{prelude::*, pybacked::PyBackedStr, sync::PyOnceLock}; +use wreq::Uri; + +use self::prefetch::{Prefetch, prefetch}; +pub use self::stream::Streamer; +use crate::{ + buffer::PyBuffer, + client::{SocketAddr, body::Json}, + cookie::Cookie, + coroutine::{self, Coroutine, EntersSelf, Start}, + error::Error, + header::HeaderMap, + http::{StatusCode, Version}, + redirect::History, + runtime::Runtime, + tls::TlsInfo, }; /// Buffered bodies of known length up to this size are read on the event loop, and on a @@ -21,7 +43,685 @@ 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 { - loop_polls(version).then_some(limit) +/// A response from a request. +/// +/// Reads take the body with [`Response::take_bytes`]: a buffered body of known length up to +/// the read's limit finishes on the caller, the event loop or a blocking thread, and any other +/// body is collected on the runtime by [`Response::collect_later`]. +#[pyclass(subclass, frozen, str)] +pub struct Response { + uri: Uri, + /// Response head; rebuilt responses share its extensions, including the connection's + /// reuse flag. + parts: Parts, + /// Shared with a read in flight, so it can cache the bytes it reads. + body: Arc>, + /// Runs body reads and keeps the runtime alive while the response does. + runtime: Runtime, + /// Captured at receipt; `None` for a body without a known length. + content_length: Option, + /// Whether the event loop may poll the unread body: an HTTP/1 body, or one held in full. + loop_polls: bool, + remote_addr: Option, + local_addr: Option, +} + +/// The response body slot. +enum Body { + /// Unread; taken by the first read or `stream()`. + Unread(wreq::Body), + /// Read with the head, but not yet consumed by a caller. + Buffered(Bytes), + /// Taken by a read in flight or by `stream()`. + Taken, + /// Read in full and shared by later `text`, `json` and `bytes` calls. + Cached(Bytes), + /// Released by `close` or a context exit; later reads fail. + Released, +} + +/// A body taken for reading. +enum BodyRead { + /// Read in full now. + Ready(Bytes), + /// Not read now: still arriving, over the read's limit or of unknown length; collected + /// on the runtime by [`Response::collect_later`]. + Pending(Collect), +} + +/// A blocking response from a request. +#[pyclass(name = "Response", subclass, frozen, str)] +pub struct BlockingResponse(Response); + +/// Forbids connection reuse on drop unless disarmed by taking the parts. Held by +/// [`Response::collect_later`] from its creation until the body is read in full, so a failed +/// or cancelled read is not pooled. +struct RecycleGuard(Option); + +// ===== impl Response ===== + +impl Response { + /// Create a response, prefetching only when polling cannot start a read timeout. + pub fn new(response: wreq::Response, runtime: Runtime, prefetch_body: bool) -> Self { + let uri = response.uri().clone(); + let content_length = response.content_length(); + let remote_addr = response.remote_addr().map(SocketAddr); + let local_addr = response.local_addr().map(SocketAddr); + let (parts, body) = HttpResponse::from(response).into_parts(); + let (body, held) = if prefetch_body { + match prefetch(body, content_length) { + Prefetch::Complete(bytes) => (Body::Buffered(bytes), true), + Prefetch::Remaining { body, held } => (Body::Unread(body), held), + } + } else { + (Body::Unread(body), false) + }; + let loop_polls = held || loop_polls(parts.version); + Response { + uri, + parts, + body: Arc::new(Mutex::new(body)), + runtime, + content_length, + loop_polls, + remote_addr, + local_addr, + } + } + + /// The largest unread body the event loop collects, `None` if it must not poll it. + #[inline] + fn loop_limit(&self, limit: u64) -> Option { + self.loop_polls.then_some(limit) + } + + #[inline] + fn slot(&self) -> MutexGuard<'_, Body> { + lock(&self.body) + } + + /// Builds a [`wreq::Response`] from the current response metadata and the given body. + #[inline] + fn build_response>(&self, body: T) -> wreq::Response { + let response = HttpResponse::from_parts(self.parts.clone(), body); + wreq::Response::from(response) + } + + /// Take the body for reading. Prefetched and cached bytes are shared at once; another + /// body of known length up to `limit` is polled on the calling thread. Overlapping reads + /// and `stream()` fail with [`Error::Memory`] while a read runs. + fn take_bytes(&self, limit: Option) -> Result { + let mut slot = self.slot(); + let body = match mem::replace(&mut *slot, Body::Taken) { + Body::Unread(body) => body, + Body::Buffered(bytes) | Body::Cached(bytes) => { + *slot = Body::Cached(bytes.clone()); + return Ok(BodyRead::Ready(bytes)); + } + other => { + *slot = other; + return Err(Error::Memory); + } + }; + drop(slot); + let attached = body + .size_hint() + .exact() + .zip(limit) + .is_some_and(|(len, limit)| len <= limit); + let mut collect = body.collect(); + let ready = if attached { + self.runtime.poll_now(&mut collect) + } else { + None + }; + match ready { + Some(Ok(collected)) => { + let bytes = collected.to_bytes(); + cache(&self.body, &bytes); + Ok(BodyRead::Ready(bytes)) + } + Some(Err(err)) => { + self.forbid_recycle(); + Err(Error::Library(err)) + } + None => Ok(BodyRead::Pending(collect)), + } + } + + /// Collect a taken body, caching its bytes and returning them with a copy of the head. + /// The connection stays out of the pool unless the body is read in full, including + /// when the future is dropped before its first poll. + fn collect_later( + &self, + collect: Collect, + ) -> impl Future> + Send + 'static { + let mut guard = RecycleGuard(Some(self.parts.clone())); + let body = self.body.clone(); + async move { + let bytes = collect + .await + .map(Collected::to_bytes) + .map_err(Error::Library)?; + let parts = guard.0.take().ok_or(Error::Memory)?; + cache(&body, &bytes); + Ok((parts, bytes)) + } + } + + /// Take the body and return a future decoding it with `read`, and whether that future + /// finishes at once: the bytes, cached or already buffered, are at most `limit`. Otherwise + /// the body is still arriving or too large to decode on the caller. + fn read_body( + &self, + limit: Option, + read: F, + ) -> Result<(impl Future> + Send + 'static, bool), Error> + where + F: FnOnce(wreq::Response) -> Fut + Send + 'static, + Fut: Future> + Send + 'static, + { + let (decode, inline) = match self.take_bytes(limit)? { + BodyRead::Ready(bytes) => { + let inline = limit.is_some_and(|limit| bytes.len() as u64 <= limit); + (Either::Left(read(self.build_response(bytes))), inline) + } + BodyRead::Pending(collect) => ( + Either::Right(self.collect_later(collect).and_then(|(parts, bytes)| { + read(wreq::Response::from(HttpResponse::from_parts(parts, bytes))) + })), + false, + ), + }; + Ok((decode.map_err(Into::into), inline)) + } + + /// Read and decode the body; the body is taken on first poll. Bytes that + /// [`read_body`](Self::read_body) can decode at once are decoded on the event loop. + async fn read(slf: Py, limit: u64, read: F) -> PyResult + where + F: FnOnce(wreq::Response) -> Fut + Send + 'static, + Fut: Future> + Send + 'static, + T: Send + 'static, + { + let this = slf.get(); + let (read, inline) = this.read_body(this.loop_limit(limit), read)?; + if inline { + read.await + } else { + this.runtime.clone().run_task(read).await + } + } + + /// The first step of an awaited [`read`](Self::read): bytes it decodes at once are + /// returned without suspending. + fn start_read(slf: Bound<'_, PyAny>, limit: u64, read: F) -> PyResult + where + F: FnOnce(wreq::Response) -> Fut + Copy + Send + 'static, + Fut: Future> + Send + 'static, + T: for<'a> IntoPyObject<'a> + Send + 'static, + { + let slf = slf.cast_into::()?; + let this = slf.get(); + let (decode, inline) = this.read_body(this.loop_limit(limit), read)?; + if inline { + // Buffered bytes decode in one poll; otherwise the read decodes the cached bytes. + return match decode.now_or_never() { + Some(value) => Start::ready(slf.py(), value?), + None => Ok(Start::pending(Self::read(slf.unbind(), limit, read))), + }; + } + Ok(Self::continue_on(slf, decode)) + } + + /// Finish a read on the runtime, holding the response until it ends. + fn continue_on(slf: Bound<'_, Self>, read: F) -> Start + where + F: Future> + Send + 'static, + T: for<'a> IntoPyObject<'a> + Send + 'static, + { + let runtime = slf.get().runtime.clone(); + let slf = slf.unbind(); + Start::pending(async move { + let _response = slf; + runtime.run_task(read).await + }) + } + + /// Keep the connection out of the pool; a rebuilt response shares its reuse flag. + fn forbid_recycle(&self) { + let mut response = HttpResponse::new(Bytes::new()); + *response.extensions_mut() = self.parts.extensions.clone(); + wreq::Response::from(response).forbid_recycle(); + } + + /// Drop the body if still held, so later reads fail. Unlike `close`, this keeps a fully + /// read connection reusable. + #[inline] + fn destroy(&self) { + let body = mem::replace(&mut *self.slot(), Body::Released); + drop(body); + } + + /// Discard the body and keep its connection out of the pool. + fn discard(&self) { + self.forbid_recycle(); + self.destroy(); + } +} + +/// Lock a body slot, recovering it from a panicked holder. +#[inline] +fn lock(body: &Mutex) -> MutexGuard<'_, Body> { + body.lock().unwrap_or_else(PoisonError::into_inner) +} + +/// Cache the bytes of a finished read; a release during the read wins over caching. +fn cache(body: &Mutex, bytes: &Bytes) { + let mut slot = lock(body); + if let Body::Taken = *slot { + *slot = Body::Cached(bytes.clone()); + } +} + +/// Decode text using the requested encoding or the response charset. +async fn decode_text( + response: wreq::Response, + encoding: Option, +) -> Result { + match encoding { + Some(encoding) => response.text_with_charset(encoding).await, + None => response.text().await, + } + .map_err(Error::Library) +} + +/// Decode a JSON body using the same conversion as request JSON values. +async fn decode_json(response: wreq::Response) -> Result { + response.json::().await.map_err(Error::Library) +} + +#[pymethods] +impl Response { + /// Get the URL of the response. + #[getter] + pub fn url(&self) -> String { + self.uri.to_string() + } + + /// Get the status code of the response. + #[getter] + pub fn status(&self) -> StatusCode { + StatusCode(self.parts.status) + } + + /// Get the HTTP version of the response. + #[getter] + pub fn version(&self) -> Version { + Version::from_ffi(self.parts.version) + } + + /// Get the headers of the response. + #[getter] + pub fn headers(&self) -> HeaderMap { + HeaderMap(self.parts.headers.clone()) + } + + /// Get the cookies of the response. + #[getter] + pub fn cookies(&self) -> Vec { + Cookie::extract_headers_cookies(&self.parts.headers) + } + + /// Get the content length of the response. + #[getter] + pub fn content_length(&self) -> Option { + self.content_length + } + + /// Get the remote address of the response. + #[getter] + pub fn remote_addr(&self) -> Option { + self.remote_addr + } + + /// Get the local address of the response. + #[getter] + pub fn local_addr(&self) -> Option { + self.local_addr + } + + /// Get the redirect history of the Response. + #[getter] + pub fn history(&self) -> Vec { + self.parts + .extensions + .get::() + .map_or_else(Vec::new, |history| { + history.into_iter().cloned().map(History).collect() + }) + } + + /// Get the TLS information of the response. + #[getter] + pub fn tls_info(&self) -> Option { + self.parts + .extensions + .get::() + .cloned() + .map(TlsInfo) + } + + /// Turn a response into an error if the server returned an error. + pub fn raise_for_status(&self) -> PyResult<()> { + let status = self.parts.status; + if !status.is_client_error() && !status.is_server_error() { + return Ok(()); + } + self.build_response(Bytes::new()) + .error_for_status() + .map(|_| ()) + .map_err(Error::Library) + .map_err(Into::into) + } + + /// Stream read-only memoryviews and any trailing headers from the body. Only an unread + /// body can be streamed; cached bytes stay readable. + pub fn stream(&self) -> PyResult { + let mut slot = self.slot(); + let body = match mem::replace(&mut *slot, Body::Taken) { + Body::Unread(body) => body, + Body::Buffered(bytes) => { + drop(slot); + return Ok(Streamer::buffered(bytes, self.runtime.clone())); + } + other => { + *slot = other; + return Err(Error::Memory.into()); + } + }; + drop(slot); + Ok(Streamer::new( + self.build_response(body), + self.runtime.clone(), + self.loop_polls, + )) + } + + /// Get the text content with the response encoding, defaulting to utf-8 when unspecified. + #[pyo3(signature = (encoding = None))] + pub fn text( + slf: Bound<'_, Self>, + encoding: Option, + ) -> PyResult> { + const QUALNAME: &str = "Response.text"; + let Some(encoding) = encoding else { + return coroutine::eager(QUALNAME, slf, |slf| { + Self::start_read(slf, READ_ATTACHED, |resp| decode_text(resp, None)) + }); + }; + let py = slf.py(); + let read = Self::read(slf.unbind(), READ_ATTACHED, |resp| { + decode_text(resp, Some(encoding)) + }); + coroutine::local(py, QUALNAME, read) + } + + /// Get the JSON content of the response. + pub fn json(slf: Bound<'_, Self>) -> PyResult> { + // Buffered JSON up to 8 KiB is parsed on the event loop; see `loop_limit`. + coroutine::eager("Response.json", slf, |slf| { + Self::start_read(slf, 8 * 1024, decode_json) + }) + } + + /// Read the body as a read-only memoryview, retaining its data after the response closes. + pub fn bytes(slf: Bound<'_, Self>) -> PyResult> { + coroutine::eager("Response.bytes", slf, |slf| { + let slf = slf.cast_into::()?; + let this = slf.get(); + let collect = match this.take_bytes(this.loop_limit(READ_ATTACHED))? { + BodyRead::Ready(bytes) => return Start::ready(slf.py(), PyBuffer::from(bytes)), + BodyRead::Pending(collect) => collect, + }; + let read = this + .collect_later(collect) + .map_ok(|(_, bytes)| PyBuffer::from(bytes)) + .map_err(Into::into); + Ok(Self::continue_on(slf, read)) + }) + } + + /// Discard the retained body and mark its connection as non-reusable. + /// This does not guarantee an immediate socket shutdown or cancel an active read. + /// Cancel and await any body-read task before closing. A body transferred to a + /// Streamer is managed separately; previously returned memoryviews remain valid. + /// `async with` instead releases the body and keeps a fully read connection reusable. + pub fn close(slf: Bound<'_, Self>) -> PyResult> { + coroutine::eager("Response.close", slf, |slf| { + slf.cast::()?.get().discard(); + Start::ready(slf.py(), ()) + }) + } +} + +#[pymethods] +impl Response { + fn __aenter__(slf: Bound<'_, Self>) -> PyResult> { + coroutine::ready("Response.__aenter__", slf) + } + + /// Release the body without forbidding reuse: a fully read connection returns + /// to the pool, while an unread HTTP/1 body drains or closes its connection. + fn __aexit__<'py>( + slf: Bound<'py, Self>, + _exc_type: Py, + _exc_val: Py, + _traceback: Py, + ) -> PyResult> { + coroutine::eager("Response.__aexit__", slf, |slf| { + slf.cast::()?.get().destroy(); + Start::ready(slf.py(), ()) + }) + } +} + +impl EntersSelf for Response { + fn native_aenter() -> &'static PyOnceLock> { + static NATIVE: PyOnceLock> = PyOnceLock::new(); + &NATIVE + } +} + +impl Display for Response { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + f, + "<{}({}) [{}] >", + stringify!(Response), + self.uri, + self.parts.status, + ) + } +} + +impl Drop for Response { + #[inline] + fn drop(&mut self) { + self.destroy(); + } +} + +// ===== impl RecycleGuard ===== + +impl Drop for RecycleGuard { + fn drop(&mut self) { + if let Some(parts) = self.0.take() { + wreq::Response::from(HttpResponse::from_parts(parts, Bytes::new())).forbid_recycle(); + } + } +} + +// ===== impl BlockingResponse ===== + +impl BlockingResponse { + /// Read the body with `read` on the calling thread. Bytes that + /// [`Response::read_body`] can decode at once are decoded without releasing the GIL. + fn read(&self, py: Python, read: F) -> PyResult + where + F: FnOnce(wreq::Response) -> Fut + Send + 'static, + 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 { + runtime.block_on_eager(py, read) + } else { + runtime.block_on(py, read) + } + } +} + +#[pymethods] +impl BlockingResponse { + /// Get the URL of the response. + #[getter] + pub fn url(&self) -> String { + self.0.url() + } + + /// Get the status code of the response. + #[getter] + pub fn status(&self) -> StatusCode { + self.0.status() + } + + /// Get the HTTP version of the response. + #[getter] + pub fn version(&self) -> Version { + self.0.version() + } + + /// Get the headers of the response. + #[getter] + pub fn headers(&self) -> HeaderMap { + self.0.headers() + } + + /// Get the cookies of the response. + #[getter] + pub fn cookies(&self) -> Vec { + self.0.cookies() + } + + /// Get the content length of the response. + #[getter] + pub fn content_length(&self) -> Option { + self.0.content_length() + } + + /// Get the remote address of the response. + #[getter] + pub fn remote_addr(&self) -> Option { + self.0.remote_addr() + } + + /// Get the local address of the response. + #[getter] + pub fn local_addr(&self) -> Option { + self.0.local_addr() + } + + /// Get the redirect history of the Response. + #[getter] + pub fn history(&self) -> Vec { + self.0.history() + } + + /// Get the TLS information of the response. + #[getter] + pub fn tls_info(&self) -> Option { + self.0.tls_info() + } + + /// Turn a response into an error if the server returned an error. + pub fn raise_for_status(&self) -> PyResult<()> { + self.0.raise_for_status() + } + + /// Stream read-only memoryviews and any trailing headers from the body. + pub fn stream(&self) -> PyResult { + self.0.stream() + } + + /// Get the text content with the response encoding, defaulting to utf-8 when unspecified. + #[pyo3(signature = (encoding = None))] + pub fn text(&self, py: Python, encoding: Option) -> PyResult { + self.read(py, |resp| decode_text(resp, encoding)) + } + + /// Get the JSON content of the response. + pub fn json(&self, py: Python) -> PyResult { + self.read(py, decode_json) + } + + /// 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)), + BodyRead::Pending(collect) => { + let read = response + .collect_later(collect) + .map_ok(|(_, bytes)| PyBuffer::from(bytes)) + .map_err(Into::into); + response.runtime.block_on(py, read) + } + } + } + + /// Discard the retained body and mark its connection as non-reusable. + /// This does not guarantee an immediate socket shutdown or interrupt an active read. + /// Do not close concurrently with a body read. A body transferred to a Streamer + /// is managed separately; previously returned memoryviews remain valid. + /// `with` instead releases the body and keeps a fully read connection reusable. + pub fn close(&self) { + self.0.discard(); + } +} + +#[pymethods] +impl BlockingResponse { + fn __enter__(slf: PyRef) -> PyRef { + slf + } + + /// Release the body without forbidding reuse: a fully read connection returns + /// to the pool, while an unread HTTP/1 body drains or closes its connection. + fn __exit__<'py>( + &self, + _exc_type: &Bound<'py, PyAny>, + _exc_value: &Bound<'py, PyAny>, + _traceback: &Bound<'py, PyAny>, + ) { + self.0.destroy(); + } +} + +impl From for BlockingResponse { + #[inline] + fn from(response: Response) -> Self { + Self(response) + } +} + +impl Display for BlockingResponse { + #[inline] + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + self.0.fmt(f) + } } diff --git a/src/client/resp/ext.rs b/src/client/resp/ext.rs deleted file mode 100644 index f516422a..00000000 --- a/src/client/resp/ext.rs +++ /dev/null @@ -1,30 +0,0 @@ -use pyo3::pybacked::PyBackedStr; -use serde::de::DeserializeOwned; - -use crate::error::Error; - -/// Body readers for [`wreq::Response`] that return crate [`Error`]s, shared by the sync -/// and async responses. -pub trait ResponseExt { - /// Decode the body with `encoding`, or the response charset defaulting to UTF-8. - async fn text(self, encoding: Option) -> Result; - - /// Deserialize the body as JSON. - async fn json(self) -> Result; -} - -impl ResponseExt for wreq::Response { - #[inline] - async fn text(self, encoding: Option) -> Result { - match encoding { - Some(encoding) => self.text_with_charset(encoding).await, - None => self.text().await, - } - .map_err(Error::Library) - } - - #[inline] - async fn json(self) -> Result { - self.json::().await.map_err(Error::Library) - } -} diff --git a/src/client/resp/http.rs b/src/client/resp/http.rs deleted file mode 100644 index 1796f8e5..00000000 --- a/src/client/resp/http.rs +++ /dev/null @@ -1,640 +0,0 @@ -use std::{ - fmt::{self, Display}, - future::Future, - mem, - sync::{Arc, Mutex, MutexGuard, PoisonError}, -}; - -use bytes::Bytes; -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}; -use pyo3::{prelude::*, pybacked::PyBackedStr, sync::PyOnceLock}; -use wreq::Uri; - -use super::{READ_ATTACHED, ext::ResponseExt, loop_limit, stream::Streamer}; -use crate::{ - buffer::PyBuffer, - client::{SocketAddr, body::Json}, - cookie::Cookie, - coroutine::{self, Coroutine, EntersSelf}, - error::Error, - header::HeaderMap, - http::{StatusCode, Version}, - redirect::History, - runtime::Runtime, - tls::TlsInfo, -}; - -/// A response from a request. -/// -/// Reads take the body with [`Response::take_bytes`]: a buffered body of known length up to -/// the read's limit finishes on the caller, the event loop or a blocking thread, and any other -/// body is collected on the runtime by [`Response::collect_later`]. -#[pyclass(subclass, frozen, str)] -pub struct Response { - uri: Uri, - /// Response head; rebuilt responses share its extensions, including the connection's - /// reuse flag. - parts: Parts, - /// Shared with a read in flight, so it can cache the bytes it reads. - body: Arc>, - /// Runs body reads and keeps the runtime alive while the response does. - runtime: Runtime, - /// Captured at receipt; `None` for a body without a known length. - content_length: Option, - remote_addr: Option, - local_addr: Option, -} - -/// The response body slot. -enum Body { - /// Unread; taken by the first read or `stream()`. - Unread(wreq::Body), - /// Taken by a read in flight or by `stream()`. - Taken, - /// Read in full and shared by later `text`, `json` and `bytes` calls. - Cached(Bytes), - /// Released by `close` or a context exit; later reads fail. - Released, -} - -/// A body taken for reading. -enum BodyRead { - /// Read in full now. - Ready(Bytes), - /// Not read now: still arriving, over the read's limit or of unknown length; collected - /// on the runtime by [`Response::collect_later`]. - Pending(Collect), -} - -/// A blocking response from a request. -#[pyclass(name = "Response", subclass, frozen, str)] -pub struct BlockingResponse(Response); - -/// Forbids connection reuse on drop unless disarmed by taking the parts. Held by -/// [`Response::collect_later`] from its creation until the body is read in full, so a failed -/// or cancelled read is not pooled. -struct RecycleGuard(Option); - -// ===== impl Response ===== - -impl Response { - /// Create a new [`Response`] instance. - pub fn new(response: wreq::Response, runtime: Runtime) -> Self { - let uri = response.uri().clone(); - let content_length = response.content_length(); - let remote_addr = response.remote_addr().map(SocketAddr); - let local_addr = response.local_addr().map(SocketAddr); - let (parts, body) = HttpResponse::from(response).into_parts(); - Response { - uri, - parts, - body: Arc::new(Mutex::new(Body::Unread(body))), - runtime, - content_length, - remote_addr, - local_addr, - } - } - - #[inline] - fn slot(&self) -> MutexGuard<'_, Body> { - lock(&self.body) - } - - /// Builds a [`wreq::Response`] from the current response metadata and the given body. - #[inline] - fn build_response>(&self, body: T) -> wreq::Response { - let response = HttpResponse::from_parts(self.parts.clone(), body); - wreq::Response::from(response) - } - - /// Take the body for reading. Cached bytes are shared at once, and a body of known - /// length up to `limit` that is already buffered is read now on the calling thread, none - /// without a `limit`; overlapping reads and `stream()` fail with [`Error::Memory`] while - /// a read runs. - fn take_bytes(&self, limit: Option) -> Result { - let mut slot = self.slot(); - let body = match mem::replace(&mut *slot, Body::Taken) { - Body::Unread(body) => body, - Body::Cached(bytes) => { - *slot = Body::Cached(bytes.clone()); - return Ok(BodyRead::Ready(bytes)); - } - other => { - *slot = other; - return Err(Error::Memory); - } - }; - drop(slot); - let attached = body - .size_hint() - .exact() - .zip(limit) - .is_some_and(|(len, limit)| len <= limit); - let mut collect = body.collect(); - let ready = if attached { - self.runtime.poll_now(&mut collect) - } else { - None - }; - match ready { - Some(Ok(collected)) => { - let bytes = collected.to_bytes(); - cache(&self.body, &bytes); - Ok(BodyRead::Ready(bytes)) - } - Some(Err(err)) => { - self.forbid_recycle(); - Err(Error::Library(err)) - } - None => Ok(BodyRead::Pending(collect)), - } - } - - /// Collect a taken body, caching its bytes and returning them with a copy of the head. - /// The connection stays out of the pool unless the body is read in full, including - /// when the future is dropped before its first poll. - fn collect_later( - &self, - collect: Collect, - ) -> impl Future> + Send + 'static { - let mut guard = RecycleGuard(Some(self.parts.clone())); - let body = self.body.clone(); - async move { - let bytes = collect - .await - .map(Collected::to_bytes) - .map_err(Error::Library)?; - let parts = guard.0.take().ok_or(Error::Memory)?; - cache(&body, &bytes); - Ok((parts, bytes)) - } - } - - /// Take the body and return a future decoding it with `read`, and whether that future - /// finishes at once: the bytes, cached or already buffered, are at most `limit`. Otherwise - /// the body is still arriving or too large to decode on the caller. - fn read_body( - &self, - limit: Option, - read: F, - ) -> Result<(impl Future> + Send + 'static, bool), Error> - where - F: FnOnce(wreq::Response) -> Fut + Send + 'static, - Fut: Future> + Send + 'static, - { - let (decode, inline) = match self.take_bytes(limit)? { - BodyRead::Ready(bytes) => { - let inline = limit.is_some_and(|limit| bytes.len() as u64 <= limit); - (Either::Left(read(self.build_response(bytes))), inline) - } - BodyRead::Pending(collect) => ( - Either::Right(self.collect_later(collect).and_then(|(parts, bytes)| { - read(wreq::Response::from(HttpResponse::from_parts(parts, bytes))) - })), - false, - ), - }; - Ok((decode.map_err(Into::into), inline)) - } - - /// Read and decode the body once awaited; the body is taken on first await. Bytes that - /// [`read_body`](Self::read_body) can decode at once are decoded on the event loop. - fn read<'py, F, Fut, T>( - slf: Bound<'py, Self>, - qualname: &'static str, - limit: u64, - read: F, - ) -> PyResult> - where - F: FnOnce(wreq::Response) -> Fut + Send + 'static, - Fut: Future> + Send + 'static, - T: for<'a> IntoPyObject<'a> + Send + 'static, - { - let py = slf.py(); - let slf = slf.unbind(); - coroutine::local(py, qualname, async move { - let this = slf.get(); - let (read, inline) = this.read_body(loop_limit(this.parts.version, limit), read)?; - if inline { - read.await - } else { - this.runtime.clone().run_task(read).await - } - }) - } - - /// Keep the connection out of the pool; a rebuilt response shares its reuse flag. - fn forbid_recycle(&self) { - let mut response = HttpResponse::new(Bytes::new()); - *response.extensions_mut() = self.parts.extensions.clone(); - wreq::Response::from(response).forbid_recycle(); - } - - /// Drop the body if still held, so later reads fail. Unlike `close`, this keeps a fully - /// read connection reusable. - #[inline] - fn destroy(&self) { - let body = mem::replace(&mut *self.slot(), Body::Released); - drop(body); - } - - /// Discard the body and keep its connection out of the pool. - fn discard(&self) { - self.forbid_recycle(); - self.destroy(); - } -} - -/// Lock a body slot, recovering it from a panicked holder. -#[inline] -fn lock(body: &Mutex) -> MutexGuard<'_, Body> { - body.lock().unwrap_or_else(PoisonError::into_inner) -} - -/// Cache the bytes of a finished read; a release during the read wins over caching. -fn cache(body: &Mutex, bytes: &Bytes) { - let mut slot = lock(body); - if let Body::Taken = *slot { - *slot = Body::Cached(bytes.clone()); - } -} - -#[pymethods] -impl Response { - /// Get the URL of the response. - #[getter] - pub fn url(&self) -> String { - self.uri.to_string() - } - - /// Get the status code of the response. - #[getter] - pub fn status(&self) -> StatusCode { - StatusCode(self.parts.status) - } - - /// Get the HTTP version of the response. - #[getter] - pub fn version(&self) -> Version { - Version::from_ffi(self.parts.version) - } - - /// Get the headers of the response. - #[getter] - pub fn headers(&self) -> HeaderMap { - HeaderMap(self.parts.headers.clone()) - } - - /// Get the cookies of the response. - #[getter] - pub fn cookies(&self) -> Vec { - Cookie::extract_headers_cookies(&self.parts.headers) - } - - /// Get the content length of the response. - #[getter] - pub fn content_length(&self) -> Option { - self.content_length - } - - /// Get the remote address of the response. - #[getter] - pub fn remote_addr(&self) -> Option { - self.remote_addr - } - - /// Get the local address of the response. - #[getter] - pub fn local_addr(&self) -> Option { - self.local_addr - } - - /// Get the redirect history of the Response. - #[getter] - pub fn history(&self) -> Vec { - self.parts - .extensions - .get::() - .map_or_else(Vec::new, |history| { - history.into_iter().cloned().map(History).collect() - }) - } - - /// Get the TLS information of the response. - #[getter] - pub fn tls_info(&self) -> Option { - self.parts - .extensions - .get::() - .cloned() - .map(TlsInfo) - } - - /// Turn a response into an error if the server returned an error. - pub fn raise_for_status(&self) -> PyResult<()> { - let status = self.parts.status; - if !status.is_client_error() && !status.is_server_error() { - return Ok(()); - } - self.build_response(Bytes::new()) - .error_for_status() - .map(|_| ()) - .map_err(Error::Library) - .map_err(Into::into) - } - - /// Stream read-only memoryviews and any trailing headers from the body. Only an unread - /// body can be streamed; cached bytes stay readable. - pub fn stream(&self) -> PyResult { - let mut slot = self.slot(); - let body = match mem::replace(&mut *slot, Body::Taken) { - Body::Unread(body) => body, - other => { - *slot = other; - return Err(Error::Memory.into()); - } - }; - drop(slot); - Ok(Streamer::new( - self.build_response(body), - self.runtime.clone(), - )) - } - - /// Get the text content with the response encoding, defaulting to utf-8 when unspecified. - #[pyo3(signature = (encoding = None))] - pub fn text( - slf: Bound<'_, Self>, - encoding: Option, - ) -> PyResult> { - Self::read(slf, "Response.text", READ_ATTACHED, |resp| { - ResponseExt::text(resp, encoding) - }) - } - - /// Get the JSON content of the response. - pub fn json(slf: Bound<'_, Self>) -> PyResult> { - // Buffered HTTP/1 JSON up to 8 KiB is parsed on the event loop; see `loop_limit`. - Self::read(slf, "Response.json", 8 * 1024, ResponseExt::json::) - } - - /// Read the body as a read-only memoryview, retaining its data after the response closes. - pub fn bytes(slf: Bound<'_, Self>) -> PyResult> { - let py = slf.py(); - let slf = slf.unbind(); - coroutine::local(py, "Response.bytes", async move { - let this = slf.get(); - match this.take_bytes(loop_limit(this.parts.version, READ_ATTACHED))? { - BodyRead::Ready(bytes) => Ok(PyBuffer::from(bytes)), - BodyRead::Pending(collect) => { - let read = this - .collect_later(collect) - .map_ok(|(_, bytes)| PyBuffer::from(bytes)) - .map_err(Into::into); - this.runtime.clone().run_task(read).await - } - } - }) - } - - /// Discard the retained body and mark its connection as non-reusable. - /// This does not guarantee an immediate socket shutdown or cancel an active read. - /// Cancel and await any body-read task before closing. A body transferred to a - /// Streamer is managed separately; previously returned memoryviews remain valid. - /// `async with` instead releases the body and keeps a fully read connection reusable. - pub fn close(slf: Bound<'_, Self>) -> PyResult> { - let py = slf.py(); - let slf = slf.unbind(); - coroutine::local(py, "Response.close", async move { - slf.get().discard(); - Ok(()) - }) - } -} - -#[pymethods] -impl Response { - fn __aenter__(slf: Bound<'_, Self>) -> PyResult> { - coroutine::ready("Response.__aenter__", slf) - } - - /// Release the body without forbidding reuse: a fully read connection returns - /// to the pool, while an unread HTTP/1 body drains or closes its connection. - fn __aexit__<'py>( - slf: Bound<'py, Self>, - _exc_type: Py, - _exc_val: Py, - _traceback: Py, - ) -> PyResult> { - let py = slf.py(); - let slf = slf.unbind(); - coroutine::local(py, "Response.__aexit__", async move { - slf.get().destroy(); - Ok(()) - }) - } -} - -impl EntersSelf for Response { - fn native_aenter() -> &'static PyOnceLock> { - static NATIVE: PyOnceLock> = PyOnceLock::new(); - &NATIVE - } -} - -impl Display for Response { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - write!( - f, - "<{}({}) [{}] >", - stringify!(Response), - self.uri, - self.parts.status, - ) - } -} - -impl Drop for Response { - #[inline] - fn drop(&mut self) { - self.destroy(); - } -} - -// ===== impl RecycleGuard ===== - -impl Drop for RecycleGuard { - fn drop(&mut self) { - if let Some(parts) = self.0.take() { - wreq::Response::from(HttpResponse::from_parts(parts, Bytes::new())).forbid_recycle(); - } - } -} - -// ===== impl BlockingResponse ===== - -impl BlockingResponse { - /// Read the body with `read` on the calling thread. Bytes that - /// [`Response::read_body`] can decode at once are decoded without releasing the GIL. - fn read(&self, py: Python, read: F) -> PyResult - where - F: FnOnce(wreq::Response) -> Fut + Send + 'static, - 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 { - runtime.block_on_eager(py, read) - } else { - runtime.block_on(py, read) - } - } -} - -#[pymethods] -impl BlockingResponse { - /// Get the URL of the response. - #[getter] - pub fn url(&self) -> String { - self.0.url() - } - - /// Get the status code of the response. - #[getter] - pub fn status(&self) -> StatusCode { - self.0.status() - } - - /// Get the HTTP version of the response. - #[getter] - pub fn version(&self) -> Version { - self.0.version() - } - - /// Get the headers of the response. - #[getter] - pub fn headers(&self) -> HeaderMap { - self.0.headers() - } - - /// Get the cookies of the response. - #[getter] - pub fn cookies(&self) -> Vec { - self.0.cookies() - } - - /// Get the content length of the response. - #[getter] - pub fn content_length(&self) -> Option { - self.0.content_length() - } - - /// Get the remote address of the response. - #[getter] - pub fn remote_addr(&self) -> Option { - self.0.remote_addr() - } - - /// Get the local address of the response. - #[getter] - pub fn local_addr(&self) -> Option { - self.0.local_addr() - } - - /// Get the redirect history of the Response. - #[getter] - pub fn history(&self) -> Vec { - self.0.history() - } - - /// Get the TLS information of the response. - #[getter] - pub fn tls_info(&self) -> Option { - self.0.tls_info() - } - - /// Turn a response into an error if the server returned an error. - pub fn raise_for_status(&self) -> PyResult<()> { - self.0.raise_for_status() - } - - /// Stream read-only memoryviews and any trailing headers from the body. - pub fn stream(&self) -> PyResult { - self.0.stream() - } - - /// Get the text content with the response encoding, defaulting to utf-8 when unspecified. - #[pyo3(signature = (encoding = None))] - pub fn text(&self, py: Python, encoding: Option) -> PyResult { - self.read(py, |resp| ResponseExt::text(resp, encoding)) - } - - /// Get the JSON content of the response. - pub fn json(&self, py: Python) -> PyResult { - self.read(py, ResponseExt::json::) - } - - /// 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)), - BodyRead::Pending(collect) => { - let read = response - .collect_later(collect) - .map_ok(|(_, bytes)| PyBuffer::from(bytes)) - .map_err(Into::into); - response.runtime.block_on(py, read) - } - } - } - - /// Discard the retained body and mark its connection as non-reusable. - /// This does not guarantee an immediate socket shutdown or interrupt an active read. - /// Do not close concurrently with a body read. A body transferred to a Streamer - /// is managed separately; previously returned memoryviews remain valid. - /// `with` instead releases the body and keeps a fully read connection reusable. - pub fn close(&self) { - self.0.discard(); - } -} - -#[pymethods] -impl BlockingResponse { - fn __enter__(slf: PyRef) -> PyRef { - slf - } - - /// Release the body without forbidding reuse: a fully read connection returns - /// to the pool, while an unread HTTP/1 body drains or closes its connection. - fn __exit__<'py>( - &self, - _exc_type: &Bound<'py, PyAny>, - _exc_value: &Bound<'py, PyAny>, - _traceback: &Bound<'py, PyAny>, - ) { - self.0.destroy(); - } -} - -impl From for BlockingResponse { - #[inline] - fn from(response: Response) -> Self { - Self(response) - } -} - -impl Display for BlockingResponse { - #[inline] - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - self.0.fmt(f) - } -} diff --git a/src/client/resp/prefetch.rs b/src/client/resp/prefetch.rs new file mode 100644 index 00000000..24e98d2c --- /dev/null +++ b/src/client/resp/prefetch.rs @@ -0,0 +1,152 @@ +//! Short response bodies read along with the head. + +use std::{ + pin::Pin, + task::{Context, Poll}, +}; + +use bytes::{Bytes, BytesMut}; +use futures_util::FutureExt; +use http_body::{Frame, SizeHint}; +use http_body_util::BodyExt; + +/// Bodies of known length up to this size are read with the head when already received. +const PREFETCH_LIMIT: u64 = 64 * 1024; + +/// A completed body, or frames still to deliver, safe to poll off-runtime only if held. +pub(super) enum Prefetch { + Complete(Bytes), + Remaining { body: wreq::Body, held: bool }, +} + +/// A body whose leading frames were read with the head, replayed before the rest. +struct Prefetched { + data: Option, + /// The frame that ended prefetching: trailers or an error. + next: Option>>, + /// The unread rest, `None` once the body failed. + rest: Option, +} + +/// Data frames read so far; the usual single frame is kept without copying. +#[derive(Default)] +enum Data { + #[default] + Empty, + One(Bytes), + Many(BytesMut), +} + +/// Read the frames of a body of known length up to [`PREFETCH_LIMIT`] that already arrived, +/// so readers on the event loop need not wait on the runtime. +/// +/// Completed bodies keep their bytes directly. Other bodies report whether all remaining +/// frames are held; an HTTP/2 body still tied to its connection stays on the runtime. +/// It never waits for the network: frames still in flight stay with the connection. +pub(super) fn prefetch(mut body: wreq::Body, len: Option) -> Prefetch { + let Some(len) = len.filter(|len| *len <= PREFETCH_LIMIT) else { + return Prefetch::Remaining { body, held: false }; + }; + let mut data = Data::default(); + let next = loop { + match body.frame().now_or_never() { + Some(Some(Ok(frame))) => match frame.into_data() { + Ok(bytes) => { + data.push(bytes, len); + } + Err(frame) => break Some(Ok(frame)), + }, + Some(Some(Err(err))) => break Some(Err(err)), + Some(None) => { + return Prefetch::Complete(data.into_bytes().unwrap_or_default()); + } + None => break None, + } + }; + let data = data.into_bytes(); + if data.is_none() && next.is_none() { + return Prefetch::Remaining { body, held: false }; + } + // A failed body holds nothing more; the error replays after the data. + let rest = (!matches!(next, Some(Err(_)))).then_some(body); + let held = rest.is_none(); + Prefetch::Remaining { + body: wreq::Body::wrap(Prefetched { data, next, rest }), + held, + } +} + +// ===== impl Data ===== + +impl Data { + fn push(&mut self, bytes: Bytes, len: u64) { + // An empty frame, such as HTTP/2's trailing END_STREAM, adds nothing worth a copy. + if bytes.is_empty() { + return; + } + match self { + Data::Empty => *self = Data::One(bytes), + Data::One(first) => { + let total = first.len() + bytes.len(); + let len = usize::try_from(len).unwrap_or(total); + let mut buf = BytesMut::with_capacity(total.max(len)); + buf.extend_from_slice(first); + buf.extend_from_slice(&bytes); + *self = Data::Many(buf); + } + Data::Many(buf) => buf.extend_from_slice(&bytes), + } + } + + fn into_bytes(self) -> Option { + match self { + Data::Empty => None, + Data::One(bytes) => Some(bytes), + Data::Many(buf) => Some(buf.freeze()), + } + } +} + +// ===== impl Prefetched ===== + +impl http_body::Body for Prefetched { + type Data = Bytes; + type Error = wreq::Error; + + fn poll_frame( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll, wreq::Error>>> { + let this = self.get_mut(); + if let Some(data) = this.data.take() { + return Poll::Ready(Some(Ok(Frame::data(data)))); + } + if let Some(next) = this.next.take() { + return Poll::Ready(Some(next)); + } + match &mut this.rest { + Some(rest) => Pin::new(rest).poll_frame(cx), + None => Poll::Ready(None), + } + } + + fn size_hint(&self) -> SizeHint { + let mut hint = self + .rest + .as_ref() + .map_or_else(|| SizeHint::with_exact(0), wreq::Body::size_hint); + let held = self.data.as_ref().map_or(0, |data| data.len() as u64); + // Raise the upper bound first: `set_lower` panics above it. + if let Some(upper) = hint.upper() { + hint.set_upper(upper.saturating_add(held)); + } + hint.set_lower(hint.lower().saturating_add(held)); + hint + } + + fn is_end_stream(&self) -> bool { + self.data.is_none() + && self.next.is_none() + && self.rest.as_ref().is_none_or(wreq::Body::is_end_stream) + } +} diff --git a/src/client/resp/stream.rs b/src/client/resp/stream.rs index 1b1ca12d..1155e835 100644 --- a/src/client/resp/stream.rs +++ b/src/client/resp/stream.rs @@ -11,7 +11,7 @@ use std::{ }, }; -use bytes::Bytes; +use bytes::{Bytes, BytesMut}; use futures_util::FutureExt; use http_body::Body as _; use http_body_util::BodyExt; @@ -22,10 +22,9 @@ use tokio::sync::{ }; use tokio_util::task::AbortOnDropHandle; -use super::loop_polls; use crate::{ buffer::PyBuffer, - coroutine::{self, Coroutine}, + coroutine::{self, Coroutine, Start}, error::Error, header::HeaderMap, runtime::Runtime, @@ -38,6 +37,9 @@ pub enum Frame { Trailers(HeaderMap), } +/// A body frame as read from the response. +type RawFrame = wreq::Result>; + /// A response stream yielding read-only memoryviews and any trailing headers. /// /// Sync and async iteration share [`Streamer::next`]; [`Iteration`] picks the ready frames @@ -46,6 +48,8 @@ pub enum Frame { pub struct Streamer { reader: Arc, runtime: Runtime, + /// Whether the event loop may poll the body: an HTTP/1 body, or one held in full. + loop_polls: bool, } /// Frames buffered between the read-ahead task and Python readers, which poll the @@ -61,6 +65,8 @@ struct Reader { enum State { /// Unread, or read only by [`Streamer::ready_frame`] on the caller. Idle(Box), + /// A successfully prefetched body, taken once without polling a transport body. + Buffered(Option), /// Held by the read-ahead task, whose frames arrive on `rx`; dropping `_task` aborts it. /// Without a task, `rx` holds only a final error. Reading { @@ -87,25 +93,47 @@ enum Iteration { // ===== impl Streamer ===== impl Streamer { - /// Frames buffered ahead of Python. A frame holds at most one transport read, up to - /// 408 KiB on HTTP/1 by default, so the 8 queued frames plus the one being sent stay - /// under about 3.6 MiB. + /// Frames buffered ahead of Python. A frame holds one transport read, up to 408 KiB on + /// HTTP/1 by default, or data frames merged up to [`Self::MERGE_LIMIT`], so the 8 queued + /// frames plus the one being sent stay under about 3.6 MiB. const READ_AHEAD: usize = 8; + /// Data frames that already arrived are merged up to this size, so a body sent in small + /// frames, such as 16 KiB HTTP/2 DATA frames, reaches Python in fewer chunks without + /// waiting for more data. + const MERGE_LIMIT: usize = 256 * 1024; + + /// Only frames smaller than this are merged; larger ones are sent without a copy. + const MERGE_FRAME: usize = 64 * 1024; + /// Buffered frames returned before `__anext__` yields to the event loop once. const YIELD_EVERY: usize = 8; /// Create a new [`Streamer`] instance. #[inline] - pub fn new(resp: wreq::Response, runtime: Runtime) -> Streamer { + pub fn new(resp: wreq::Response, runtime: Runtime, loop_polls: bool) -> Streamer { + Self::with_state(State::Idle(Box::new(resp)), runtime, loop_polls) + } + + /// A body already read to EOF, with no trailers or errors left to deliver. + pub fn buffered(bytes: Bytes, runtime: Runtime) -> Streamer { + Self::with_state( + State::Buffered((!bytes.is_empty()).then_some(bytes)), + runtime, + true, + ) + } + + fn with_state(state: State, runtime: Runtime, loop_polls: bool) -> Streamer { let reader = Reader { - state: Mutex::new(State::Idle(Box::new(resp))), + state: Mutex::new(state), arrived: Arc::new(Notify::new()), since_yield: AtomicUsize::new(0), }; Streamer { reader: Arc::new(reader), runtime, + loop_polls, } } @@ -133,8 +161,21 @@ 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(Self::next_frame(&mut resp), arrived, &mut unannounced).await - { + // A frame read past merged data, sent next; `Some(None)` is the end of the body. + let mut pending = None; + loop { + let frame = match pending.take() { + Some(frame) => frame, + None => burst(resp.frame(), arrived, &mut unannounced).await, + }; + let Some(frame) = frame else { + break; + }; + let (frame, next) = Self::merge(frame, &mut resp); + pending = next; + let Some(frame) = Self::convert(frame, &mut resp) else { + continue; + }; let failed = frame.is_err(); if burst(tx.send(frame), arrived, &mut unannounced) .await @@ -147,6 +188,52 @@ impl Streamer { } } + /// Merge the data frames `resp` already holds into a data `frame`, up to + /// [`Self::MERGE_LIMIT`]. Returns the merged frame and the frame read past it, if any. + fn merge(frame: RawFrame, resp: &mut wreq::Response) -> (RawFrame, Option>) { + let first = match frame.map(http_body::Frame::into_data) { + Ok(Ok(data)) if data.len() < Self::MERGE_FRAME => data, + Ok(Ok(data)) => return (Ok(http_body::Frame::data(data)), None), + Ok(Err(frame)) => return (Ok(frame), None), + Err(err) => return (Err(err), None), + }; + // Frames are gathered first and copied once into a buffer of their exact size, so a + // chunk holds no spare capacity; a single frame is sent without a copy. + let mut rest = Vec::new(); + let mut len = first.len(); + let next = loop { + if len >= Self::MERGE_LIMIT { + break None; + } + let more = match resp.frame().now_or_never() { + Some(Some(Ok(frame))) => match frame.into_data() { + Ok(more) => more, + Err(frame) => break Some(Some(Ok(frame))), + }, + Some(other) => break Some(other), + None => break None, + }; + // A large frame, or one that would pass the limit, starts the next chunk, so no + // chunk grows beyond the limit or a single frame. + if more.len() >= Self::MERGE_FRAME || len + more.len() > Self::MERGE_LIMIT { + break Some(Some(Ok(http_body::Frame::data(more)))); + } + if !more.is_empty() { + len += more.len(); + rest.push(more); + } + }; + let data = if rest.is_empty() { + first + } else { + let mut buf = BytesMut::with_capacity(len); + buf.extend_from_slice(&first); + rest.iter().for_each(|frame| buf.extend_from_slice(frame)); + buf.freeze() + }; + (Ok(http_body::Frame::data(data)), next) + } + /// Read the next frame, skipping unknown kinds; `None` at the end of the body. async fn next_frame(resp: &mut wreq::Response) -> Option> { loop { @@ -158,10 +245,7 @@ impl Streamer { } /// Convert a body frame, skipping unknown kinds. - fn convert( - frame: wreq::Result>, - resp: &mut wreq::Response, - ) -> Option> { + fn convert(frame: RawFrame, resp: &mut wreq::Response) -> Option> { match frame.map(|frame| frame.into_data()) { Ok(Ok(bytes)) => Some(Ok(Frame::Bytes(PyBuffer::from(bytes)))), Ok(Err(frame)) => frame @@ -180,10 +264,13 @@ impl Streamer { /// of a short body never waits on another thread. fn ready_frame(&self, iteration: Iteration) -> Option> { let mut state = self.reader.lock(); + if let Some(frame) = state.take_buffered(iteration) { + return Some(frame); + } let State::Idle(resp) = &mut *state else { return None; }; - if !iteration.takes_ready(resp) { + if !iteration.takes_ready(resp, self.loop_polls) { return None; } let Some(frame) = self.runtime.poll_now(Self::next_frame(resp))? else { @@ -287,8 +374,13 @@ impl Streamer { } fn __next__(&self, py: Python) -> PyResult { - // Refused before the first poll, which returns frames already received without + // Refused before the first read, which returns frames already received without // releasing the GIL, so a nested read never leaves the stream half consumed. + Runtime::refuse_nested()?; + // A prefetched body is read without entering the runtime. + if let Some(frame) = self.reader.lock().take_buffered(Iteration::Sync) { + return frame; + } self.runtime.block_on_eager(py, self.next(Iteration::Sync)) } @@ -317,17 +409,28 @@ impl Streamer { /// 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 { + coroutine::eager("Streamer.__anext__", slf, |slf| { + let slf = slf.cast_into::()?; let this = slf.get(); // 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 { + let yields = + this.reader.since_yield.fetch_add(1, Ordering::Relaxed) >= Self::YIELD_EVERY; + if yields { this.reader.since_yield.store(0, Ordering::Relaxed); - coroutine::yield_now().await; + } else if let Some(frame) = this + .ready_frame(Iteration::Async) + .or_else(|| this.try_next(Iteration::Async)) + { + return Start::ready(slf.py(), frame?); } - this.next(Iteration::Async).await + let slf = slf.unbind(); + Ok(Start::pending(async move { + if yields { + coroutine::yield_now().await; + } + slf.get().next(Iteration::Async).await + })) }) } @@ -337,17 +440,15 @@ impl Streamer { } /// Release the body and end any pending read; returned views stay valid. - fn __aexit__<'py>( - &self, - py: Python<'py>, + fn __aexit__( + slf: Bound<'_, Self>, _exc_type: Py, _exc_val: Py, _traceback: Py, - ) -> PyResult> { - let reader = self.reader.clone(); - coroutine::local(py, "Streamer.__aexit__", async move { - reader.close(); - Ok(()) + ) -> PyResult> { + coroutine::eager("Streamer.__aexit__", slf, |slf| { + slf.cast::()?.get().reader.close(); + Start::ready(slf.py(), ()) }) } } @@ -383,6 +484,23 @@ impl Reader { } } +// ===== impl State ===== + +impl State { + /// Take a prefetched body's frame, then the end of `iteration`; `None` for any other body. + fn take_buffered(&mut self, iteration: Iteration) -> Option> { + let State::Buffered(bytes) = self else { + return None; + }; + Some( + bytes + .take() + .map(|bytes| Frame::Bytes(PyBuffer::from(bytes))) + .ok_or_else(|| iteration.end().into()), + ) + } +} + // ===== impl NotifyOnDrop ===== impl Drop for NotifyOnDrop { @@ -394,13 +512,13 @@ impl Drop for NotifyOnDrop { // ===== 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 { + /// Whether the caller takes a frame `resp` already holds. On the event loop only a body the + /// loop may poll and 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, loop_polls: bool) -> bool { match self { Iteration::Sync => true, - Iteration::Async => loop_polls(resp.version()) && resp.size_hint().exact().is_some(), + Iteration::Async => loop_polls && resp.size_hint().exact().is_some(), } } diff --git a/src/client/resp/ws.rs b/src/client/ws.rs similarity index 100% rename from src/client/resp/ws.rs rename to src/client/ws.rs diff --git a/src/client/resp/ws/cmd.rs b/src/client/ws/cmd.rs similarity index 100% rename from src/client/resp/ws/cmd.rs rename to src/client/ws/cmd.rs diff --git a/src/client/resp/ws/msg.rs b/src/client/ws/msg.rs similarity index 100% rename from src/client/resp/ws/msg.rs rename to src/client/ws/msg.rs diff --git a/src/coroutine.rs b/src/coroutine.rs index 0397661a..d4fd1a01 100644 --- a/src/coroutine.rs +++ b/src/coroutine.rs @@ -6,13 +6,15 @@ //! A [`Coroutine`] polls its future on the event loop thread, and [`spawn`] moves //! the work to Tokio. A wake marks the coroutine ready and queues it on the [`Port`] //! of the awaiting loop, ringing a socket the loop watches only when the queue was -//! empty, so waking Tokio threads never take the GIL. A loop that cannot watch the -//! socket is woken with `call_soon_threadsafe`, which does attach the waking thread. -//! The loop thread then resolves every queued asyncio future in one batch. +//! empty, so waking threads never take the GIL. A loop that cannot watch the socket is +//! woken with `call_soon_threadsafe` from a dedicated thread, as wakers may run under +//! runtime locks that a GIL holder waits for. The loop thread then resolves every queued +//! [`Waiter`](waiter::Waiter) in one batch, running the woken tasks' wakeups at once. mod asyncio; mod awaitable; mod scope; +mod waiter; use std::{ any::TypeId, @@ -21,6 +23,7 @@ use std::{ task::Poll, }; +use futures_util::{FutureExt, future::BoxFuture}; use pyo3::{IntoPyObjectExt, PyTypeInfo, intern, prelude::*, sync::PyOnceLock}; pub(crate) use self::asyncio::running_loop; @@ -57,6 +60,29 @@ where Bound::new(py, coroutine(qualname, fut)) } +/// Run `start` with `arg` on the first await. A result ready then finishes the coroutine +/// without the boxed future and waker that [`local`] allocates for a suspension. +#[inline] +pub fn eager<'py, A>( + qualname: &'static str, + arg: Bound<'py, A>, + start: fn(Bound<'_, PyAny>) -> PyResult, +) -> PyResult> { + let py = arg.py(); + Bound::new( + py, + Coroutine::eager(qualname, arg.into_any().unbind(), start), + ) +} + +/// What the first step of an [`eager`] coroutine produced. +pub enum Start { + /// The result, returned without suspending. + Ready(Py), + /// A future polled from then on, as [`local`] polls one. + Pending(BoxFuture<'static, PyResult>>), +} + /// A result type whose native `__aenter__` returns the object itself, so `async with` on a /// [`managed`] coroutine enters it without awaiting `__aenter__`. pub trait EntersSelf: PyTypeInfo { @@ -102,9 +128,7 @@ pub fn ready<'py, T>( qualname: &'static str, value: Bound<'py, T>, ) -> PyResult> { - let py = value.py(); - let value = value.into_any().unbind(); - Bound::new(py, Coroutine::new(qualname, async move { Ok(value) })) + eager(qualname, value, |value| Ok(Start::Ready(value.unbind()))) } /// Yield to the event loop once, like `asyncio.sleep(0)`. @@ -122,20 +146,54 @@ pub async fn yield_now() { .await; } -/// A coroutine awaiting `fut` and converting its output to a Python object. Like a sync -/// method, one with no result returns `None`, where `()` would convert to an empty tuple. +/// A coroutine awaiting `fut` and converting its output to a Python object. fn coroutine(qualname: &'static str, fut: F) -> Coroutine where F: Future> + Send + 'static, T: for<'a> IntoPyObject<'a> + 'static, { - Coroutine::new(qualname, async move { - let value = fut.await?; - Python::attach(|py| { - if TypeId::of::() == TypeId::of::<()>() { - return Ok(py.None()); - } - value.into_py_any(py) - }) - }) + Coroutine::new(qualname, converted(fut)) +} + +/// Await `fut` and convert its output to a Python object. +async fn converted(fut: F) -> PyResult> +where + F: Future>, + T: for<'a> IntoPyObject<'a> + 'static, +{ + let value = fut.await?; + Python::attach(|py| to_object(py, value)) +} + +/// Like a sync method's result, `()` converts to `None`, not an empty tuple. +#[inline] +fn to_object(py: Python<'_>, value: T) -> PyResult> +where + T: for<'a> IntoPyObject<'a> + 'static, +{ + if TypeId::of::() == TypeId::of::<()>() { + return Ok(py.None()); + } + value.into_py_any(py) +} + +impl Start { + /// Finish with `value`, converted as [`local`] converts a result. + #[inline] + pub fn ready(py: Python<'_>, value: T) -> PyResult + where + T: for<'a> IntoPyObject<'a> + 'static, + { + to_object(py, value).map(Start::Ready) + } + + /// Continue by polling `fut` on the event loop thread, as [`local`] does. + #[inline] + pub fn pending(fut: F) -> Start + where + F: Future> + Send + 'static, + T: for<'a> IntoPyObject<'a> + 'static, + { + Start::Pending(converted(fut).boxed()) + } } diff --git a/src/coroutine/asyncio.rs b/src/coroutine/asyncio.rs index 677dee07..6a7572aa 100644 --- a/src/coroutine/asyncio.rs +++ b/src/coroutine/asyncio.rs @@ -1,19 +1,21 @@ -//! The asyncio event loop side: the futures a suspended task waits on, and the per-loop -//! [`Port`] that resolves them when Rust work wakes. +//! The asyncio event loop side: the per-loop [`Port`] that resolves the futures suspended +//! tasks wait on when Rust work wakes. use std::{ cell::RefCell, io::{self, ErrorKind, Read, Write}, iter, mem, sync::{ - Arc, Mutex, MutexGuard, PoisonError, Weak, + Arc, Mutex, MutexGuard, OnceLock, PoisonError, Weak, atomic::{AtomicBool, Ordering}, + mpsc::{self, Sender}, }, + thread, }; use pyo3::{PyTraverseError, PyVisit, intern, prelude::*, sync::PyOnceLock}; -use super::awaitable::Slot; +use super::{awaitable::Slot, waiter::Waiter}; /// Wakes queued for one event loop, delivered by a single bell per batch. pub(crate) struct Port { @@ -28,7 +30,7 @@ enum Bell { /// A socket pair whose read end the loop watches. The read end is `None` /// when Python owns it, as with a proactor loop. Socket { tx: Socket, rx: Option }, - /// A loop that cannot watch the socket: the waking thread schedules the drain, + /// A loop that cannot watch the socket: a dedicated thread schedules the drain, /// and a [`Keeper`] holds the port until the loop closes. Loop(Py), } @@ -39,7 +41,7 @@ type Socket = std::os::unix::net::UnixStream; #[cfg(windows)] type Socket = std::net::TcpStream; -/// Resolves the asyncio futures of queued coroutines on the loop thread. +/// Resolves the waiters of queued coroutines on the loop thread, resuming their tasks. #[pyclass(frozen, module = "wreq")] struct Drain(Arc); @@ -56,13 +58,11 @@ thread_local! { // ===== impl Port ===== impl Port { - /// A new asyncio future for a task to wait on, and the port of the loop that runs it. - pub(super) fn waiter(py: Python<'_>) -> PyResult<(Py, Arc)> { + /// A new future for a task to wait on, and the port of the loop that runs it. + pub(super) fn waiter(py: Python<'_>) -> PyResult<(Py, Arc)> { let (event_loop, port) = Self::current(py)?; - let waiter = event_loop.call_method0(intern!(py, "create_future"))?; - // Tasks only accept futures marked as yielded by `await`. - waiter.setattr(intern!(py, "_asyncio_future_blocking"), true)?; - Ok((waiter.unbind(), port)) + let waiter = Py::new(py, Waiter::new(event_loop))?; + Ok((waiter, port)) } /// Return the running loop and its port, opening the port on first use. @@ -173,21 +173,52 @@ impl Port { _ => break, } }, - Bell::Loop(event_loop) => { - Python::try_attach(|py| { - let drain = Drain(self.clone()); - let scheduled = event_loop - .bind(py) - .call_method1(intern!(py, "call_soon_threadsafe"), (drain,)); - // A closed loop never runs the drain. - if scheduled.is_err() { - self.close(); - } - }); + Bell::Loop(_) => self.clone().schedule_off_thread(), + } + } + + /// Schedule a drain from a thread that only waits for the GIL, started on first use. + /// + /// Wakers may run under a runtime's internal locks, such as a semaphore's while it closes, + /// that a GIL holder may be waiting for, so they never attach themselves. Without the + /// thread, the waking thread schedules the drain. + fn schedule_off_thread(self: Arc) { + static SCHEDULER: OnceLock>>> = OnceLock::new(); + let scheduler = SCHEDULER.get_or_init(|| { + let (tx, rx) = mpsc::channel::>(); + thread::Builder::new() + .name("wreq-loop-wake".into()) + .spawn(move || rx.into_iter().for_each(|port| port.schedule_drain())) + .ok() + .map(|_| tx) + }); + match scheduler { + Some(tx) => { + if let Err(mpsc::SendError(port)) = tx.send(self) { + port.schedule_drain(); + } } + None => self.schedule_drain(), } } + /// Schedule a drain on the port's loop, attached to Python. + fn schedule_drain(self: &Arc) { + let Bell::Loop(event_loop) = &self.bell else { + return; + }; + Python::try_attach(|py| { + let drain = Drain(self.clone()); + let scheduled = event_loop + .bind(py) + .call_method1(intern!(py, "call_soon_threadsafe"), (drain,)); + // A closed loop never runs the drain. + if scheduled.is_err() { + self.close(); + } + }); + } + /// Discard queued wakes, breaking the cycle between queued slots and the port. fn close(&self) { let queue = { @@ -224,21 +255,10 @@ impl Drain { let Some(waiter) = slot.take_waiter() else { continue; }; - // Resume the task awaiting the waiter, unless it was already cancelled. - let done = waiter - .bind(py) - .call_method0(intern!(py, "done")) - .and_then(|done| done.is_truthy()); - let released = match done { - Ok(false) => waiter - .bind(py) - .call_method1(intern!(py, "set_result"), (py.None(),)) - .map(drop), - done => done.map(drop), - }; - if let Err(err) = released { - // A raising callback, such as a signal handler, ends this drain. Requeue - // this wake and the rest; the next drain skips futures already done. + // Resume the task awaiting the waiter now, unless it was already cancelled. + if let Err(err) = Waiter::resolve(waiter.bind(py)) { + // A `KeyboardInterrupt`, `SystemExit` or busy waiter ends this drain. Requeue + // this wake and the rest; the next drain skips a waiter already resolved. slot.set_waiter(waiter); port.lock().splice(..0, iter::once(slot).chain(slots)); port.ring(); diff --git a/src/coroutine/awaitable.rs b/src/coroutine/awaitable.rs index a609707e..7a300e1a 100644 --- a/src/coroutine/awaitable.rs +++ b/src/coroutine/awaitable.rs @@ -16,12 +16,12 @@ use pyo3::{ types::{PyTraceback, PyTuple, PyType}, }; -use super::{Port, scope::Scope}; +use super::{Port, Start, scope::Scope, waiter::Waiter}; /// 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 it accepts, or a close, drops the Rust +/// Each suspension hands the task a waiter following the asyncio future protocol, so +/// task cancellation works unchanged. A throw it accepts, or a close, drops the Rust /// future, which aborts any spawned work. /// /// A request coroutine also works as `async with`, as `async with await` would. @@ -31,32 +31,50 @@ use super::{Port, scope::Scope}; #[pyclass(module = "wreq")] pub struct Coroutine { qualname: &'static str, + /// Run on the first step, which may finish without polling a future. + eager: Option, /// Polled only through `&mut self`, so `get_mut` reaches it without locking. future: Mutex>>>>, + /// Created on the first poll, so a coroutine finishing on its first step needs none. + wakeup: Option, + /// A poll reached suspension, which may have published a port or waiter. + suspended: bool, + pub(super) scope: Scope, +} + +/// The first step of an eager coroutine, taken on first await. +struct Eager { + start: fn(Bound<'_, PyAny>) -> PyResult, + arg: Py, +} + +/// The waker a coroutine polls its future with, and the wake state behind it. +/// +/// Reusing one waker lets Tokio skip re-registering it on every poll. +struct Wakeup { slot: Arc, waker: Waker, - pub(super) scope: Scope, } /// The outcome of a step: the protocol's `StopIteration` is built only for Python. pub(super) enum Step { - /// Suspend, handing the task an asyncio future to wait on, or `None`. + /// Suspend, handing the task a future to wait on, or `None`. Yield(Py), /// Finish with this value. Return(Py), } -/// Wake state shared with Rust wakers for the coroutine's whole life. +/// Wake state shared with Rust wakers from the first poll on. /// -/// Reusing one waker lets Tokio skip re-registering it on every poll. Wakes only -/// change the state and queue the slot; Python objects are touched on the loop thread. +/// Wakes only change the state and queue the slot; Python objects are touched on the loop +/// thread. #[derive(Default)] pub(super) struct Slot { state: AtomicU8, /// The port of the loop the coroutine last waited on, set before each wait. port: Mutex>>, - /// The asyncio future the task awaits while the coroutine waits. - waiter: Mutex>>, + /// The future the task awaits while the coroutine waits. + waiter: Mutex>>, } const POLLING: u8 = 0; @@ -71,12 +89,27 @@ impl Coroutine { where F: Future>> + Send + 'static, { - let slot = Arc::new(Slot::default()); Coroutine { qualname, + eager: None, future: Mutex::new(Some(future.boxed())), - waker: Waker::from(slot.clone()), - slot, + wakeup: None, + suspended: false, + scope: Scope::Unsupported, + } + } + + pub(super) fn eager( + qualname: &'static str, + arg: Py, + start: fn(Bound<'_, PyAny>) -> PyResult, + ) -> Self { + Coroutine { + qualname, + eager: Some(Eager { start, arg }), + future: Mutex::new(None), + wakeup: None, + suspended: false, scope: Scope::Unsupported, } } @@ -88,6 +121,12 @@ impl Coroutine { .unwrap_or_else(PoisonError::into_inner) } + /// Whether the coroutine was not yet awaited to its end. + #[inline] + pub(super) fn is_pending(&mut self) -> bool { + self.eager.is_some() || self.future().is_some() + } + pub(super) fn step( &mut self, py: Python<'_>, @@ -98,6 +137,19 @@ impl Coroutine { } self.scope.start(); + if let Some(Eager { start, arg }) = self.eager.take() { + match start(arg.into_bound(py)) { + Ok(Start::Ready(value)) => { + return self.complete(py, value).inspect_err(|_| self.abandon()); + } + Ok(Start::Pending(future)) => *self.future() = Some(future), + Err(err) => { + self.abandon(); + return Err(err); + } + } + } + let future = self .future .get_mut() @@ -105,35 +157,46 @@ impl Coroutine { .as_mut() .ok_or_else(|| PyRuntimeError::new_err("cannot reuse already awaited coroutine"))?; + let Wakeup { slot, waker } = self.wakeup.get_or_insert_with(Wakeup::new); + // A pending waiter means another task is suspended on this coroutine. Once // it is taken, wakes from now on resume this poll. - if let Some(waiter) = self.slot.take_waiter() - && let Err(err) = ensure_done(waiter.bind(py)) + if self.suspended + && let Some(waiter) = slot.take_waiter() { - self.slot.set_waiter(waiter); - return Err(err); + let done = waiter.get().is_done(); + if !matches!(done, Ok(true)) { + slot.set_waiter(waiter); + return Err(done.err().unwrap_or_else(awaited_already)); + } } - self.slot.state.store(POLLING, Ordering::Release); - if let Poll::Ready(result) = future.as_mut().poll(&mut Context::from_waker(&self.waker)) { + slot.state.store(POLLING, Ordering::Release); + if let Poll::Ready(result) = future.as_mut().poll(&mut Context::from_waker(waker)) { return result .and_then(|value| self.complete(py, value)) .inspect_err(|_| self.abandon()); } // A coroutine that cannot wait ends here, releasing any spawned work. - self.slot - .suspend(py) - .map(Step::Yield) + self.suspended = true; + slot.suspend(py) + .map(|waiter| Step::Yield(waiter.map_or_else(|| py.None(), Py::into_any))) .inspect_err(|_| self.abandon()) } pub(super) fn finish(&mut self) { + self.eager = None; *self.future() = None; - self.slot.state.store(DONE, Ordering::Release); + let Some(Wakeup { slot, .. }) = &self.wakeup else { + return; + }; + slot.state.store(DONE, Ordering::Release); // A finished coroutine must not keep its loop's port open. - let port = self.slot.lock_port().take(); - drop(port); - drop(self.slot.take_waiter()); + if self.suspended { + let port = slot.lock_port().take(); + drop(port); + drop(slot.take_waiter()); + } } } @@ -189,14 +252,19 @@ impl Coroutine { fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { self.scope.traverse(&visit)?; + if let Some(eager) = &self.eager { + visit.call(&eager.arg)?; + } + let Some(Wakeup { slot, .. }) = &self.wakeup else { + return Ok(()); + }; // Pending Rust work owns the waiter until its loop closes the port; only then // can a cycle through the waiter be collected without aborting live requests. - let closed = self - .slot + let closed = slot .port .try_lock() .is_ok_and(|port| port.as_ref().is_some_and(|port| !port.is_open())); - if closed && let Ok(waiter) = self.slot.waiter.try_lock() { + if closed && let Ok(waiter) = slot.waiter.try_lock() { visit.call(&*waiter)?; } Ok(()) @@ -316,9 +384,23 @@ pub(super) fn ensure_done(waiter: &Bound<'_, PyAny>) -> PyResult<()> { { Ok(()) } else { - Err(PyRuntimeError::new_err( - "coroutine is being awaited already", - )) + Err(awaited_already()) + } +} + +fn awaited_already() -> PyErr { + PyRuntimeError::new_err("coroutine is being awaited already") +} + +// ===== impl Wakeup ===== + +impl Wakeup { + fn new() -> Self { + let slot = Arc::new(Slot::default()); + Wakeup { + waker: Waker::from(slot.clone()), + slot, + } } } @@ -327,9 +409,9 @@ pub(super) fn ensure_done(waiter: &Bound<'_, PyAny>) -> PyResult<()> { impl Slot { /// Suspend after a pending poll, returning the future the task should await, /// or `None` for a bare yield when a wake already arrived. - fn suspend(&self, py: Python<'_>) -> PyResult> { + fn suspend(&self, py: Python<'_>) -> PyResult>> { if self.state.load(Ordering::Acquire) == NOTIFIED { - return Ok(py.None()); + return Ok(None); } let (waiter, port) = Port::waiter(py)?; // Publish the port and waiter before waiting, so a wake that sees WAITING @@ -340,11 +422,11 @@ impl Slot { .state .compare_exchange(POLLING, WAITING, Ordering::AcqRel, Ordering::Acquire) { - Ok(_) => Ok(waiter), + Ok(_) => Ok(Some(waiter)), // Woken after the poll but before waiting: resume on the next loop turn. Err(_) => { drop(self.take_waiter()); - Ok(py.None()) + Ok(None) } } } @@ -355,12 +437,16 @@ impl Slot { } #[inline] - pub(super) fn set_waiter(&self, waiter: Py) { - *self.waiter.lock().unwrap_or_else(PoisonError::into_inner) = Some(waiter); + pub(super) fn set_waiter(&self, waiter: Py) { + let mut slot = self.waiter.lock().unwrap_or_else(PoisonError::into_inner); + let replaced = slot.replace(waiter); + // Released after the lock: dropping a waiter may run Python, which may wake this slot. + drop(slot); + drop(replaced); } #[inline] - pub(super) fn take_waiter(&self) -> Option> { + pub(super) fn take_waiter(&self) -> Option> { self.waiter .lock() .unwrap_or_else(PoisonError::into_inner) diff --git a/src/coroutine/scope.rs b/src/coroutine/scope.rs index 6a2ef446..77c7db43 100644 --- a/src/coroutine/scope.rs +++ b/src/coroutine/scope.rs @@ -207,7 +207,7 @@ impl Coroutine { impl Coroutine { /// Enter by awaiting the coroutine itself, which also awaits the result's `__aenter__`. fn __aenter__(mut slf: PyRefMut<'_, Self>) -> PyResult> { - let pending = slf.future().is_some(); + let pending = slf.is_pending(); match slf.scope { Scope::Unsupported => Err(PyTypeError::new_err( "coroutine does not support the asynchronous context manager protocol", diff --git a/src/coroutine/waiter.rs b/src/coroutine/waiter.rs new file mode 100644 index 00000000..8558055c --- /dev/null +++ b/src/coroutine/waiter.rs @@ -0,0 +1,322 @@ +//! The future an asyncio task waits on while a coroutine's Rust work is pending. + +use std::sync::{ + Mutex, MutexGuard, TryLockError, + atomic::{AtomicBool, Ordering}, +}; + +use pyo3::{ + PyTraverseError, PyVisit, + exceptions::{ + PyKeyboardInterrupt, PyRuntimeError, PySystemExit, + asyncio::{CancelledError, InvalidStateError}, + }, + intern, + prelude::*, + sync::PyOnceLock, + types::PyDict, +}; + +/// The future a task waits on while a coroutine's Rust work is pending. +/// +/// It implements the protocol `asyncio.Task` follows for futures from other libraries. A +/// drain resolves it on the loop thread and runs the task's wakeup at once, where an +/// `asyncio.Future` would schedule the wakeup for the next loop iteration. +/// +/// Its state is only locked on the loop thread and never while Python runs, so a lock +/// already held means a reentrant call; it fails instead of deadlocking the loop. +#[pyclass(frozen, module = "wreq")] +pub(super) struct Waiter { + event_loop: Py, + state: Mutex, + /// `_asyncio_future_blocking`: set when yielded and cleared by the task accepting it. + blocking: AtomicBool, +} + +#[derive(Default)] +struct State { + status: Status, + /// Usually one: the wakeup of the task awaiting the coroutine. + callbacks: Vec, + cancel_message: Option>, +} + +#[derive(Clone, Copy, Default, PartialEq, Eq)] +enum Status { + #[default] + Pending, + Done, + Cancelled, +} + +struct Callback { + func: Py, + context: Py, +} + +// ===== impl Waiter ===== + +impl Waiter { + /// A pending waiter on `event_loop`, marked as yielded by `await`. + pub(super) fn new(event_loop: Bound<'_, PyAny>) -> Self { + Waiter { + event_loop: event_loop.unbind(), + state: Mutex::default(), + blocking: AtomicBool::new(true), + } + } + + #[inline] + fn lock(&self) -> PyResult> { + match self.state.try_lock() { + Ok(state) => Ok(state), + Err(TryLockError::Poisoned(err)) => Ok(err.into_inner()), + Err(TryLockError::WouldBlock) => { + Err(PyRuntimeError::new_err("waiter state is already in use")) + } + } + } + + #[inline] + fn status(&self) -> PyResult { + Ok(self.lock()?.status) + } + + #[inline] + pub(super) fn is_done(&self) -> PyResult { + Ok(self.status()? != Status::Pending) + } + + /// Mark the waiter done and run its callbacks now, as `Handle` would run them. + /// + /// A `KeyboardInterrupt` or `SystemExit` from a callback is returned after the remaining + /// callbacks are scheduled, so it stops the loop like one raised by a handle. + pub(super) fn resolve(slf: &Bound<'_, Self>) -> PyResult<()> { + let callbacks = { + let mut state = slf.get().lock()?; + if state.status != Status::Pending { + return Ok(()); + } + state.status = Status::Done; + std::mem::take(&mut state.callbacks) + }; + let py = slf.py(); + let mut callbacks = callbacks.into_iter(); + while let Some(callback) = callbacks.next() { + let ran = callback + .context + .bind(py) + .call_method1(intern!(py, "run"), (callback.func.bind(py), slf)); + let Err(err) = ran else { + continue; + }; + if !err.is_instance_of::(py) + && !err.is_instance_of::(py) + { + slf.get().report(py, &callback, err); + continue; + } + for callback in callbacks { + if let Err(err) = Self::schedule(slf, &callback) { + err.write_unraisable(py, Some(slf.as_any())); + } + } + return Err(err); + } + Ok(()) + } + + /// Run `callback` with this waiter on the next loop iteration. + fn schedule(slf: &Bound<'_, Self>, callback: &Callback) -> PyResult<()> { + let py = slf.py(); + let kwargs = PyDict::new(py); + kwargs.set_item(intern!(py, "context"), callback.context.bind(py))?; + slf.get() + .event_loop + .bind(py) + .call_method( + intern!(py, "call_soon"), + (callback.func.bind(py), slf), + Some(&kwargs), + ) + .map(drop) + } + + /// Pass a callback's error to the loop's exception handler, as `Handle` does. + fn report(&self, py: Python<'_>, callback: &Callback, err: PyErr) { + let reported = (|| { + let context = PyDict::new(py); + let message = format!("Exception in callback {}", callback.func.bind(py).repr()?); + context.set_item(intern!(py, "message"), message)?; + context.set_item(intern!(py, "exception"), err.value(py))?; + self.event_loop + .bind(py) + .call_method1(intern!(py, "call_exception_handler"), (context,)) + .map(drop) + })(); + if let Err(err) = reported { + err.write_unraisable(py, Some(callback.func.bind(py))); + } + } + + fn cancelled_error(&self, py: Python<'_>) -> PyErr { + let message = self + .lock() + .map(|state| state.cancel_message.as_ref().map(|m| m.clone_ref(py))); + match message { + Ok(Some(message)) => CancelledError::new_err((message,)), + Ok(None) => CancelledError::new_err(()), + Err(err) => err, + } + } +} + +#[pymethods] +impl Waiter { + #[getter(_asyncio_future_blocking)] + fn get_blocking(&self) -> bool { + self.blocking.load(Ordering::Relaxed) + } + + #[setter(_asyncio_future_blocking)] + fn set_blocking(&self, value: bool) { + self.blocking.store(value, Ordering::Relaxed); + } + + fn get_loop(&self, py: Python<'_>) -> Py { + self.event_loop.clone_ref(py) + } + + #[pyo3(signature = (func, *, context = None))] + fn add_done_callback( + slf: &Bound<'_, Self>, + func: Bound<'_, PyAny>, + context: Option>, + ) -> PyResult<()> { + static COPY_CONTEXT: PyOnceLock> = PyOnceLock::new(); + let py = slf.py(); + let context = match context { + Some(context) if !context.is_none() => context, + _ => COPY_CONTEXT + .import(py, "contextvars", "copy_context")? + .call0()?, + }; + let callback = Callback { + func: func.unbind(), + context: context.unbind(), + }; + { + let mut state = slf.get().lock()?; + if state.status == Status::Pending { + state.callbacks.push(callback); + return Ok(()); + } + } + Self::schedule(slf, &callback) + } + + /// Remove callbacks equal to `func`, comparing outside the lock as `__eq__` may run Python. + fn remove_done_callback(&self, func: &Bound<'_, PyAny>) -> PyResult { + let py = func.py(); + let candidates: Vec> = self + .lock()? + .callbacks + .iter() + .map(|callback| callback.func.clone_ref(py)) + .collect(); + let mut matched = Vec::new(); + for candidate in candidates { + if candidate.is(func) || candidate.bind(py).eq(func)? { + matched.push(candidate); + } + } + if matched.is_empty() { + return Ok(0); + } + // Removed callbacks are released after the lock, as dropping them may run Python. + let removed: Vec = { + let mut state = self.lock()?; + let (removed, kept) = std::mem::take(&mut state.callbacks) + .into_iter() + .partition(|callback| matched.iter().any(|func| func.is(&callback.func))); + state.callbacks = kept; + removed + }; + Ok(removed.len()) + } + + #[pyo3(signature = (msg = None))] + fn cancel(slf: &Bound<'_, Self>, msg: Option>) -> PyResult { + let callbacks = { + let mut state = slf.get().lock()?; + if state.status != Status::Pending { + return Ok(false); + } + state.status = Status::Cancelled; + state.cancel_message = msg; + std::mem::take(&mut state.callbacks) + }; + for callback in &callbacks { + Self::schedule(slf, callback)?; + } + Ok(true) + } + + fn cancelled(&self) -> PyResult { + Ok(self.status()? == Status::Cancelled) + } + + fn done(&self) -> PyResult { + self.is_done() + } + + fn result(&self, py: Python<'_>) -> PyResult<()> { + match self.status()? { + Status::Pending => Err(InvalidStateError::new_err("Result is not set.")), + Status::Cancelled => Err(self.cancelled_error(py)), + Status::Done => Ok(()), + } + } + + fn exception(&self, py: Python<'_>) -> PyResult>> { + match self.status()? { + Status::Pending => Err(InvalidStateError::new_err("Exception is not set.")), + Status::Cancelled => Err(self.cancelled_error(py)), + Status::Done => Ok(None), + } + } + + fn __repr__(&self) -> &'static str { + match self.status() { + Ok(Status::Pending) => "", + Ok(Status::Done) => "", + Ok(Status::Cancelled) => "", + Err(_) => "", + } + } + + fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.event_loop)?; + let Ok(state) = self.lock() else { + return Ok(()); + }; + for callback in &state.callbacks { + visit.call(&callback.func)?; + visit.call(&callback.context)?; + } + visit.call(&state.cancel_message) + } + + fn __clear__(&self) { + let Ok(mut state) = self.lock() else { + return; + }; + let cleared = ( + std::mem::take(&mut state.callbacks), + state.cancel_message.take(), + ); + // Dropping references may run Python, so release the lock first. + drop(state); + drop(cleared); + } +} diff --git a/src/lib.rs b/src/lib.rs index 796ab7d6..a7aaed94 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -25,7 +25,8 @@ mod tls; use client::{ BlockingClient, Client, SocketAddr, body::multipart::{Multipart, Part}, - resp::{BlockingResponse, BlockingWebSocket, Message, Response, Streamer, WebSocket}, + resp::{BlockingResponse, Response, Streamer}, + ws::{BlockingWebSocket, WebSocket, msg::Message}, }; use cookie::{Cookie, Jar, SameSite}; use dns::{DnsOptions, LookupIpStrategy}; @@ -212,7 +213,8 @@ mod blocking { client::{ BlockingClient, req::{Request, WebSocketRequest}, - resp::{BlockingResponse, BlockingWebSocket}, + resp::BlockingResponse, + ws::BlockingWebSocket, }, http::Method, }; diff --git a/src/macros.rs b/src/macros.rs index f8e2d486..4ac46c08 100644 --- a/src/macros.rs +++ b/src/macros.rs @@ -29,6 +29,20 @@ where .transpose() } +/// Take one exact-string keyword without invoking key or value callbacks. +pub(crate) fn single_keyword<'py>( + ob: &Bound<'py, PyAny>, +) -> Option<(Bound<'py, PyString>, Bound<'py, PyAny>)> { + let dict = ob.cast_exact::().ok()?; + pyo3::sync::critical_section::with_critical_section(dict.as_any(), || { + if dict.len() != 1 { + return None; + } + let (key, value) = dict.iter().next()?; + Some((key.cast_into_exact::().ok()?, value)) + }) +} + /// Look up an optional keyword; a dict answers a missing key without raising KeyError. pub(crate) fn lookup<'py>( ob: &Bound<'py, PyAny>, @@ -42,7 +56,9 @@ pub(crate) fn lookup<'py>( derive_deftly::define_derive_deftly! { /// Derive keyword names and extraction from the parameter fields. - /// Reject unknown keys when the type specifies `deny_unknown`. + /// Reject unknown keys when the type specifies `deny_unknown`. With `single_keyword`, a + /// lone exact-string keyword is converted directly, and `from_single_keyword` exposes that + /// dispatch to wrapping types. /// Marked fields are converted last, after other parameters are validated. ExtractOptions for struct: @@ -50,6 +66,31 @@ derive_deftly::define_derive_deftly! { pub(crate) const FIELDS: &'static [&'static str] = &[ ${for fields { ${concat $fname}, }} ]; + + ${if tmeta(single_keyword) { + /// The parameters for a lone keyword `name` set to `value`; `None` for an unknown + /// name. Converting the value fails the same way as a full extraction. + pub(crate) fn from_single_keyword( + name: &str, + value: &pyo3::Bound<'_, pyo3::PyAny>, + ) -> Option> { + use pyo3::prelude::*; + match name { + ${for fields { + ${concat $fname} => Some( + value + .extract() + .map(|$fname| Self { + $fname, + ..Self::default() + }) + .map_err(pyo3::PyErr::from), + ), + }} + _ => None, + } + } + }} } impl<'__extract_a, '__extract_py, $tgens> @@ -60,6 +101,16 @@ derive_deftly::define_derive_deftly! { fn extract(ob: pyo3::Borrowed<'__extract_a, '__extract_py, pyo3::PyAny>) -> pyo3::PyResult { use pyo3::prelude::*; + ${if tmeta(single_keyword) { + // Nested: the template expands under the proc-macro crate's edition. + if let Some((key, value)) = $crate::macros::single_keyword(&ob) { + if let Ok(name) = key.to_str() { + if let Some(params) = Self::from_single_keyword(name, &value) { + return params; + } + } + } + }} ${if tmeta(deny_unknown) { let _ = $crate::macros::reject_unknown_kwargs( ob.cast::().ok(), @@ -101,6 +152,64 @@ derive_deftly::define_derive_deftly! { } } +derive_deftly::define_derive_deftly! { + /// Extract inline fields directly and delegate other parameters to boxed options. + /// The field marked `flatten` names the full options type. + ExtractOptionsWrapper for struct: + + impl<$tgens> $ttype where $twheres { + ${for fields { + ${when fmeta(flatten)} + pub(crate) const FIELDS: &'static [&'static str] = + ${fmeta(flatten) as ty}::FIELDS; + }} + } + + impl<'__extract_a, '__extract_py, $tgens> + pyo3::FromPyObject<'__extract_a, '__extract_py> for $ttype + where $twheres + { + type Error = pyo3::PyErr; + + fn extract(ob: pyo3::Borrowed<'__extract_a, '__extract_py, pyo3::PyAny>) -> pyo3::PyResult { + use pyo3::prelude::*; + let wrap = |options| Self { + ${for fields { + ${when fmeta(flatten)} + $fname: Some(Box::new(options)), + }} + ..Self::default() + }; + if let Some((key, value)) = $crate::macros::single_keyword(&ob) { + if let Ok(name) = key.to_str() { + match name { + ${for fields { + ${when not(fmeta(flatten))} + ${concat $fname} => return value.extract().map(|value| Self { + $fname: value, + ..Self::default() + }), + }} + _ => {} + } + ${for fields { + ${when fmeta(flatten)} + if let Some(options) = + ${fmeta(flatten) as ty}::from_single_keyword(name, &value) + { + return options.map(wrap); + } + }} + } + } + ${for fields { + ${when fmeta(flatten)} + ${fmeta(flatten) as ty}::extract(ob).map(wrap) + }} + } + } +} + macro_rules! apply_option { (set_if_some, $builder:expr, $option:expr, $method:ident) => { if let Some(value) = $option.take() { diff --git a/src/runtime.rs b/src/runtime.rs index a23a44d9..bfcce092 100644 --- a/src/runtime.rs +++ b/src/runtime.rs @@ -1,10 +1,11 @@ //! 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. +//! Async requests and reads that must wait run as tasks on the client's selected worker +//! ([`Runtime::join`], [`Runtime::run_task`]). Blocking requests run on the caller with work +//! stealing, or on the selected worker with per-worker scheduling ([`Runtime::block_on_task`]); +//! blocking reads are polled by the caller ([`Runtime::block_on`], [`Runtime::block_on_eager`]). +//! Blocking calls release the GIL while waiting. A current-thread runtime also drives its +//! connection tasks and IO on the caller, with no background IO workers. //! //! Every other runtime operation goes through [`Runtime`] too; `clippy.toml` rejects direct //! Tokio handles, spawns and yields elsewhere. @@ -51,9 +52,13 @@ pub enum Scheduler { #[derive(Clone)] #[pyclass(frozen)] pub struct Runtime { - /// Shared owner; `None` only once `Drop` has taken it to shut the runtime down. + binding: Arc, +} + +/// The selected worker and its owner, shared by a client and its responses. +/// Selecting another client creates a binding; cloning keeps that worker. +struct Binding { inner: Option>, - /// The worker this copy runs on; per-worker clients keep their tasks there. handle: Handle, } @@ -79,25 +84,28 @@ thread_local! { impl Runtime { /// Whether blocking calls drive this runtime on their own thread. pub fn is_current_thread(&self) -> bool { - matches!(self.inner.as_deref(), Some(Inner::CurrentThread(_))) + matches!(self.binding.inner.as_deref(), Some(Inner::CurrentThread(_))) } /// Share the runtime and select a worker for a new client. pub fn select(&self) -> PyResult { let inner = self + .binding .inner .as_ref() .ok_or_else(|| PyRuntimeError::new_err("Runtime is unavailable"))?; Ok(Self { - handle: inner.handle(), - inner: Some(inner.clone()), + binding: Arc::new(Binding { + handle: inner.handle(), + inner: Some(inner.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(); + let _runtime = self.binding.handle.enter(); future.now_or_never() } @@ -127,8 +135,9 @@ impl Runtime { } } - /// 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. + /// Run a request detached from Python: on the caller with work stealing, or on the + /// selected worker with per-worker scheduling. A current-thread runtime also drives + /// its connection tasks and IO on the caller. pub fn block_on_task(&self, py: Python<'_>, future: F) -> PyResult where F: Future> + Send + 'static, @@ -145,11 +154,17 @@ impl Runtime { F: Future> + Send + 'static, T: Send + 'static, { - let task = self.handle.clone().spawn(async move { + let task = self.binding.handle.clone().spawn(async move { let _owner = self; future.await }); - AbortOnDropHandle::new(task).await.map_err(join_error)? + Self::join(task).await + } + + /// Wait for a task from [`spawn`](Self::spawn) from any executor. Dropping the returned + /// future aborts the task; the task must keep the runtime alive itself. + pub fn join(task: JoinHandle>) -> impl Future> { + AbortOnDropHandle::new(task).map(|result| result.map_err(join_error)?) } /// Spawn `future` on the selected worker. Dropping the handle detaches the task, which @@ -159,28 +174,28 @@ impl Runtime { F: Future + Send + 'static, F::Output: Send + 'static, { - self.handle.spawn(future) + self.binding.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() { + match self.binding.inner.as_deref() { Some(Inner::CurrentThread(runtime)) => Driving::drive(runtime, future), - _ => self.handle.block_on(future), + _ => self.binding.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. + /// polling or spawning, 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() { + match self.binding.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 @@ -189,9 +204,13 @@ impl Runtime { task::yield_now().await; future.await }), + Some(Inner::Workers(PingoraRuntime::Steal { .. })) => { + self.binding.handle.block_on(future) + } _ => self + .binding .handle - .block_on(self.handle.spawn(future)) + .block_on(self.binding.handle.spawn(future)) .map_err(join_error)?, } } @@ -274,8 +293,10 @@ impl From for Runtime { fn from(runtime: Inner) -> Self { let handle = runtime.handle(); Self { - inner: Some(Arc::new(runtime)), - handle, + binding: Arc::new(Binding { + inner: Some(Arc::new(runtime)), + handle, + }), } } } @@ -369,7 +390,9 @@ impl Runtime { } } -impl Drop for Runtime { +// ===== impl Binding ===== + +impl Drop for Binding { fn drop(&mut self) { // 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) { @@ -430,13 +453,42 @@ fn parallelism() -> usize { /// The runtime this thread is in, else the shared one. fn current() -> Handle { - Handle::try_current().unwrap_or_else(|_| Runtime::shared().handle.clone()) + Handle::try_current().unwrap_or_else(|_| Runtime::shared().binding.handle.clone()) } #[cfg(test)] mod tests { use super::*; + #[test] + fn work_stealing_polls_requests_on_callers_and_keeps_background_workers() { + let runtime = Runtime::from(RuntimeBuilder::new(2, "wreq-caller").build()); + let caller = thread::current().id(); + let (release, released) = tokio::sync::oneshot::channel(); + let (completed, completion) = std::sync::mpsc::channel(); + let child = runtime + .wait_task(async move { + assert!(!Runtime::driving()); + assert_eq!(thread::current().id(), caller); + task::yield_now().await; + assert_eq!(thread::current().id(), caller); + Ok(task::spawn(async move { + released.await.unwrap(); + tokio::time::sleep(Duration::from_millis(1)).await; + completed.send(thread::current().id()).unwrap(); + })) + }) + .unwrap(); + // Only release the child after the blocking call returns. Its timer must keep + // progressing without the caller entering the runtime again. + release.send(()).unwrap(); + assert_ne!( + completion.recv_timeout(Duration::from_secs(5)).unwrap(), + caller, + ); + runtime.binding.handle.block_on(child).unwrap(); + } + #[test] fn fixed_worker() { let runtime = Runtime::from( @@ -521,8 +573,8 @@ mod tests { .work_steal(steal) .build(), ); - let weak = Arc::downgrade(runtime.inner.as_ref().unwrap()); - let handle = runtime.handle.clone(); + let weak = Arc::downgrade(runtime.binding.inner.as_ref().unwrap()); + let handle = runtime.binding.handle.clone(); let (dropped, released) = std::sync::mpsc::channel(); let guard = NotifyOnDrop(dropped); let background = handle.spawn(async move { diff --git a/tests/aio_test.py b/tests/aio_test.py index d0761b11..02e5f64b 100644 --- a/tests/aio_test.py +++ b/tests/aio_test.py @@ -339,3 +339,49 @@ async def slow(): for writer in writers: writer.close() await server.wait_closed() + + +@pytest.mark.asyncio +async def test_short_bodies_keep_their_frames(): + # A short body may be read with its head; data still arriving and a body cut short must + # read as before, on the event loop and on a blocking caller. + async def serve(reader, writer): + try: + request = await reader.readuntil(b"\r\n\r\n") + writer.write(b"HTTP/1.1 200 OK\r\nContent-Length: 10\r\n\r\nhello") + await writer.drain() + if request.startswith(b"GET /delayed"): + await asyncio.sleep(0.05) + writer.write(b"world") + await writer.drain() + finally: + writer.close() + + def read_blocking(url): + with wreq.blocking.Client(proxies=[]) as client: + chunks = [] + with pytest.raises(wreq.exceptions.DecodingError): + with client.get(url).stream() as stream: + chunks.extend(bytes(chunk) for chunk in stream) + return chunks + + server = await asyncio.start_server(serve, "127.0.0.1", 0) + url = f"http://127.0.0.1:{server.sockets[0].getsockname()[1]}" + try: + async with wreq.Client(proxies=[]) as client: + async with (await client.get(f"{url}/delayed")).stream() as stream: + assert b"".join([bytes(chunk) async for chunk in stream]) == b"helloworld" + + chunks = [] + with pytest.raises(wreq.exceptions.DecodingError): + async with (await client.get(f"{url}/cut")).stream() as stream: + async for chunk in stream: + chunks.append(bytes(chunk)) + assert chunks == [b"hello"] + with pytest.raises(wreq.exceptions.DecodingError): + await (await client.get(f"{url}/cut")).text() + + assert await asyncio.to_thread(read_blocking, f"{url}/cut") == [b"hello"] + finally: + server.close() + await server.wait_closed() diff --git a/tests/cancellation_test.py b/tests/cancellation_test.py index d6c70b34..024e0033 100644 --- a/tests/cancellation_test.py +++ b/tests/cancellation_test.py @@ -221,7 +221,7 @@ async def test_cancellation_after_rust_completion(operation): try: # Complete the Rust work without resuming its Python coroutine. - assert isinstance(waiter, asyncio.Future) + assert asyncio.isfuture(waiter) if operation == "request_error": writer.write(b"invalid HTTP response\r\n\r\n") writer.close() @@ -284,7 +284,7 @@ async def test_pending_stream_cancellation(action): coroutine = anext(stream) if action == "close_coroutine": - assert isinstance(coroutine.send(None), asyncio.Future) + assert asyncio.isfuture(coroutine.send(None)) coroutine.close() else: started = asyncio.Event() diff --git a/tests/kwargs_test.py b/tests/kwargs_test.py new file mode 100644 index 00000000..0a5778e2 --- /dev/null +++ b/tests/kwargs_test.py @@ -0,0 +1,203 @@ +"""Request extraction regressions; closed clients never poll a network future.""" + +import asyncio +import gc +import sys +import sysconfig +import threading +from datetime import timedelta + +import pytest +import wreq +from wreq.runtime import Runtime + + +URL = "not a URL" + + +@pytest.fixture +def closed_clients(): + runtime = Runtime(workers=1) + clients = (wreq.Client(runtime=runtime), wreq.blocking.Client(runtime=runtime)) + for client in clients: + client.close() + try: + yield clients + finally: + for client in clients: + client.close() + + +def captured_kwargs(value, fields): + matches = [ + item + for item in gc.get_referrers(value) + if type(item) is dict + and set(item) == set(fields) + and any(entry is value for entry in item.values()) + ] + assert len(matches) == 1 + return matches[0] + + +class IndexValue: + def __init__(self): + self.target = None + self.calls = 0 + + def __index__(self): + self.calls += 1 + if self.target is not None: + self.target["headers"] = object() + self.target["body"] = object() + return 3 + + +@pytest.mark.asyncio +async def test_single_kwargs_conversion_and_error_timing(closed_clients): + client, blocking = closed_clients + for options in ( + {"body": None}, + {"body": b"data"}, + {"json": None}, + {"json": {"answer": 3}}, + {"headers": {"x-test": "value"}}, + {"timeout": timedelta(seconds=1)}, + {"zstd": False}, + ): + with pytest.raises(asyncio.CancelledError): + await client.post(URL, **options) + with pytest.raises(asyncio.CancelledError): + blocking.post(URL, **options) + + for route in (client, wreq): + for name in ("complete_nonsense", "\ud800"): + with pytest.raises(TypeError, match="unexpected keyword argument"): + route.post(URL, **{name: object()}) + pending = route.post(URL, headers=object()) + with pytest.raises(TypeError): + await pending + for route in (blocking, wreq.blocking): + with pytest.raises(TypeError, match="unexpected keyword argument"): + route.post(URL, complete_nonsense=object()) + with pytest.raises(TypeError): + route.post(URL, headers=object()) + + +@pytest.mark.asyncio +@pytest.mark.skipif( + sys.implementation.name != "cpython", + reason="gc.get_referrers cannot see a dict held only by the extension", +) +async def test_kwargs_callback_mutation_keeps_extraction_order(closed_clients): + client, _ = closed_clients + for multiple in (False, True): + value = IndexValue() + pending = ( + client.post(URL, json=value, auth=None) + if multiple + else client.post(URL, json=value) + ) + value.target = captured_kwargs(value, ("json", "auth") if multiple else ("json",)) + try: + # With one field, remaining reaches zero after json conversion. + # With two, the newly added invalid headers are encountered next. + with pytest.raises(TypeError if multiple else asyncio.CancelledError): + await pending + assert value.calls == 1 + finally: + value.target = None + pending.close() + + +@pytest.mark.asyncio +@pytest.mark.skipif( + sys.implementation.name != "cpython", + reason="gc.get_referrers cannot see a dict held only by the extension", +) +async def test_captured_kwargs_are_rechecked_on_first_await(closed_clients): + client, _ = closed_clients + value = IndexValue() + pending = client.post(URL, json=value) + captured = captured_kwargs(value, ("json",)) + captured["complete_nonsense"] = None + with pytest.raises(TypeError, match="unexpected keyword argument"): + await pending + assert value.calls == 0 + + value = IndexValue() + pending = client.post(URL, json=value) + pending.close() + assert value.calls == 0 + + +@pytest.mark.asyncio +async def test_string_subclass_keyword_keeps_lookup_fallback(closed_clients): + client, blocking = closed_clients + + class Keyword(str): + def __new__(cls, value): + instance = super().__new__(cls, value) + instance.comparisons = 0 + return instance + + __hash__ = str.__hash__ + + def __eq__(self, other): + self.comparisons += 1 + return super().__eq__(other) + + key = Keyword("body") + with pytest.raises(asyncio.CancelledError): + await client.post(URL, **{key: b"data"}) + assert key.comparisons > 0 + key.comparisons = 0 + with pytest.raises(asyncio.CancelledError): + blocking.post(URL, **{key: b"data"}) + assert key.comparisons > 0 + + key = "".join(("bo", "dy")) + with pytest.raises(asyncio.CancelledError): + await client.post(URL, **{key: None}) + + +@pytest.mark.asyncio +@pytest.mark.skipif( + not sysconfig.get_config_var("Py_GIL_DISABLED"), + reason="Concurrent native extraction requires free-threaded Python", +) +async def test_kwargs_value_replacement_is_safe_without_gil(closed_clients): + assert not sys._is_gil_enabled() + client, _ = closed_clients + values = [IndexValue() for _ in range(12)] + pending = [client.post(URL, json=value) for value in values] + targets = [captured_kwargs(value, ("json",)) for value in values] + stopped, started = threading.Event(), threading.Event() + failures = [] + + def replace_values(): + try: + started.set() + while not stopped.is_set(): + for target in targets: + # Preserve the key set: baseline dictionary iteration can + # reject size changes; values remain valid json options. + target["json"] = None + target["json"] = 3 + except BaseException as error: + failures.append(error) + + thread = threading.Thread(target=replace_values, daemon=True) + thread.start() + try: + assert started.wait(2) + for operation in pending: + with pytest.raises(asyncio.CancelledError): + await operation + finally: + stopped.set() + thread.join(3) + for operation in pending: + operation.close() + assert not thread.is_alive() + assert not failures diff --git a/tests/response_test.py b/tests/response_test.py index bd89ba4d..b73028b1 100644 --- a/tests/response_test.py +++ b/tests/response_test.py @@ -1,6 +1,8 @@ import asyncio +import gc import json import socket +import sys import threading import time from contextlib import asynccontextmanager @@ -195,6 +197,63 @@ def read_blocking(url): assert isinstance(value[0]["n"][3], float) +@pytest.mark.asyncio +@pytest.mark.parametrize("blocking", [False, True], ids=["async", "blocking"]) +@pytest.mark.parametrize("scope", ["client", "request"]) +@pytest.mark.parametrize("reader", ["bytes", "stream"]) +@pytest.mark.parametrize("prefix", [0, 5], ids=["head-only", "partial-body"]) +async def test_read_timeout_starts_when_the_caller_reads(blocking, scope, reader, prefix): + body = b"0123456789" + timeout = {"read_timeout": timedelta(milliseconds=150)} + client_options = timeout if scope == "client" else {} + request_options = timeout if scope == "request" else {} + received = asyncio.Event() + release = threading.Event() + loop = asyncio.get_running_loop() + + def read_blocking(url): + with wreq.blocking.Client(proxies=[], **client_options) as client: + with client.get(url, **request_options) as response: + loop.call_soon_threadsafe(received.set) + assert release.wait(5) + if reader == "bytes": + return bytes(response.bytes()) + with response.stream() as stream: + return b"".join(stream) + + async def read_async(url): + async with wreq.Client(proxies=[], **client_options) as client: + async with await client.get(url, **request_options) as response: + received.set() + assert await asyncio.to_thread(release.wait, 5) + if reader == "bytes": + return bytes(await response.bytes()) + async with response.stream() as stream: + return b"".join([chunk async for chunk in stream]) + + async with local_server() as (url, connections): + task = asyncio.create_task( + asyncio.to_thread(read_blocking, url) if blocking else read_async(url) + ) + try: + _, writer = await asyncio.wait_for(connections.get(), 5) + writer.write(b"HTTP/1.1 200 OK\r\nContent-Length: 10\r\n\r\n" + body[:prefix]) + await writer.drain() + await asyncio.wait_for(received.wait(), 5) + # The rest arrives within the timeout, while application code waits longer. + await asyncio.sleep(0.05) + writer.write(body[prefix:]) + await writer.drain() + await asyncio.sleep(0.3) + release.set() + assert await asyncio.wait_for(task, 5) == body + finally: + release.set() + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + @pytest.mark.asyncio async def test_empty_body_and_unit_results(): async with local_server() as (url, connections): @@ -224,6 +283,117 @@ async def get(): assert await client.__aexit__(None, None, None) is None +@pytest.mark.asyncio +async def test_ready_awaits_finish_on_first_step_and_stay_lazy(): + payload = b'{"ready": 1}' + head = b"HTTP/1.1 200 OK\r\nConnection: close\r\nContent-Length: %d\r\n\r\n" + # Only CPython's GC sees what an extension object holds. + cpython = sys.implementation.name == "cpython" + async with local_server() as (url, connections), wreq.Client(proxies=[]) as client: + + async def get(): + task = asyncio.create_task(client.get(url)) + _, writer = await asyncio.wait_for(connections.get(), 5) + writer.write(head % len(payload) + payload) + return await asyncio.wait_for(task, 5) + + def finish(coroutine): + with pytest.raises(StopIteration) as stop: + coroutine.send(None) + with pytest.raises(RuntimeError, match="already awaited"): + coroutine.send(None) + return stop.value.value + + response = await get() + # Nothing runs before the first step: an unstarted close or exit leaves the body + # readable. + response.close().close() + response.__aexit__(None, None, None).close() + # A buffered body finishes on the first step. + assert bytes(finish(response.bytes())) == payload + assert finish(response.json()) == {"ready": 1} + assert finish(response.text()) == payload.decode() + assert finish(response.text(encoding="utf-8")) == payload.decode() + assert finish(response.__aenter__()) is response + assert finish(response.__aexit__(None, None, None)) is None + with pytest.raises(RuntimeError): + await response.bytes() + + # Nor does an unstarted read take the body: a read caches it, and a cached body + # cannot be streamed. Until it ends, the read holds the response. + response = await get() + read = response.bytes() + if cpython: + assert response in gc.get_referents(read) + with pytest.raises(ValueError): + read.throw(ValueError("unread")) + if cpython: + assert response not in gc.get_referents(read) + response.json().close() + response.text().close() + streamer = response.stream() + unread = streamer.__anext__() + if cpython: + assert streamer in gc.get_referents(unread) + unread.close() + streamer.__aexit__(None, None, None).close() + assert bytes(finish(streamer.__anext__())) == payload + with pytest.raises(StopAsyncIteration): + streamer.__anext__().send(None) + assert finish(streamer.__aexit__(None, None, None)) is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("blocking", [False, True], ids=["async", "blocking"]) +async def test_prefetched_stream_keeps_lazy_reads_and_owned_views(blocking): + payload = b"buffered" * 128 + reply = b"HTTP/1.1 200 OK\r\nConnection: close\r\nContent-Length: 1024\r\n\r\n" + payload + + def read_blocking(url): + with wreq.blocking.Client(proxies=[]) as client: + response = client.get(url) + with response.stream() as stream: + views = list(stream) + # Every later read reports the end the same way. + for _ in range(2): + with pytest.raises(StopIteration) as stop: + next(stream) + assert stop.value.args == ("The iterator is exhausted",) + with pytest.raises(RuntimeError): + response.stream() + assert response.status.as_int() == 200 + assert response.content_length == len(payload) + response.close() + return views + + async def read_async(url): + async with wreq.Client(proxies=[]) as client: + response = await client.get(url) + async with response.stream() as stream: + # Creating and discarding a read must not consume a prefetched frame. + unread = stream.__anext__() + unread.close() + views = [view async for view in stream] + with pytest.raises(RuntimeError): + response.stream() + assert response.status.as_int() == 200 + assert response.content_length == len(payload) + await response.close() + return views + + async with local_server() as (url, connections): + task = asyncio.create_task( + asyncio.to_thread(read_blocking, url) if blocking else read_async(url) + ) + _, writer = await asyncio.wait_for(connections.get(), 5) + writer.write(reply) + views = await asyncio.wait_for(task, 5) + + assert all(isinstance(view, memoryview) and view.readonly for view in views) + # Both the stream and response have closed; each view still owns its bytes. + assert b"".join(views) == payload + + @pytest.mark.asyncio async def test_pending_read_holds_the_body_until_it_ends(): head = b"HTTP/1.1 200 OK\r\nContent-Length: 10\r\n\r\n" @@ -270,8 +440,9 @@ async def steps(coroutine): @asynccontextmanager -async def body_server(http2, body): - """Answer every request with `body` and its Content-Length, over HTTP/1 or h2c.""" +async def body_server(http2, body, delay=0): + """Answer every request with `body` and its Content-Length, over HTTP/1 or h2c, + sending the body `delay` seconds after the head.""" length = str(len(body)).encode() handlers, writers = set(), [] @@ -286,7 +457,11 @@ async def accept(reader, writer): if not http2: await reader.readuntil(b"\r\n\r\n") head = b"HTTP/1.1 200 OK\r\nConnection: close\r\nContent-Length: " - writer.write(head + length + b"\r\n\r\n" + body) + writer.write(head + length + b"\r\n\r\n") + if delay: + await writer.drain() + await asyncio.sleep(delay) + writer.write(body) await writer.drain() await reader.read() return @@ -302,7 +477,12 @@ async def accept(reader, writer): elif kind == 1: # `:status: 200` and `content-length`, then the body ending the stream. block = b"\x88\x0f\x0d" + bytes((len(length),)) + length - if body: + if delay: + writer.write(frame(1, 4, stream, block)) + await writer.drain() + await asyncio.sleep(delay) + writer.write(frame(0, 1, stream, body)) + elif body: writer.write( frame(1, 4, stream, block) + frame(0, 1, stream, body) ) @@ -328,12 +508,15 @@ async def accept(reader, writer): @pytest.mark.asyncio @pytest.mark.parametrize("read", ["bytes", "stream"]) +@pytest.mark.parametrize("delay", [0, 0.02], ids=["with-head", "later"]) @pytest.mark.parametrize("body", [b"ok", b""], ids=["body", "empty"]) @pytest.mark.parametrize("http2", [False, True], ids=["http1", "http2"]) -async def test_only_http1_bodies_are_read_on_the_event_loop(http2, body, read): - # A buffered HTTP/1 body is read in the first step of the coroutine. An HTTP/2 body, - # even an empty one, is always read on the runtime, so the coroutine suspends first. - async with body_server(http2, body) as url: +async def test_event_loop_reads_only_bodies_it_may_poll(http2, body, delay, read): + # A buffered HTTP/1 body, or a short HTTP/2 body that arrived with its head and was + # read along with it, is read in the first step of the coroutine. An HTTP/2 body still + # tied to its stream, even an empty one, is read on the runtime, so the coroutine + # suspends first. + async with body_server(http2, body, delay) as url: async with wreq.Client(http2_only=http2, proxies=[]) as client: response = await asyncio.wait_for(client.get(url), 5) assert response.version == (Version.HTTP_2 if http2 else Version.HTTP_11) @@ -345,7 +528,7 @@ async def test_only_http1_bodies_are_read_on_the_event_loop(http2, body, read): else: result, suspended = await steps(response.stream().__anext__()) assert (bytes(result) if result is not None else b"") == body - assert (suspended > 0) == http2 + assert (suspended > 0) == (http2 and delay > 0) @pytest.mark.asyncio diff --git a/tests/runtime_test.py b/tests/runtime_test.py index 09c44ca6..3d7aef5d 100644 --- a/tests/runtime_test.py +++ b/tests/runtime_test.py @@ -1,5 +1,6 @@ import asyncio import base64 +import gc import hashlib import threading import time @@ -142,6 +143,25 @@ async def test_response_and_stream_keep_runtime_alive(scheduler): del stream +@pytest.mark.asyncio +@pytest.mark.parametrize("scheduler", [Scheduler.PER_WORKER, Scheduler.WORK_STEALING]) +async def test_pending_request_keeps_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)) + _, writer = await asyncio.wait_for(connections.get(), 5) + # Only the pending request now owns the runtime. + del client, runtime + gc.collect() + writer.write(b"HTTP/1.1 200 OK\r\nContent-Length: 4\r\n\r\nbody") + await writer.drain() + response = await asyncio.wait_for(task, 5) + del task + assert await response.bytes() == b"body" + del response + + @pytest.mark.asyncio async def test_shared_runtime_cancellation_and_upload(): runtime = Runtime(scheduler=Scheduler.PER_WORKER, workers=2, max_blocking_threads=2) @@ -243,7 +263,8 @@ def settle(): @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.""" + must drive the runtime, except `/buffered`, which can be prefetched. `/hold` waits for + `release`; `arrived` records paths.""" arrived, release, uploads, writers = asyncio.Queue(), asyncio.Event(), [], [] async def accept(reader, writer): @@ -260,8 +281,9 @@ async def accept(reader, writer): 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) + if path != b"/buffered": + await writer.drain() + await asyncio.sleep(0.05) writer.write(path) await writer.drain() except (asyncio.IncompleteReadError, ConnectionError, ValueError): @@ -307,6 +329,9 @@ def run(): caller, readers = threading.get_ident(), [] pending = client.get(url + "/pending") streamed = client.get(url + "/streamed").stream() + buffered = client.get(url + "/buffered").stream() + drained = client.get(url + "/buffered").stream() + assert b"".join(drained) == b"/buffered" def chunks(): # Read on the thread driving the runtime, which refuses blocking calls @@ -316,8 +341,10 @@ def chunks(): 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) + # Refused before a ready frame or the end of a drained body is returned. + for stream in (streamed, buffered, drained): + with pytest.raises(RuntimeError, match="CURRENT_THREAD"): + next(stream) yield b" upload" with client.post(url + "/upload", body=chunks()) as response: @@ -331,7 +358,13 @@ def chunks(): # itself. with pytest.raises(Interrupt): upload_interrupted() - return readers == [caller], uploaded, pending.text(), b"".join(streamed) + return ( + readers == [caller], + uploaded, + pending.text(), + b"".join(streamed), + b"".join(buffered), + ) try: async with path_server() as (url, arrived, release, uploads): @@ -340,6 +373,7 @@ def chunks(): "/upload", "/pending", b"/streamed", + b"/buffered", ) assert uploads == [b"sync upload"] seen = [] @@ -444,14 +478,15 @@ def fetch_after_idle(): @pytest.mark.asyncio +@pytest.mark.parametrize("scheduler", [Scheduler.WORK_STEALING, Scheduler.PER_WORKER]) @pytest.mark.parametrize( ("blocking", "operation"), - [(True, "get"), (False, "websocket"), (True, "websocket")], + [(False, "get"), (True, "get"), (False, "websocket"), (True, "websocket")], ) -async def test_client_close_cancels_requests(blocking, operation): +async def test_client_close_cancels_requests(blocking, operation, scheduler): factory = wreq.blocking.Client if blocking else wreq.Client client = factory( - runtime=custom_runtime(Scheduler.PER_WORKER), + runtime=custom_runtime(scheduler), proxies=[], timeout=timedelta(seconds=10), ) diff --git a/tests/shutdown_test.py b/tests/shutdown_test.py index 54efdc1a..5c87a4b4 100644 --- a/tests/shutdown_test.py +++ b/tests/shutdown_test.py @@ -11,7 +11,7 @@ import os import socket import sys -import threading +import time from types import FunctionType import wreq @@ -23,19 +23,17 @@ url = f"http://127.0.0.1:{listener.getsockname()[1]}/" -wait_lock = threading.Lock() -wait_lock.acquire() - class HoldTeardown: def __init__(self, peer, task): self.peer = peer self.task = task - def __del__(self, write=os.write, wait=wait_lock.acquire): + # Python 3.15 refuses lock acquisition during finalization, so wait with sleep. + def __del__(self, write=os.write, sleep=time.sleep): self.peer.close() write(2, b"teardown: connection closed\\n") - wait(timeout=1) + sleep(1) write(2, b"teardown: wait complete\\n") diff --git a/tests/upload_test.py b/tests/upload_test.py index 699f0bf1..e84c7410 100644 --- a/tests/upload_test.py +++ b/tests/upload_test.py @@ -1,6 +1,7 @@ import asyncio import contextvars import gc +import inspect import threading import pytest @@ -66,6 +67,97 @@ async def chunks(): assert closed.is_set() +@pytest.mark.asyncio +@pytest.mark.parametrize("hook", ["instance", "base", "factory"]) +@pytest.mark.parametrize("keywords", [False, True], ids=["legacy", "kwargs"]) +async def test_upload_task_creation_hooks(hook, keywords, monkeypatch): + loop = asyncio.get_running_loop() + original = loop.create_task + context = contextvars.ContextVar("task_hook", default="missing") + created = [] + + def tracked(coro, **kwargs): + # The loop also closes collected async generators here, which PyPy does late. + if getattr(coro, "__qualname__", None) == "forward": + created.append(coro) + task_context = contextvars.copy_context() + task_context.run(context.set, "wrapped") + kwargs["context"] = task_context + if hook == "factory": + return asyncio.Task(coro, loop=loop, **kwargs) + return original(coro, **kwargs) + + def legacy(coro, *, name=None, context=None): + return tracked(coro, name=name, context=context) + + create = tracked if keywords else legacy + if hook == "instance": + monkeypatch.setattr(loop, "create_task", create) + else: + if keywords: + + def create_on_loop(self, coro, **kwargs): + return tracked(coro, **kwargs) + else: + + def create_on_loop(self, coro, *, name=None, context=None): + return tracked(coro, name=name, context=context) + if hook == "base": + monkeypatch.setattr(asyncio.BaseEventLoop, "create_task", create_on_loop) + else: + factory = loop.get_task_factory() + loop.set_task_factory(create_on_loop) + + async def chunks(): + assert context.get() == "wrapped" + yield b"body" + + try: + async with local_server() as (url, connections), wreq.Client(proxies=[]) as client: + request = asyncio.create_task(client.post(url, body=chunks())) + try: + await asyncio.sleep(0) + assert len(created) == 1 + reader, writer = await asyncio.wait_for(connections.get(), 5) + assert await asyncio.wait_for(read_chunked(reader), 5) == b"body" + writer.write(b"HTTP/1.1 200 OK\r\nContent-Length: 0\r\n\r\n") + await writer.drain() + await (await asyncio.wait_for(request, 5)).close() + assert context.get() == "missing" + finally: + request.cancel() + await asyncio.gather(request, return_exceptions=True) + finally: + if hook == "factory": + loop.set_task_factory(factory) + + +@pytest.mark.asyncio +@pytest.mark.skipif( + inspect.getfullargspec(asyncio.BaseEventLoop.create_task).varkw is None, + reason="loop create_task does not pass eager_start", +) +async def test_upload_starts_eagerly_on_standard_loop(): + started = False + + async def chunks(): + nonlocal started + started = True + yield b"body" + + async with local_server() as (url, connections), wreq.Client(proxies=[]) as client: + request = asyncio.create_task(client.post(url, body=chunks())) + # Let the request's first poll extract the body, before a separately queued + # forwarding task could run on the following loop iteration. + await asyncio.sleep(0) + assert started + reader, writer = await asyncio.wait_for(connections.get(), 5) + assert await asyncio.wait_for(read_chunked(reader), 5) == b"body" + writer.write(b"HTTP/1.1 200 OK\r\nContent-Length: 0\r\n\r\n") + await writer.drain() + await (await asyncio.wait_for(request, 5)).close() + + @pytest.mark.asyncio @pytest.mark.parametrize("size", [0, 1 << 24], ids=["empty", "queued_tail"]) async def test_finish_delivers_queued_chunks(size): @@ -247,3 +339,38 @@ async def body(): finally: server.close() await server.wait_closed() + + +@pytest.mark.asyncio +async def test_upload_task_code_replacement_keeps_hook_arguments(monkeypatch): + loop = asyncio.get_running_loop() + original = loop.create_task + calls = [] + + def create(coro, **kwargs): + if getattr(coro, "__qualname__", None) == "forward": + calls.append(dict(kwargs)) + return original(coro, **kwargs) + + def legacy(coro, *, name=None, context=None): + if getattr(coro, "__qualname__", None) == "forward": + calls.append({}) + return original(coro, name=name, context=context) + + monkeypatch.setattr(loop, "create_task", create) + async with local_server() as (url, connections), wreq.Client(proxies=[]) as client: + for code in (create.__code__, legacy.__code__): + create.__code__ = code + create.__kwdefaults__ = legacy.__kwdefaults__ + + async def chunks(): + yield b"body" + + task = original(client.post(url, body=chunks())) + reader, writer = await asyncio.wait_for(connections.get(), 5) + assert await asyncio.wait_for(read_chunked(reader), 5) == b"body" + writer.write(b"HTTP/1.1 200 OK\r\nContent-Length: 0\r\n\r\n") + await writer.drain() + await (await asyncio.wait_for(task, 5)).close() + assert len(calls) == 2 + assert calls[-1] == {}