diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 91f88646..d68d4c36 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -72,6 +72,8 @@ jobs: run: | uv pip install .[test] uv pip install maturin + # uvloop is optional and can lag a new Python release; never build it from source. + uv pip install --only-binary uvloop uvloop || echo "::warning::no uvloop wheel; skipping uvloop tests" - name: Build wheels run: | source .venv/bin/activate @@ -366,7 +368,8 @@ jobs: timeout-minutes: 10 run: | python -m pip install --no-index --find-links dist wreq - python -m pip install pytest pytest-asyncio uvloop + python -m pip install pytest pytest-asyncio + python -m pip install --only-binary uvloop uvloop || echo "::warning::no uvloop wheel; skipping uvloop tests" python -m pytest -v -o faulthandler_timeout=60 tests/aio_test.py tests/cancellation_test.py tests/upload_test.py tests/runtime_test.py - name: Upload wheels uses: actions/upload-artifact@v7 diff --git a/Cargo.lock b/Cargo.lock index 1a219e01..45288969 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1178,8 +1178,7 @@ checksum = "27b02d87554356db9e9a873add8782d4ea6e3e58ea071a9adb9a2e8ddb884a8b" [[package]] name = "netty" version = "0.2.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a6d4d60c6b6193ac462e05a1c8ae826a624eff20368e208e59101f5a1fbfc055" +source = "git+https://github.com/0x676e67/netty?rev=f8652ddf5e386d66b7f03cceedbc6b319d849762#f8652ddf5e386d66b7f03cceedbc6b319d849762" dependencies = [ "bytes", "futures-channel", diff --git a/Cargo.toml b/Cargo.toml index 0c9cf0d5..b96ba3f1 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -36,9 +36,8 @@ pyo3 = { version = "0.29.0", features = [ "multiple-pymethods", "extension-module", "generate-import-lib", - "experimental-async", ] } -futures-util = { version = "0.3.33", default-features = false } +futures-util = { version = "0.3.33", default-features = false, features = ["std"] } serde = { version = "1.0.228", features = ["derive"] } serde_json = "1" serde_urlencoded = "0.7.1" @@ -84,6 +83,8 @@ wreq = { git = "https://github.com/0x676e67/wreq", rev = "54fdc7c9e0c1290e2f70f3 btls = { git = "https://github.com/0x676e67/btls", rev = "ffaf0e8c23d455af4e306986ce81bb57ebd03bc2" } btls-sys = { git = "https://github.com/0x676e67/btls", rev = "ffaf0e8c23d455af4e306986ce81bb57ebd03bc2" } tokio-btls = { git = "https://github.com/0x676e67/btls", rev = "ffaf0e8c23d455af4e306986ce81bb57ebd03bc2" } +# main (0x676e67/netty#145) +netty = { git = "https://github.com/0x676e67/netty", rev = "f8652ddf5e386d66b7f03cceedbc6b319d849762" } [profile.release] codegen-units = 1 diff --git a/docs/source/getting-started/quickstart.md b/docs/source/getting-started/quickstart.md index e7679a17..55849262 100644 --- a/docs/source/getting-started/quickstart.md +++ b/docs/source/getting-started/quickstart.md @@ -88,9 +88,9 @@ print(data) ```python import hashlib -view = await response.bytes() -await response.close() -print(view.readonly) # True; closing the response does not invalidate the view +async with response: + view = await response.bytes() +print(view.readonly) # True; releasing the response does not invalidate the view print(hashlib.sha256(view).hexdigest()) # Reads the buffer directly ``` diff --git a/docs/source/guide/basic.md b/docs/source/guide/basic.md index c284043f..e6afe24e 100644 --- a/docs/source/guide/basic.md +++ b/docs/source/guide/basic.md @@ -184,6 +184,6 @@ async def main(): Data chunks are read-only `memoryview` objects; trailer frames are `HeaderMap` objects. The example writes data directly through the buffer protocol. Each view stays valid after the stream is closed. See [Binary data](../getting-started/quickstart.md#binary-data) for copying and releasing views. -Finish a body read before closing the response. If a read task is still running, cancel it and await its cancellation first: `response.close()` does not cancel active reads or guarantee an immediate socket shutdown. `response.stream()` transfers the body to the streamer, so manage the streamer with its own context manager as shown above. +Finish a body read before closing the response. If a read task is still running, cancel it and await its cancellation first: `response.close()` does not cancel active reads or guarantee an immediate socket shutdown. It also marks the connection as non-reusable; leaving `async with response:` instead releases the body and keeps a fully read connection reusable. `response.stream()` transfers the body to the streamer, so manage the streamer with its own context manager as shown above. --- diff --git a/docs/source/guide/blocking.md b/docs/source/guide/blocking.md index 69008ac1..67383a9f 100644 --- a/docs/source/guide/blocking.md +++ b/docs/source/guide/blocking.md @@ -210,4 +210,4 @@ if __name__ == "__main__": Data chunks are read-only `memoryview` objects that stay valid after the stream is closed. `resp.bytes()` returns the same type. Pass views directly to APIs that accept the buffer protocol; use `bytes(view)` or `view.tobytes()` only when you need a copy. -Do not close a response while another thread is reading its body. `resp.close()` discards the retained body and marks its connection as non-reusable, but does not guarantee an immediate socket shutdown or interrupt an active read. A body transferred by `resp.stream()` belongs to the streamer and needs its own context manager, as shown above. +Do not close a response while another thread is reading its body. `resp.close()` discards the retained body and marks its connection as non-reusable, but does not guarantee an immediate socket shutdown or interrupt an active read. Leaving `with resp:` instead keeps a fully read connection reusable. A body transferred by `resp.stream()` belongs to the streamer and needs its own context manager, as shown above. diff --git a/pyproject.toml b/pyproject.toml index 301b556c..e72ad38f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -41,7 +41,6 @@ test = [ "pytest", "pytest-asyncio", "pytest-rerunfailures", - "uvloop; sys_platform != 'win32' and platform_python_implementation == 'CPython'", ] docs = [ "zensical", diff --git a/python/wreq/_compat.py b/python/wreq/_compat.py index 9b610e7d..d1cc1f8d 100644 --- a/python/wreq/_compat.py +++ b/python/wreq/_compat.py @@ -1,4 +1,4 @@ -"""Compatibility for PyPy's legacy exception delegation to PyO3 coroutines.""" +"""Compatibility for PyPy's legacy exception delegation to wreq coroutines.""" from types import TracebackType @@ -8,7 +8,8 @@ def _install(coroutine_type): if getattr(original, "_wreq_pypy_throw_compat", False): return - # Remove this shim once PyO3 accepts throw(type, value, traceback). + # Coroutine.throw takes one exception, as asyncio passes; PyPy delegates + # throws from awaiting generators as the legacy (type, value, traceback). def throw(self, exc, value=None, traceback=None): if traceback is not None and type(traceback) is not TracebackType: raise TypeError("throw() third argument must be a traceback object") diff --git a/python/wreq/blocking.py b/python/wreq/blocking.py index c5770812..b8e21314 100644 --- a/python/wreq/blocking.py +++ b/python/wreq/blocking.py @@ -123,7 +123,12 @@ def close(self) -> None: """ def __enter__(self) -> Any: ... - def __exit__(self, _exc_type: Any, _exc_value: Any, _traceback: Any) -> None: ... + def __exit__(self, _exc_type: Any, _exc_value: Any, _traceback: Any) -> None: + r""" + 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. + """ + def __str__(self) -> str: ... @@ -205,7 +210,11 @@ def close( """ def __enter__(self) -> Any: ... - def __exit__(self, _exc_type: Any, _exc_value: Any, _traceback: Any) -> None: ... + def __exit__(self, _exc_type: Any, _exc_value: Any, _traceback: Any) -> None: + r""" + Close the WebSocket connection without a close code or reason, unless already closed. + """ + def __str__(self) -> str: ... @@ -489,7 +498,10 @@ def get( ... def __enter__(self) -> Any: ... - def __exit__(self, _exc_type: Any, _exc_value: Any, _traceback: Any) -> None: ... + def __exit__(self, _exc_type: Any, _exc_value: Any, _traceback: Any) -> None: + r""" + Close the client like `close()`: cancel pending requests and reject new ones. + """ def delete( diff --git a/python/wreq/wreq.py b/python/wreq/wreq.py index dc7ab0ce..44e49543 100644 --- a/python/wreq/wreq.py +++ b/python/wreq/wreq.py @@ -318,7 +318,10 @@ def __next__(self) -> memoryview | HeaderMap: ... def __enter__(self) -> Any: ... - def __exit__(self, _exc_type: Any, _exc_value: Any, _traceback: Any) -> None: ... + def __exit__(self, _exc_type: Any, _exc_value: Any, _traceback: Any) -> None: + r""" + Release the body and end any pending read; returned views stay valid. + """ def __aiter__(self) -> "Streamer": ... @@ -326,9 +329,10 @@ async def __anext__(self) -> memoryview | HeaderMap: ... async def __aenter__(self) -> Any: ... - async def __aexit__( - self, _exc_type: Any, _exc_value: Any, _traceback: Any - ) -> None: ... + async def __aexit__(self, _exc_type: Any, _exc_value: Any, _traceback: Any) -> None: + r""" + Release the body and end any pending read; returned views stay valid. + """ class Response: @@ -348,7 +352,6 @@ async def main(): print("Response URL: ", response.url) print("Headers: ", response.headers) print("Content-Length: ", response.content_length) - print("Encoding: ", response.encoding) print("Remote Address: ", response.remote_addr) text_content = await response.text() @@ -449,9 +452,12 @@ async def close(self) -> None: """ async def __aenter__(self) -> Any: ... - async def __aexit__( - self, _exc_type: Any, _exc_value: Any, _traceback: Any - ) -> Any: ... + async def __aexit__(self, _exc_type: Any, _exc_value: Any, _traceback: Any) -> Any: + r""" + 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. + """ + def __str__(self) -> str: ... @@ -519,8 +525,12 @@ async def close( Close the WebSocket connection. """ - def __aenter__(self) -> Any: ... - def __aexit__(self, _exc_type: Any, _exc_value: Any, _traceback: Any) -> Any: ... + async def __aenter__(self) -> Any: ... + async def __aexit__(self, _exc_type: Any, _exc_value: Any, _traceback: Any) -> Any: + r""" + Close the WebSocket connection without a close code or reason, unless already closed. + """ + def __str__(self) -> str: ... @@ -1433,9 +1443,10 @@ async def main(): ... async def __aenter__(self) -> Any: ... - async def __aexit__( - self, _exc_type: Any, _exc_value: Any, _traceback: Any - ) -> Any: ... + async def __aexit__(self, _exc_type: Any, _exc_value: Any, _traceback: Any) -> Any: + r""" + Close the client like `close()`: cancel pending requests and reject new ones. + """ async def delete( diff --git a/src/aio.rs b/src/aio.rs index 1b1ef28c..f65f7381 100644 --- a/src/aio.rs +++ b/src/aio.rs @@ -1,9 +1,11 @@ -//! Awaitables that resume asyncio tasks without attaching Tokio threads to Python. +//! Awaitables that drive Rust futures from asyncio tasks. //! -//! Rust futures complete on Tokio workers, which never take the GIL. A wake marks -//! its [`Coroutine`] ready and queues it on the [`Port`] of the awaiting event loop, -//! ringing a socket the loop polls only when the queue was empty. The loop thread -//! then resolves every queued asyncio future in one batch. +//! 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. mod coroutine; mod port; diff --git a/src/aio/port.rs b/src/aio/port.rs index a95ce0b5..dba9f0a3 100644 --- a/src/aio/port.rs +++ b/src/aio/port.rs @@ -7,7 +7,7 @@ use std::{ }, }; -use pyo3::{intern, prelude::*, sync::PyOnceLock}; +use pyo3::{PyTraverseError, PyVisit, intern, prelude::*, sync::PyOnceLock}; use super::coroutine::Slot; @@ -15,7 +15,8 @@ use super::coroutine::Slot; pub(crate) struct Port { queue: Mutex>>, bell: Bell, - /// Cleared once the loop drops its drain; later wakes are discarded. + /// Cleared once the loop drops its drain or keeper, or a scheduled drain cannot be + /// queued; later wakes are discarded. open: AtomicBool, } @@ -23,7 +24,8 @@ 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: the waking thread schedules the drain, + /// and a [`Keeper`] holds the port until the loop closes. Loop(Py), } @@ -34,9 +36,14 @@ type Socket = std::os::unix::net::UnixStream; type Socket = std::net::TcpStream; /// Resolves the asyncio futures of queued coroutines on the loop thread. -#[pyclass(frozen)] +#[pyclass(frozen, module = "wreq")] struct Drain(Arc); +/// Owns a scheduled-drain port through a renewing timer, which the loop discards when +/// closed. It reports the port's loop reference so an unclosed loop can be collected. +#[pyclass(module = "wreq")] +struct Keeper(Option>); + thread_local! { /// Ports of the loops that ran on this thread, keyed by loop address. static PORTS: RefCell)>> = const { RefCell::new(Vec::new()) }; @@ -57,7 +64,8 @@ impl Port { .bind(py) .call0()?; - // A loop drops its drain when closed or freed, so a reused address never matches. + // A loop drops its drain or keeper when closed or freed, so a reused address + // never matches. let key = event_loop.as_ptr() as usize; let cached = PORTS.with_borrow_mut(|ports| { ports.retain(|(_, port)| port.upgrade().is_some_and(|port| port.is_open())); @@ -92,7 +100,7 @@ impl Port { if watched.is_ok() { return Ok(port); } - return Ok(Self::new(Bell::Loop(event_loop.clone().unbind()))); + return Self::scheduled(event_loop); } let fd = raw(&rx); @@ -106,7 +114,15 @@ impl Port { return Ok(port); } } - Ok(Self::new(Bell::Loop(event_loop.clone().unbind()))) + Self::scheduled(event_loop) + } + + /// Open a port whose wakes schedule drains on `event_loop`. + fn scheduled(event_loop: &Bound<'_, PyAny>) -> PyResult> { + let port = Self::new(Bell::Loop(event_loop.clone().unbind())); + let keeper = Bound::new(event_loop.py(), Keeper(Some(port.clone())))?; + Keeper::arm(&keeper, event_loop)?; + Ok(port) } fn new(bell: Bell) -> Arc { @@ -182,6 +198,13 @@ impl Port { #[pymethods] impl Drain { + fn __repr__(&self) -> &'static str { + match self.0.bell { + Bell::Socket { .. } => "", + Bell::Loop(_) => "", + } + } + fn __call__(&self, py: Python<'_>) -> PyResult<()> { let port = &self.0; // Clear the bell before taking the queue so no ring is lost; a short read @@ -197,12 +220,13 @@ impl Drain { continue; }; // Resume the task awaiting the waiter, unless it was already cancelled. - let future = waiter.bind(py); - let done = future + let done = waiter + .bind(py) .call_method0(intern!(py, "done")) .and_then(|done| done.is_truthy()); let released = match done { - Ok(false) => future + Ok(false) => waiter + .bind(py) .call_method1(intern!(py, "set_result"), (py.None(),)) .map(drop), done => done.map(drop), @@ -229,6 +253,58 @@ impl Drop for Drain { } } +// ===== impl Keeper ===== + +impl Keeper { + /// Seconds between renewals; any delay works, as only a closing loop drops the timer. + const RENEW_SECS: f64 = 3600.0; + + fn arm(keeper: &Bound<'_, Keeper>, event_loop: &Bound<'_, PyAny>) -> PyResult<()> { + let py = keeper.py(); + event_loop + .call_method1(intern!(py, "call_later"), (Self::RENEW_SECS, keeper)) + .map(drop) + } +} + +#[pymethods] +impl Keeper { + fn __call__(slf: &Bound<'_, Self>) -> PyResult<()> { + let event_loop = match slf.borrow().0.as_deref() { + Some(Port { + bell: Bell::Loop(event_loop), + .. + }) => event_loop.clone_ref(slf.py()), + _ => return Ok(()), + }; + Self::arm(slf, event_loop.bind(slf.py())) + } + + fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { + // The keeper is the only object reporting this reference, so it is counted once. + if let Some(Port { + bell: Bell::Loop(event_loop), + .. + }) = self.0.as_deref() + { + visit.call(event_loop)?; + } + Ok(()) + } + + fn __clear__(&mut self) { + if let Some(port) = self.0.take() { + port.close(); + } + } +} + +impl Drop for Keeper { + fn drop(&mut self) { + self.__clear__(); + } +} + #[cfg(unix)] fn socket_pair() -> std::io::Result<(Socket, Socket)> { let (tx, rx) = Socket::pair()?; @@ -306,12 +382,16 @@ class Watch: def watch(loop, fileno, drain): - sock = socket.socket(fileno=fileno) + sock = None try: + sock = socket.socket(fileno=fileno) # A first receive that fails lets the caller fall back to scheduled drains. Watch(loop, sock, drain).arm() except BaseException: - sock.close() + if sock is None: + socket.close(fileno) + else: + sock.close() raise ", c"wreq/_proactor.py", diff --git a/src/client.rs b/src/client.rs index f02b6408..3e527614 100644 --- a/src/client.rs +++ b/src/client.rs @@ -670,6 +670,7 @@ impl Client { aio::ready("Client.__aenter__", slf) } + /// Close the client like `close()`: cancel pending requests and reject new ones. fn __aexit__<'py>( &self, py: Python<'py>, @@ -844,6 +845,7 @@ impl BlockingClient { slf } + /// Close the client like `close()`: cancel pending requests and reject new ones. fn __exit__<'py>( &self, _py: Python<'py>, diff --git a/src/client/body/stream.rs b/src/client/body/stream.rs index 76e31c41..cd59d04d 100644 --- a/src/client/body/stream.rs +++ b/src/client/body/stream.rs @@ -1,5 +1,6 @@ use std::{ future::Future, + panic::AssertUnwindSafe, pin::{Pin, pin}, sync::{ Arc, Mutex, MutexGuard, PoisonError, @@ -11,8 +12,14 @@ use std::{ use bytes::Bytes; use futures_util::{FutureExt, Stream, future::poll_fn}; use http_body_util::BodyExt; -use pyo3::{exceptions::PyStopIteration, intern, prelude::*, sync::PyOnceLock}; +use pyo3::{ + exceptions::{PyRuntimeError, PyStopIteration}, + intern, + prelude::*, + sync::PyOnceLock, +}; use tokio::{ + runtime::Handle, sync::{ Notify, mpsc::{self, error::TryRecvError}, @@ -69,7 +76,6 @@ struct PyAsyncStream { struct Sender(mpsc::Sender>>); /// A response stream yielding read-only memoryviews and any trailing headers. -#[derive(Clone)] #[pyclass(subclass, frozen, skip_from_py_object)] pub struct Streamer(Arc, Runtime); @@ -80,7 +86,7 @@ struct Reader { state: Mutex, arrived: Arc, /// Buffered frames returned since `__anext__` last yielded to the event loop. - ready: AtomicUsize, + since_yield: AtomicUsize, } enum State { @@ -98,7 +104,9 @@ struct NotifyOnDrop(Arc); // ===== impl Streamer ===== impl Streamer { - /// Frames buffered ahead of Python, about 128 KiB of full TLS records. + /// 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. const READ_AHEAD: usize = 8; /// Buffered frames returned before `__anext__` yields to the event loop once. @@ -110,7 +118,7 @@ impl Streamer { let reader = Reader { state: Mutex::new(State::Idle(Box::new(resp))), arrived: Arc::new(Notify::new()), - ready: AtomicUsize::new(0), + since_yield: AtomicUsize::new(0), }; Streamer(Arc::new(reader), runtime) } @@ -120,12 +128,21 @@ impl Streamer { tx: mpsc::Sender>, arrived: Arc, ) { - // Readers wake only after `pump` drops the sender, so they see the end. - let _end = NotifyOnDrop(arrived.clone()); - Self::pump(resp, tx, &arrived).await; + let end = NotifyOnDrop(arrived); + // Without this a panic would drop `tx` and read as the end of the body. + if AssertUnwindSafe(Self::pump(resp, &tx, &end.0)) + .catch_unwind() + .await + .is_err() + { + let panicked = Err(PyRuntimeError::new_err("response body reader panicked")); + let _ = burst(tx.send(panicked), &end.0).await; + } + // Readers wake only after the sender drops, so they see the end. + drop(tx); } - async fn pump(mut resp: wreq::Response, tx: mpsc::Sender>, arrived: &Notify) { + async fn pump(mut resp: wreq::Response, tx: &mpsc::Sender>, arrived: &Notify) { while let Some(frame) = burst(resp.frame(), arrived).await { let frame = match frame.map(|frame| frame.into_data()) { Ok(Ok(bytes)) => Ok(Frame::Bytes(PyBuffer::from(bytes))), @@ -134,7 +151,7 @@ impl Streamer { Err(_) => continue, }, Err(err) => { - // A failed body may leave a stalled HTTP/2 connection behind. + // Conservatively keep the connection out of the pool after a body error. resp.forbid_recycle(); Err(Error::Library(err).into()) } @@ -202,6 +219,7 @@ impl Streamer { slf } + /// Release the body and end any pending read; returned views stay valid. fn __exit__<'py>( &self, py: Python, @@ -227,8 +245,8 @@ impl Streamer { let this = slf.get(); // Buffered frames complete without suspending; yield to the event loop // periodically so timeouts, cancellation and other tasks run. - if this.0.ready.fetch_add(1, Ordering::Relaxed) >= Self::YIELD_EVERY { - this.0.ready.store(0, Ordering::Relaxed); + if this.0.since_yield.fetch_add(1, Ordering::Relaxed) >= Self::YIELD_EVERY { + this.0.since_yield.store(0, Ordering::Relaxed); aio::yield_now().await; } loop { @@ -238,7 +256,7 @@ impl Streamer { if let Some(frame) = this.try_next(|| Error::StopAsyncIteration) { return frame; } - this.0.ready.store(0, Ordering::Relaxed); + this.0.since_yield.store(0, Ordering::Relaxed); arrived.await; } }) @@ -248,6 +266,7 @@ impl Streamer { aio::ready("Streamer.__aenter__", slf) } + /// Release the body and end any pending read; returned views stay valid. fn __aexit__<'py>( &self, py: Python<'py>, @@ -461,8 +480,11 @@ impl Drop for PyAsyncStream { fn drop(&mut self) { self.rx.close(); if let Some((task, event_loop)) = self.task.take() { - // Body drop can run on Tokio: acquire the interpreter on a blocking thread. - crate::runtime::get().handle().spawn_blocking(move || { + // Body drop can run on Tokio: cancel from a blocking thread, preferring the current + // runtime so a drop on a client's own runtime does not start the shared one. + let handle = + Handle::try_current().unwrap_or_else(|_| crate::runtime::get().handle().clone()); + handle.spawn_blocking(move || { Python::try_attach(|py| { if let Ok(cancel) = task.bind(py).getattr(intern!(py, "cancel")) { let _ = event_loop.call_method1( diff --git a/src/client/resp/http.rs b/src/client/resp/http.rs index f34e429d..4dc772b7 100644 --- a/src/client/resp/http.rs +++ b/src/client/resp/http.rs @@ -92,8 +92,8 @@ impl Response { match Arc::into_inner(arc) { Some(Body::Streamable(stream)) => { return Box::pin(async move { - // A failed or cancelled read may leave a stalled HTTP/2 - // connection behind, so only a complete body keeps it reusable. + // Conservatively keep the connection out of the pool unless the + // body is read in full. let mut guard = RecycleGuard(Some(parts)); let bytes = stream .collect() @@ -473,6 +473,8 @@ impl BlockingResponse { 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, py: Python<'py>, @@ -480,7 +482,6 @@ impl BlockingResponse { _exc_value: &Bound<'py, PyAny>, _traceback: &Bound<'py, PyAny>, ) { - // Like `__aexit__`, release the body and leave reuse to the protocol. py.detach(|| self.0.destroy()) } } diff --git a/src/client/resp/ws.rs b/src/client/resp/ws.rs index bd81a459..4c0cd3bc 100644 --- a/src/client/resp/ws.rs +++ b/src/client/resp/ws.rs @@ -166,6 +166,7 @@ impl WebSocket { aio::ready("WebSocket.__aenter__", slf) } + /// Close the WebSocket connection without a close code or reason, unless already closed. fn __aexit__<'py>( &self, py: Python<'py>, @@ -177,7 +178,7 @@ impl WebSocket { py, "WebSocket.__aexit__", &self.runtime, - cmd::close(self.cmd.clone(), None, None), + cmd::close_on_exit(self.cmd.clone()), ) } } @@ -279,6 +280,7 @@ impl BlockingWebSocket { slf } + /// Close the WebSocket connection without a close code or reason, unless already closed. fn __exit__<'py>( &self, py: Python<'py>, @@ -286,7 +288,9 @@ impl BlockingWebSocket { _exc_value: &Bound<'py, PyAny>, _traceback: &Bound<'py, PyAny>, ) -> PyResult<()> { - self.close(py, None, None) + py.detach(|| { + crate::client::nogil::block_on(&self.0.runtime, cmd::close_on_exit(self.0.cmd.clone())) + }) } } diff --git a/src/client/resp/ws/cmd.rs b/src/client/resp/ws/cmd.rs index 6f381cf7..ae68556a 100644 --- a/src/client/resp/ws/cmd.rs +++ b/src/client/resp/ws/cmd.rs @@ -93,16 +93,21 @@ pub async fn task(ws: WebSocket, mut cmd: UnboundedReceiver) { } } Command::Close(code, reason, tx) => { - let code = code - .map(ws::message::CloseCode::from) - .unwrap_or(ws::message::CloseCode::NORMAL); let reason = reason .map(|reason| reason.0) .map(Utf8Bytes::try_from) .transpose(); + // A reason requires a code (RFC 6455 ยง5.5.1), so a lone reason closes normally. let close_frame = match reason { - Ok(Some(reason)) => Some(ws::message::CloseFrame { code, reason }), + Ok(reason) if code.is_some() || reason.is_some() => { + Some(ws::message::CloseFrame { + code: code + .map(ws::message::CloseCode::from) + .unwrap_or(ws::message::CloseCode::NORMAL), + reason: reason.unwrap_or_default(), + }) + } _ => None, }; @@ -127,7 +132,9 @@ pub async fn recv( cmd: UnboundedSender, timeout: Option, ) -> PyResult> { - send_command(cmd, |tx| Command::Recv(timeout, tx)).await? + send_command(cmd, |tx| Command::Recv(timeout, tx)) + .await + .ok_or(Error::WebSocketDisconnected)? } /// Sends a [`Command::Send`] to the background task to transmit a message over the WebSocket. @@ -135,7 +142,9 @@ pub async fn recv( /// Returns Ok if the message was sent successfully, or an error otherwise. #[inline] pub async fn send(cmd: UnboundedSender, message: Message) -> PyResult<()> { - send_command(cmd, |tx| Command::Send(message, tx)).await? + send_command(cmd, |tx| Command::Send(message, tx)) + .await + .ok_or(Error::WebSocketDisconnected)? } /// Send as [`Command::SendMany`] to the background task to transmit multiple messages over the @@ -147,7 +156,9 @@ pub async fn send_all(cmd: UnboundedSender, messages: Vec) -> if messages.is_empty() { return Ok(()); } - send_command(cmd, |tx| Command::SendMany(messages, tx)).await? + send_command(cmd, |tx| Command::SendMany(messages, tx)) + .await + .ok_or(Error::WebSocketDisconnected)? } /// Sends a [`Command::Close`] to the background task to gracefully close the WebSocket connection. @@ -159,18 +170,29 @@ pub async fn close( code: Option, reason: Option, ) -> PyResult<()> { - send_command(cmd, |tx| Command::Close(code, reason, tx)).await? + send_command(cmd, |tx| Command::Close(code, reason, tx)) + .await + .ok_or(Error::WebSocketDisconnected)? +} + +/// Closes the WebSocket like [`close`], treating an already closed connection as done, as a +/// context manager exit does. +#[inline] +pub async fn close_on_exit(cmd: UnboundedSender) -> PyResult<()> { + send_command(cmd, |tx| Command::Close(None, None, tx)) + .await + .unwrap_or(Ok(())) } +/// Run a command on the background task, or return `None` once it has ended. async fn send_command( cmd: UnboundedSender, make: impl FnOnce(oneshot::Sender) -> Command, -) -> PyResult { +) -> Option { if cmd.is_closed() { - return Err(Error::WebSocketDisconnected.into()); + return None; } let (tx, rx) = oneshot::channel(); - cmd.send(make(tx)) - .map_err(|_| Error::WebSocketDisconnected)?; - Ok(rx.await.map_err(|_| Error::WebSocketDisconnected)?) + cmd.send(make(tx)).ok()?; + rx.await.ok() } diff --git a/tests/aio_test.py b/tests/aio_test.py index d09ab0b6..83bec3f3 100644 --- a/tests/aio_test.py +++ b/tests/aio_test.py @@ -2,6 +2,7 @@ import gc import socket import sys +import time import weakref import pytest @@ -11,7 +12,11 @@ class NoReaderLoop(asyncio.SelectorEventLoop): """A loop that cannot watch the wake socket, so wakes are scheduled thread-safely.""" + ports = 0 + def add_reader(self, *args): + # Each wake port first tries to register its socket. + self.ports += 1 raise NotImplementedError @@ -63,7 +68,8 @@ async def serve(reader, writer): finally: writer.close() - server = await asyncio.start_server(serve, "127.0.0.1", 0) + # The default backlog of 100 could overflow with 300 concurrent stream connections. + server = await asyncio.start_server(serve, "127.0.0.1", 0, backlog=512) url = f"http://127.0.0.1:{server.sockets[0].getsockname()[1]}" try: async with wreq.Client(proxies=[]) as client: @@ -118,6 +124,8 @@ def test_wakes_resume_tasks_on_event_loops(new_loop): loop = new_loop() try: loop.run_until_complete(asyncio.wait_for(exchange(), 30)) + # A loop keeps one wake port across all of its awaits. + assert getattr(loop, "ports", 1) == 1 finally: loop.close() @@ -126,15 +134,23 @@ def test_wakes_resume_tasks_on_event_loops(new_loop): sys.implementation.name != "cpython", reason="PyPy's cpyext does not collect cycles through extension objects", ) -@pytest.mark.parametrize("new_loop", [p for p in LOOPS if p.id != "no-reader"]) -def test_closed_loop_releases_pending_requests(new_loop): - # The listener accepts connections into its backlog but never answers. +@pytest.mark.parametrize("new_loop", LOOPS) +@pytest.mark.parametrize("woken", [False, True], ids=["pending", "woken"]) +def test_closed_loop_releases_pending_requests(new_loop, woken): + # The listener accepts connections into its backlog and answers only when woken. with socket.create_server(("127.0.0.1", 0)) as server: url = f"http://127.0.0.1:{server.getsockname()[1]}/" client = wreq.Client(proxies=[]) loop = new_loop() task = loop.create_task(client.get(url)) loop.run_until_complete(asyncio.sleep(0.2)) + if woken: + # Answer while the loop is stopped, so its wake is queued but never runs. + conn, _ = server.accept() + conn.recv(65536) + conn.sendall(b"HTTP/1.1 200 OK\r\nContent-Length: 0\r\n\r\n") + time.sleep(0.2) + conn.close() loop.close() # Once its loop is closed, a task left pending can be collected with its request. ref = weakref.ref(task) diff --git a/tests/response_test.py b/tests/response_test.py index 62a2175c..29e042fc 100644 --- a/tests/response_test.py +++ b/tests/response_test.py @@ -162,10 +162,17 @@ async def test_stream_read_ahead_yields_and_closes_waiting_readers(): def serve(listener): for size in (64 << 20, None): - conn, _ = listener.accept() + try: + conn, _ = listener.accept() + except OSError: + return with conn: - while b"\r\n\r\n" not in conn.recv(65536): - pass + head = b"" + while b"\r\n\r\n" not in head: + data = conn.recv(65536) + if not data: + return + head += data length = size or 1 << 20 conn.sendall(b"HTTP/1.1 200 OK\r\nContent-Length: %d\r\n\r\n" % length) try: @@ -176,6 +183,8 @@ def serve(listener): stalled.wait(10) listener = socket.create_server(("127.0.0.1", 0)) + # Closing the listener does not wake a blocked accept after a failed test. + listener.settimeout(10) url = f"http://127.0.0.1:{listener.getsockname()[1]}/" thread = threading.Thread(target=serve, args=(listener,), daemon=True) thread.start() diff --git a/tests/upload_test.py b/tests/upload_test.py index 4288132c..cbf79917 100644 --- a/tests/upload_test.py +++ b/tests/upload_test.py @@ -88,6 +88,38 @@ async def chunks(): await asyncio.wait_for(closed.wait(), 5) +@pytest.mark.asyncio +@pytest.mark.parametrize("action", ["cancel", "early_response"]) +async def test_abandoned_upload_to_stalled_peer(action): + # The peer reads only the head, so the upload stalls with the socket buffers full. + closed = asyncio.Event() + + async def chunks(): + try: + while True: + yield b"x" * 65536 + finally: + closed.set() + + async with local_server() as (url, connections), wreq.Client(proxies=[]) as client: + task = asyncio.create_task(client.post(url, body=chunks())) + _, writer = await asyncio.wait_for(connections.get(), 5) + await asyncio.sleep(0.5) + if action == "cancel": + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + else: + # Close an early response whose body never completes, mid-upload. + writer.write( + b"HTTP/1.1 413 Payload Too Large\r\nContent-Length: 10\r\n\r\nab" + ) + await writer.drain() + response = await asyncio.wait_for(task, 5) + await response.close() + await asyncio.wait_for(closed.wait(), 5) + + @pytest.mark.asyncio @pytest.mark.parametrize("action", ["cancel", "close_client"]) async def test_upload_cancellation(action): diff --git a/tests/websocket_test.py b/tests/websocket_test.py new file mode 100644 index 00000000..39cb19f9 --- /dev/null +++ b/tests/websocket_test.py @@ -0,0 +1,91 @@ +import asyncio +import base64 +import hashlib +import struct +import threading + +import pytest +import wreq +import wreq.blocking + +GUID = b"258EAFA5-E914-47DA-95CA-C5AB0DC85B11" + + +async def read_close(reader): + """Read one masked client frame and return its close code and reason.""" + opcode, size = await reader.readexactly(2) + assert opcode == 0x88 + mask = await reader.readexactly(4) + payload = bytes( + b ^ mask[i % 4] for i, b in enumerate(await reader.readexactly(size & 0x7F)) + ) + if not payload: + return None, "" + return struct.unpack("!H", payload[:2])[0], payload[2:].decode() + + +async def serve(reader, writer, frames): + head = await reader.readuntil(b"\r\n\r\n") + key = next( + line.split(b":", 1)[1].strip() + for line in head.split(b"\r\n") + if line.lower().startswith(b"sec-websocket-key:") + ) + accept = base64.b64encode(hashlib.sha1(key + GUID).digest()) + writer.write( + b"HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\n" + b"Connection: Upgrade\r\nSec-WebSocket-Accept: " + accept + b"\r\n\r\n" + ) + frames.put_nowait(await read_close(reader)) + writer.write(b"\x88\x00") + await writer.drain() + writer.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "code, reason, expected", + [ + (None, None, (None, "")), + (4000, None, (4000, "")), + (4000, "bye", (4000, "bye")), + (None, "bye", (1000, "bye")), + ], +) +async def test_websocket_close_frame(code, reason, expected): + frames = asyncio.Queue() + server = await asyncio.start_server( + lambda r, w: serve(r, w, frames), "127.0.0.1", 0 + ) + url = f"ws://127.0.0.1:{server.sockets[0].getsockname()[1]}/" + try: + async with wreq.Client(proxies=[]) as client: + # Leaving the block after an explicit close must not fail. + async with await client.websocket(url) as ws: + await ws.close(code, reason) + assert await asyncio.wait_for(frames.get(), 5) == expected + finally: + server.close() + await server.wait_closed() + + +def test_blocking_websocket_exit_after_close(): + frames = asyncio.Queue() + loop = asyncio.new_event_loop() + server = loop.run_until_complete( + asyncio.start_server(lambda r, w: serve(r, w, frames), "127.0.0.1", 0) + ) + url = f"ws://127.0.0.1:{server.sockets[0].getsockname()[1]}/" + thread = threading.Thread(target=loop.run_forever, daemon=True) + thread.start() + try: + with wreq.blocking.Client(proxies=[]).websocket(url) as ws: + ws.close(4001) + frame = asyncio.run_coroutine_threadsafe(frames.get(), loop).result(5) + assert frame == (4001, "") + finally: + loop.call_soon_threadsafe(loop.stop) + thread.join(5) + server.close() + loop.run_until_complete(server.wait_closed()) + loop.close()