From e3d6ff6a3f2c3397051d11018abc09bee7f49c39 Mon Sep 17 00:00:00 2001 From: 0x676e67 Date: Sat, 3 Oct 2026 13:27:22 +0000 Subject: [PATCH 01/17] fix(deps): patch netty to close canceled HTTP/1 uploads --- Cargo.lock | 5 ++--- Cargo.toml | 2 ++ 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 1a219e01..03f62b16 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -111,7 +111,7 @@ dependencies = [ "regex", "rustc-hash", "shlex", - "syn 3.0.5", + "syn 2.0.119", ] [[package]] @@ -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=7b35fed8886df5618ecfa55456634e5627e9e442#7b35fed8886df5618ecfa55456634e5627e9e442" dependencies = [ "bytes", "futures-channel", diff --git a/Cargo.toml b/Cargo.toml index 0c9cf0d5..782201cf 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -84,6 +84,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" } +# demo/http1-cancel-shutdown +netty = { git = "https://github.com/0x676e67/netty", rev = "7b35fed8886df5618ecfa55456634e5627e9e442" } [profile.release] codegen-units = 1 From c104c14ee3c0146bfc1989efdfeed2039a441973 Mon Sep 17 00:00:00 2001 From: 0x676e67 Date: Sat, 3 Oct 2026 13:27:22 +0000 Subject: [PATCH 02/17] fix(aio): keep fallback wake ports until their loop closes --- src/aio/port.rs | 94 +++++++++++++++++++++++++++++++++++++++++------ tests/aio_test.py | 11 +++++- 2 files changed, 92 insertions(+), 13 deletions(-) diff --git a/src/aio/port.rs b/src/aio/port.rs index a95ce0b5..00933f46 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; @@ -23,7 +23,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 +35,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 +63,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 +99,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 +113,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 { @@ -197,12 +212,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 +245,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 +374,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/tests/aio_test.py b/tests/aio_test.py index d09ab0b6..984e4f50 100644 --- a/tests/aio_test.py +++ b/tests/aio_test.py @@ -11,7 +11,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 +67,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 would drop the 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 +123,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,7 +133,7 @@ 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"]) +@pytest.mark.parametrize("new_loop", LOOPS) def test_closed_loop_releases_pending_requests(new_loop): # The listener accepts connections into its backlog but never answers. with socket.create_server(("127.0.0.1", 0)) as server: From 9056093b815dcb41301cb55f9225c28dc65b7d49 Mon Sep 17 00:00:00 2001 From: 0x676e67 Date: Sat, 3 Oct 2026 13:27:22 +0000 Subject: [PATCH 03/17] refactor(stream): report read-ahead panics and tidy the reader --- Cargo.toml | 2 +- src/client/body/stream.rs | 41 ++++++++++++++++++++++++++------------- 2 files changed, 29 insertions(+), 14 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 782201cf..1bfeadfb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -38,7 +38,7 @@ pyo3 = { version = "0.29.0", features = [ "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" diff --git a/src/client/body/stream.rs b/src/client/body/stream.rs index 76e31c41..83010679 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,7 +12,12 @@ 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::{ sync::{ Notify, @@ -69,7 +75,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 +85,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 +103,8 @@ 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, so the buffer stays under about 3.3 MiB. const READ_AHEAD: usize = 8; /// Buffered frames returned before `__anext__` yields to the event loop once. @@ -110,7 +116,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 +126,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 +149,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()) } @@ -227,8 +242,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 +253,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; } }) From e8a19a0e6de4b4e26166cbae2338980a6e490c35 Mon Sep 17 00:00:00 2001 From: 0x676e67 Date: Sat, 3 Oct 2026 13:27:22 +0000 Subject: [PATCH 04/17] chore(deps): drop unused pyo3 experimental-async --- Cargo.toml | 1 - 1 file changed, 1 deletion(-) diff --git a/Cargo.toml b/Cargo.toml index 1bfeadfb..281e2bd6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -36,7 +36,6 @@ pyo3 = { version = "0.29.0", features = [ "multiple-pymethods", "extension-module", "generate-import-lib", - "experimental-async", ] } futures-util = { version = "0.3.33", default-features = false, features = ["std"] } serde = { version = "1.0.228", features = ["derive"] } From e728028a66b7ef5acf4c743e67c83089d7e7ebaa Mon Sep 17 00:00:00 2001 From: 0x676e67 Date: Sat, 3 Oct 2026 13:27:22 +0000 Subject: [PATCH 05/17] docs: correct aio, stream and response exit docs --- docs/source/getting-started/quickstart.md | 6 +++--- python/wreq/_compat.py | 5 +++-- python/wreq/blocking.py | 6 +++++- python/wreq/wreq.py | 7 ++++++- src/aio.rs | 12 +++++++----- src/client/resp/http.rs | 7 ++++--- 6 files changed, 28 insertions(+), 15 deletions(-) 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/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..7bcf9288 100644 --- a/python/wreq/blocking.py +++ b/python/wreq/blocking.py @@ -123,7 +123,11 @@ 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: ... diff --git a/python/wreq/wreq.py b/python/wreq/wreq.py index dc7ab0ce..7e0336d0 100644 --- a/python/wreq/wreq.py +++ b/python/wreq/wreq.py @@ -451,7 +451,12 @@ async def close(self) -> None: async def __aenter__(self) -> Any: ... async def __aexit__( self, _exc_type: Any, _exc_value: Any, _traceback: Any - ) -> 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: ... 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/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()) } } From 436f3248025146dc6504c0360c9220de78b5002a Mon Sep 17 00:00:00 2001 From: 0x676e67 Date: Sat, 3 Oct 2026 13:27:22 +0000 Subject: [PATCH 06/17] test(response): stop the read-ahead server on client failure --- tests/response_test.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) 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() From a3d02c6d50c82c359f32b4453fdbe196a48c1843 Mon Sep 17 00:00:00 2001 From: 0x676e67 Date: Sat, 3 Oct 2026 13:27:23 +0000 Subject: [PATCH 07/17] ci: install uvloop only from prebuilt wheels --- .github/workflows/ci.yml | 5 ++++- pyproject.toml | 1 - 2 files changed, 4 insertions(+), 2 deletions(-) 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/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", From 8f1e30d4170474dd135fc8b8b878ac67e7d9bb8f Mon Sep 17 00:00:00 2001 From: 0x676e67 Date: Sat, 3 Oct 2026 14:00:22 +0000 Subject: [PATCH 08/17] fix(deps): bump the netty patch for abandoned uploads --- Cargo.lock | 2 +- Cargo.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 03f62b16..2f185d3d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1178,7 +1178,7 @@ checksum = "27b02d87554356db9e9a873add8782d4ea6e3e58ea071a9adb9a2e8ddb884a8b" [[package]] name = "netty" version = "0.2.5" -source = "git+https://github.com/0x676e67/netty?rev=7b35fed8886df5618ecfa55456634e5627e9e442#7b35fed8886df5618ecfa55456634e5627e9e442" +source = "git+https://github.com/0x676e67/netty?rev=1b07384f8e517d835b0d336a1a3874aaeb08bc08#1b07384f8e517d835b0d336a1a3874aaeb08bc08" dependencies = [ "bytes", "futures-channel", diff --git a/Cargo.toml b/Cargo.toml index 281e2bd6..ad07ceb0 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -84,7 +84,7 @@ btls = { git = "https://github.com/0x676e67/btls", rev = "ffaf0e8c23d455af4e3069 btls-sys = { git = "https://github.com/0x676e67/btls", rev = "ffaf0e8c23d455af4e306986ce81bb57ebd03bc2" } tokio-btls = { git = "https://github.com/0x676e67/btls", rev = "ffaf0e8c23d455af4e306986ce81bb57ebd03bc2" } # demo/http1-cancel-shutdown -netty = { git = "https://github.com/0x676e67/netty", rev = "7b35fed8886df5618ecfa55456634e5627e9e442" } +netty = { git = "https://github.com/0x676e67/netty", rev = "1b07384f8e517d835b0d336a1a3874aaeb08bc08" } [profile.release] codegen-units = 1 From 272fbb887ea48e93ed2bf3f4d639b5c4913e7f72 Mon Sep 17 00:00:00 2001 From: 0x676e67 Date: Sat, 3 Oct 2026 14:00:22 +0000 Subject: [PATCH 09/17] fix(stream): cancel upload sources on the current runtime --- src/client/body/stream.rs | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/src/client/body/stream.rs b/src/client/body/stream.rs index 83010679..32093e4f 100644 --- a/src/client/body/stream.rs +++ b/src/client/body/stream.rs @@ -19,6 +19,7 @@ use pyo3::{ sync::PyOnceLock, }; use tokio::{ + runtime::Handle, sync::{ Notify, mpsc::{self, error::TryRecvError}, @@ -476,8 +477,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: acquire the interpreter on a blocking thread of the + // current runtime, so a client's own runtime never starts 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( From 0ea791b86957bb1d775627b9ce4bf6c9d1e95d16 Mon Sep 17 00:00:00 2001 From: 0x676e67 Date: Sat, 3 Oct 2026 14:09:31 +0000 Subject: [PATCH 10/17] docs: document context manager exits --- python/wreq/blocking.py | 11 +++++++++-- python/wreq/wreq.py | 25 ++++++++++++++++++++----- src/client.rs | 2 ++ src/client/body/stream.rs | 2 ++ src/client/resp/ws.rs | 2 ++ 5 files changed, 35 insertions(+), 7 deletions(-) diff --git a/python/wreq/blocking.py b/python/wreq/blocking.py index 7bcf9288..78de0088 100644 --- a/python/wreq/blocking.py +++ b/python/wreq/blocking.py @@ -209,7 +209,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. + """ + def __str__(self) -> str: ... @@ -493,7 +497,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 7e0336d0..b9306efe 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": ... @@ -328,7 +331,10 @@ async def __aenter__(self) -> Any: ... async def __aexit__( self, _exc_type: Any, _exc_value: Any, _traceback: Any - ) -> None: ... + ) -> None: + r""" + Release the body and end any pending read; returned views stay valid. + """ class Response: @@ -524,8 +530,14 @@ 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. + """ + def __str__(self) -> str: ... @@ -1440,7 +1452,10 @@ async def main(): async def __aenter__(self) -> Any: ... async def __aexit__( self, _exc_type: Any, _exc_value: Any, _traceback: Any - ) -> Any: ... + ) -> Any: + r""" + Close the client like `close()`: cancel pending requests and reject new ones. + """ async def delete( 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 32093e4f..12e3b6dc 100644 --- a/src/client/body/stream.rs +++ b/src/client/body/stream.rs @@ -218,6 +218,7 @@ impl Streamer { slf } + /// Release the body and end any pending read; returned views stay valid. fn __exit__<'py>( &self, py: Python, @@ -264,6 +265,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>, diff --git a/src/client/resp/ws.rs b/src/client/resp/ws.rs index bd81a459..99773f9e 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. fn __aexit__<'py>( &self, py: Python<'py>, @@ -279,6 +280,7 @@ impl BlockingWebSocket { slf } + /// Close the WebSocket connection without a close code or reason. fn __exit__<'py>( &self, py: Python<'py>, From cd37fedd44365919b4187816772261996906d22e Mon Sep 17 00:00:00 2001 From: 0x676e67 Date: Sat, 3 Oct 2026 14:09:31 +0000 Subject: [PATCH 11/17] chore(aio): give the wake drain a readable repr --- src/aio/port.rs | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/src/aio/port.rs b/src/aio/port.rs index 00933f46..8d6da909 100644 --- a/src/aio/port.rs +++ b/src/aio/port.rs @@ -197,6 +197,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 From 622927115c31fef65fc3f4a0fb52ba052e8f1064 Mon Sep 17 00:00:00 2001 From: 0x676e67 Date: Sat, 3 Oct 2026 14:23:20 +0000 Subject: [PATCH 12/17] test: cover abandoned uploads and wakes after the loop stops --- tests/aio_test.py | 15 ++++++++++++--- tests/upload_test.py | 32 ++++++++++++++++++++++++++++++++ 2 files changed, 44 insertions(+), 3 deletions(-) diff --git a/tests/aio_test.py b/tests/aio_test.py index 984e4f50..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 @@ -67,7 +68,7 @@ async def serve(reader, writer): finally: writer.close() - # The default backlog of 100 would drop the 300 concurrent stream connections. + # 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: @@ -134,14 +135,22 @@ def test_wakes_resume_tasks_on_event_loops(new_loop): reason="PyPy's cpyext does not collect cycles through extension objects", ) @pytest.mark.parametrize("new_loop", LOOPS) -def test_closed_loop_releases_pending_requests(new_loop): - # The listener accepts connections into its backlog but never answers. +@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/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): From 927f1e56e2f2ecf143e4a8fe6929db7c66e0744b Mon Sep 17 00:00:00 2001 From: 0x676e67 Date: Sat, 3 Oct 2026 14:23:20 +0000 Subject: [PATCH 13/17] docs: correct follow-up comments and note connection reuse --- docs/source/guide/basic.md | 2 +- docs/source/guide/blocking.md | 2 +- python/wreq/blocking.py | 1 + python/wreq/wreq.py | 16 ++++------------ src/aio/port.rs | 3 ++- src/client/body/stream.rs | 9 +++++---- 6 files changed, 14 insertions(+), 19 deletions(-) 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/python/wreq/blocking.py b/python/wreq/blocking.py index 78de0088..b9b804c9 100644 --- a/python/wreq/blocking.py +++ b/python/wreq/blocking.py @@ -128,6 +128,7 @@ def __exit__(self, _exc_type: Any, _exc_value: Any, _traceback: Any) -> None: 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: ... diff --git a/python/wreq/wreq.py b/python/wreq/wreq.py index b9306efe..848e5836 100644 --- a/python/wreq/wreq.py +++ b/python/wreq/wreq.py @@ -329,9 +329,7 @@ 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. """ @@ -455,9 +453,7 @@ 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. @@ -531,9 +527,7 @@ async def close( """ 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 WebSocket connection without a close code or reason. """ @@ -1450,9 +1444,7 @@ 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. """ diff --git a/src/aio/port.rs b/src/aio/port.rs index 8d6da909..dba9f0a3 100644 --- a/src/aio/port.rs +++ b/src/aio/port.rs @@ -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, } diff --git a/src/client/body/stream.rs b/src/client/body/stream.rs index 12e3b6dc..cd59d04d 100644 --- a/src/client/body/stream.rs +++ b/src/client/body/stream.rs @@ -104,8 +104,9 @@ struct NotifyOnDrop(Arc); // ===== impl Streamer ===== impl Streamer { - /// Frames buffered ahead of Python. A frame holds at most one transport read, up - /// to 408 KiB on HTTP/1, so the buffer stays under about 3.3 MiB. + /// 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. @@ -479,8 +480,8 @@ 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 of the - // current runtime, so a client's own runtime never starts the shared one. + // 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 || { From 29dd30e054bdcc3fbbccda5fddc638323dfdfb94 Mon Sep 17 00:00:00 2001 From: 0x676e67 Date: Sat, 3 Oct 2026 14:23:20 +0000 Subject: [PATCH 14/17] chore(deps): limit the lockfile change to netty --- Cargo.lock | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Cargo.lock b/Cargo.lock index 2f185d3d..a394fe43 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -111,7 +111,7 @@ dependencies = [ "regex", "rustc-hash", "shlex", - "syn 2.0.119", + "syn 3.0.5", ] [[package]] From fdc76e134b88e03c1d6f051c0cc48fc25b80d662 Mon Sep 17 00:00:00 2001 From: 0x676e67 Date: Sat, 3 Oct 2026 14:35:24 +0000 Subject: [PATCH 15/17] docs: drop a nonexistent attribute from the response example --- python/wreq/wreq.py | 1 - 1 file changed, 1 deletion(-) diff --git a/python/wreq/wreq.py b/python/wreq/wreq.py index 848e5836..6dd45e30 100644 --- a/python/wreq/wreq.py +++ b/python/wreq/wreq.py @@ -352,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() From d2badcdc52066e0505020926ad843e4a2c09a884 Mon Sep 17 00:00:00 2001 From: 0x676e67 Date: Sat, 3 Oct 2026 14:35:24 +0000 Subject: [PATCH 16/17] fix(ws): send close codes without a reason and allow exit after close --- python/wreq/blocking.py | 2 +- python/wreq/wreq.py | 2 +- src/client/resp/ws.rs | 10 +++-- src/client/resp/ws/cmd.rs | 48 +++++++++++++++------ tests/websocket_test.py | 91 +++++++++++++++++++++++++++++++++++++++ 5 files changed, 134 insertions(+), 19 deletions(-) create mode 100644 tests/websocket_test.py diff --git a/python/wreq/blocking.py b/python/wreq/blocking.py index b9b804c9..b8e21314 100644 --- a/python/wreq/blocking.py +++ b/python/wreq/blocking.py @@ -212,7 +212,7 @@ def close( def __enter__(self) -> Any: ... def __exit__(self, _exc_type: Any, _exc_value: Any, _traceback: Any) -> None: r""" - Close the WebSocket connection without a close code or reason. + Close the WebSocket connection without a close code or reason, unless already closed. """ def __str__(self) -> str: ... diff --git a/python/wreq/wreq.py b/python/wreq/wreq.py index 6dd45e30..44e49543 100644 --- a/python/wreq/wreq.py +++ b/python/wreq/wreq.py @@ -528,7 +528,7 @@ async def close( 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. + Close the WebSocket connection without a close code or reason, unless already closed. """ def __str__(self) -> str: ... diff --git a/src/client/resp/ws.rs b/src/client/resp/ws.rs index 99773f9e..4c0cd3bc 100644 --- a/src/client/resp/ws.rs +++ b/src/client/resp/ws.rs @@ -166,7 +166,7 @@ impl WebSocket { aio::ready("WebSocket.__aenter__", slf) } - /// Close the WebSocket connection without a close code or reason. + /// Close the WebSocket connection without a close code or reason, unless already closed. fn __aexit__<'py>( &self, py: Python<'py>, @@ -178,7 +178,7 @@ impl WebSocket { py, "WebSocket.__aexit__", &self.runtime, - cmd::close(self.cmd.clone(), None, None), + cmd::close_on_exit(self.cmd.clone()), ) } } @@ -280,7 +280,7 @@ impl BlockingWebSocket { slf } - /// Close the WebSocket connection without a close code or reason. + /// Close the WebSocket connection without a close code or reason, unless already closed. fn __exit__<'py>( &self, py: Python<'py>, @@ -288,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/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() From 9b22133194f122629a073e24d690dd850db07bf0 Mon Sep 17 00:00:00 2001 From: 0x676e67 Date: Sat, 3 Oct 2026 15:16:37 +0000 Subject: [PATCH 17/17] chore(deps): patch netty to the merged cancel fix --- Cargo.lock | 2 +- Cargo.toml | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index a394fe43..45288969 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1178,7 +1178,7 @@ checksum = "27b02d87554356db9e9a873add8782d4ea6e3e58ea071a9adb9a2e8ddb884a8b" [[package]] name = "netty" version = "0.2.5" -source = "git+https://github.com/0x676e67/netty?rev=1b07384f8e517d835b0d336a1a3874aaeb08bc08#1b07384f8e517d835b0d336a1a3874aaeb08bc08" +source = "git+https://github.com/0x676e67/netty?rev=f8652ddf5e386d66b7f03cceedbc6b319d849762#f8652ddf5e386d66b7f03cceedbc6b319d849762" dependencies = [ "bytes", "futures-channel", diff --git a/Cargo.toml b/Cargo.toml index ad07ceb0..b96ba3f1 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -83,8 +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" } -# demo/http1-cancel-shutdown -netty = { git = "https://github.com/0x676e67/netty", rev = "1b07384f8e517d835b0d336a1a3874aaeb08bc08" } +# main (0x676e67/netty#145) +netty = { git = "https://github.com/0x676e67/netty", rev = "f8652ddf5e386d66b7f03cceedbc6b319d849762" } [profile.release] codegen-units = 1