From 9db27f4a2ac772881aa503b71158500f3bac899f Mon Sep 17 00:00:00 2001 From: Yehor Smoliakov Date: Sat, 10 Oct 2026 03:48:04 -0700 Subject: [PATCH 1/4] test: reproduce writer drain cancellation and shutdown gaps --- tests/test_stream_writer_lifecycle.py | 193 ++++++++++++++++++++++++++ 1 file changed, 193 insertions(+) create mode 100644 tests/test_stream_writer_lifecycle.py diff --git a/tests/test_stream_writer_lifecycle.py b/tests/test_stream_writer_lifecycle.py new file mode 100644 index 0000000..a748cf3 --- /dev/null +++ b/tests/test_stream_writer_lifecycle.py @@ -0,0 +1,193 @@ +"""Stream writer contracts under cancellation, backpressure and shutdown.""" + +import asyncio +import gc +import socket +import weakref +from typing import Any, cast + +import pytest +import rsloop + + +@pytest.fixture(params=["asyncio", "rsloop", "uvloop"]) +def loop(request): + if request.param == "uvloop": + factory = pytest.importorskip("uvloop").new_event_loop + elif request.param == "rsloop": + factory = rsloop.new_event_loop + else: + factory = asyncio.new_event_loop + event_loop = factory() + errors = [] + event_loop.set_exception_handler(lambda _, context: errors.append(context)) + yield event_loop + pending = asyncio.all_tasks(event_loop) + for task in pending: + task.cancel() + if pending: + event_loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True)) + event_loop.run_until_complete(event_loop.shutdown_asyncgens()) + event_loop.close() + assert not pending, "test leaked tasks" + assert not errors, errors + + +def run(loop, coro): + return loop.run_until_complete(asyncio.wait_for(coro, 20)) + + +@pytest.mark.parametrize("cancel_mode", ["cancel", "timeout"]) +def test_cancelled_drains_release_waiters_while_still_paused(loop, cancel_mode): + async def main(): + _, writer, peer = await connect(loop) + protocol = writer.transport.get_protocol() + paused = False + try: + protocol.pause_writing() + paused = True + + async def cancelled_drain(): + entered = asyncio.Event() + + async def drain(): + entered.set() + await writer.drain() + + task = asyncio.create_task(drain()) + await entered.wait() + assert not task.done() + # Inspect only ownership, without relying on either protocol's container. + reference = weakref.ref(cast(Any, task)._fut_waiter) + if cancel_mode == "cancel": + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + else: + with pytest.raises(TimeoutError): + await asyncio.wait_for(task, 0) + return reference + + references = [await cancelled_drain() for _ in range(32)] + checkpoint = asyncio.Event() + loop.call_soon(checkpoint.set) + await checkpoint.wait() + gc.collect() + assert all(reference() is None for reference in references), ( + "cancelled drains retained until the stalled peer resumes" + ) + # Cancellation must not unpause the protocol or poison the next drain. + entered = asyncio.Event() + + async def live_drain(): + entered.set() + await writer.drain() + + live = asyncio.create_task(live_drain()) + await entered.wait() + assert not live.done() + protocol.resume_writing() + paused = False + await live + await writer.drain() + + finally: + if paused: + protocol.resume_writing() + peer.close() + writer.close() + await writer.wait_closed() + + run(loop, main()) + + +async def connect(loop): + listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + listener.bind(("127.0.0.1", 0)) + listener.listen() + listener.setblocking(False) + try: + reader, writer = await asyncio.open_connection(*listener.getsockname()) + peer, _ = await loop.sock_accept(listener) + peer.setblocking(False) + return reader, writer, peer + finally: + listener.close() + + +def test_drain_yields_when_transport_is_closing(loop): + async def main(): + _, writer, peer = await connect(loop) + try: + writer.close() + assert writer.is_closing() + # A closing drain must let connection_lost run, even when unpaused. + with pytest.raises(ConnectionResetError): + await writer.drain() + finally: + peer.close() + writer.close() + await writer.wait_closed() + + run(loop, main()) + + +@pytest.mark.parametrize("finish", ["close", "write_eof"]) +def test_slow_peer_concurrent_drains_cancellation_and_ordered_shutdown(loop, finish): + async def main(): + reader, writer, peer = await connect(loop) + tasks: list[asyncio.Task[Any]] = [] + try: + writer.get_extra_info("socket").setsockopt( + socket.SOL_SOCKET, socket.SO_SNDBUF, 4096 + ) + writer.transport.set_write_buffer_limits(high=1024, low=0) + payload = bytes(range(256)) * 4096 + copies = 0 + while writer.transport.get_write_buffer_size() <= 1024: + writer.write(payload) + copies += 1 + assert copies * len(payload) <= 32 * 1024 * 1024 + started = [asyncio.Event() for _ in range(3)] + + async def drain(index): + started[index].set() + await writer.drain() + + tasks = [asyncio.create_task(drain(i)) for i in range(3)] + for event in started: + await event.wait() + assert all(not task.done() for task in tasks) + tasks[1].cancel() + with pytest.raises(asyncio.CancelledError): + await tasks[1] + assert not tasks[0].done() and not tasks[2].done() + + async def receive(): + received = bytearray() + while chunk := await loop.sock_recv(peer, 65536): + received.extend(chunk) + return received + + receiver = asyncio.create_task(receive()) + tasks.append(receiver) + await asyncio.gather(tasks[0], tasks[2]) + await writer.drain() + writer.write(b"tail") + getattr(writer, finish)() + assert await receiver == payload * copies + b"tail" + if finish == "write_eof": + # Sending EOF must preserve the other direction of the stream. + await loop.sock_sendall(peer, b"ack") + peer.shutdown(socket.SHUT_WR) + assert await reader.read() == b"ack" + finally: + for task in tasks: + if not task.done(): + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + peer.close() + writer.close() + await writer.wait_closed() + + run(loop, main()) From e01af16623a3a0a7dd6cd824fe7e95fa9e3cc5c5 Mon Sep 17 00:00:00 2001 From: Yehor Smoliakov Date: Sat, 10 Oct 2026 03:48:09 -0700 Subject: [PATCH 2/4] fix: release cancelled drains and yield during writer shutdown --- docs/development.md | 15 +++++ python/rsloop/_stream_writer.py | 10 ++++ src/transport/stream/fast.rs | 80 +++++++++++++++++++++++++-- tests/test_stream_writer_lifecycle.py | 80 ++++++++++++++++++++++++--- 4 files changed, 172 insertions(+), 13 deletions(-) create mode 100644 python/rsloop/_stream_writer.py diff --git a/docs/development.md b/docs/development.md index ea48280..8c82283 100644 --- a/docs/development.md +++ b/docs/development.md @@ -238,6 +238,21 @@ uv run python -m pytest tests/test_compat.py uv run python -m pytest tests/test_tls.py -k start_tls ``` +Writer lifecycle contracts are exercised with: + +```bash +uv run python -m pytest tests/test_stream_writer_lifecycle.py +``` + +These tests compare asyncio, rsloop and optional uvloop over loopback TCP. +They cover concurrent drains against a slow peer, cancellation and timeout +cleanup while writing remains paused, ordered delivery through close/half-close, +shutdown of pending drains, and error precedence. Events establish when drains +are pending; a raw peer controls when bytes can be consumed. Protocol callbacks +exercise prolonged pause deterministically without depending on kernel buffer +sizes. Weak references check that cancelled waits are released before resume. +Tests also check for leaked tasks and unhandled loop exceptions. + CI and `just test` use `scripts/run_python_tests.py`, which forwards pytest arguments and prints recurring stack dumps when tests stop making progress. For the same diagnostics locally, run `uv run python scripts/run_python_tests.py`. diff --git a/python/rsloop/_stream_writer.py b/python/rsloop/_stream_writer.py new file mode 100644 index 0000000..ce100b6 --- /dev/null +++ b/python/rsloop/_stream_writer.py @@ -0,0 +1,10 @@ +"""Cold-path orchestration for native stream writer shutdown.""" + +import asyncio + + +async def drain_closing(protocol): + # StreamWriter.drain() yields on a closing transport so its queued + # connection_lost callback can make the drain fail instead of busy-looping. + await asyncio.sleep(0) + await protocol._drain_helper() diff --git a/src/transport/stream/fast.rs b/src/transport/stream/fast.rs index 767ac88..1bb084b 100644 --- a/src/transport/stream/fast.rs +++ b/src/transport/stream/fast.rs @@ -2164,7 +2164,57 @@ pub struct PyFastStreamProtocol { over_ssl: bool, } +/// Remove completed drains even if a stalled peer never resumes writing. +/// The weak reference avoids a protocol -> Future -> callback -> protocol +/// cycle. +#[pyclass(module = "rsloop._loop")] +struct PyDrainWaiterDone { + protocol: Py, +} + +#[pymethods] +impl PyDrainWaiterDone { + #[cfg_attr(feature = "profile", hotpath::measure(impl_type = "PyDrainWaiterDone"))] + fn __call__(&self, py: Python<'_>, future: &Bound<'_, PyAny>) -> PyResult<()> { + if let Some(protocol) = self.protocol.bind(py).upgrade() { + protocol + .cast::()? + .try_borrow_mut()? + .drain_waiters + .retain(|waiter| !waiter.bind(py).is(future)); + } + Ok(()) + } +} + impl PyFastStreamProtocol { + #[cfg_attr( + feature = "profile", + hotpath::measure(impl_type = "PyFastStreamProtocol") + )] + fn watch_drain_waiter( + slf: PyRefMut<'_, Self>, + py: Python<'_>, + future: Py, + ) -> PyResult> { + // Ready drains keep their allocation-free fast path. + if slf + .drain_waiters + .last() + .is_some_and(|waiter| waiter.bind(py).is(future.bind(py))) + { + let protocol = slf.into_pyobject(py)?; + let callback = Py::new( + py, + PyDrainWaiterDone { + protocol: PyWeakrefReference::new(&protocol)?.unbind(), + }, + )?; + future.call_method1(py, "add_done_callback", (callback,))?; + } + Ok(future) + } + #[cfg_attr( feature = "profile", hotpath::measure(impl_type = "PyFastStreamProtocol") @@ -2506,8 +2556,9 @@ impl PyFastStreamProtocol { feature = "profile", hotpath::measure(impl_type = "PyFastStreamProtocol") )] - fn _drain_helper(&mut self, py: Python<'_>) -> PyResult> { - self.build_drain_future(py, None) + fn _drain_helper(mut slf: PyRefMut<'_, Self>, py: Python<'_>) -> PyResult> { + let future = slf.build_drain_future(py, None)?; + Self::watch_drain_waiter(slf, py, future) } #[cfg_attr( @@ -2715,9 +2766,28 @@ impl PyFastStreamWriter { .exception .as_ref() .map(|exc| exc.clone_ref(py)); - self.protocol - .borrow_mut(py) - .build_drain_future(py, reader_exception) + // Match StreamWriter.drain(): reader errors take precedence, and a + // closing transport must yield so connection_lost can be delivered. + if reader_exception.is_none() { + let closing = + if let Ok(transport) = self.transport.bind(py).cast_exact::() { + transport.borrow().core.is_closing() + } else { + self.transport + .call_method0(py, "is_closing")? + .extract::(py)? + }; + if closing { + return Ok(py + .import("rsloop._stream_writer")? + .getattr("drain_closing")? + .call1((self.protocol.clone_ref(py),))? + .unbind()); + } + } + let mut protocol = self.protocol.borrow_mut(py); + let future = protocol.build_drain_future(py, reader_exception)?; + PyFastStreamProtocol::watch_drain_waiter(protocol, py, future) } } diff --git a/tests/test_stream_writer_lifecycle.py b/tests/test_stream_writer_lifecycle.py index a748cf3..f49d9ac 100644 --- a/tests/test_stream_writer_lifecycle.py +++ b/tests/test_stream_writer_lifecycle.py @@ -37,11 +37,15 @@ def run(loop, coro): return loop.run_until_complete(asyncio.wait_for(coro, 20)) +@pytest.mark.parametrize("drain_api", ["writer", "protocol"]) @pytest.mark.parametrize("cancel_mode", ["cancel", "timeout"]) -def test_cancelled_drains_release_waiters_while_still_paused(loop, cancel_mode): +def test_cancelled_drains_release_waiters_while_still_paused( + loop, cancel_mode, drain_api +): async def main(): _, writer, peer = await connect(loop) - protocol = writer.transport.get_protocol() + protocol = cast(Any, writer.transport.get_protocol()) + drain_call = writer.drain if drain_api == "writer" else protocol._drain_helper paused = False try: protocol.pause_writing() @@ -52,7 +56,7 @@ async def cancelled_drain(): async def drain(): entered.set() - await writer.drain() + await drain_call() task = asyncio.create_task(drain()) await entered.wait() @@ -81,7 +85,7 @@ async def drain(): async def live_drain(): entered.set() - await writer.drain() + await drain_call() live = asyncio.create_task(live_drain()) await entered.wait() @@ -89,14 +93,14 @@ async def live_drain(): protocol.resume_writing() paused = False await live - await writer.drain() + await drain_call() finally: if paused: protocol.resume_writing() peer.close() writer.close() - await writer.wait_closed() + await asyncio.wait_for(writer.wait_closed(), 5) run(loop, main()) @@ -108,6 +112,8 @@ async def connect(loop): listener.setblocking(False) try: reader, writer = await asyncio.open_connection(*listener.getsockname()) + if isinstance(loop, rsloop.Loop): + assert type(writer).__name__ == "PyFastStreamWriter" peer, _ = await loop.sock_accept(listener) peer.setblocking(False) return reader, writer, peer @@ -127,7 +133,7 @@ async def main(): finally: peer.close() writer.close() - await writer.wait_closed() + await asyncio.wait_for(writer.wait_closed(), 5) run(loop, main()) @@ -188,6 +194,64 @@ async def receive(): await asyncio.gather(*tasks, return_exceptions=True) peer.close() writer.close() - await writer.wait_closed() + await asyncio.wait_for(writer.wait_closed(), 5) + + run(loop, main()) + + +def test_reader_error_takes_precedence_over_closing_drain(loop): + async def main(): + reader, writer, peer = await connect(loop) + error = OSError("peer failed while closing") + try: + reader.set_exception(error) + writer.close() + with pytest.raises(OSError) as raised: + await writer.drain() + assert raised.value is error + finally: + peer.close() + writer.close() + await asyncio.wait_for(writer.wait_closed(), 5) + + run(loop, main()) + + +@pytest.mark.parametrize("finish", ["close", "abort"]) +def test_shutdown_wakes_all_paused_drains(loop, finish): + async def main(): + _, writer, peer = await connect(loop) + tasks = [] + try: + writer.transport.get_protocol().pause_writing() + started = [asyncio.Event() for _ in range(3)] + + async def drain(index): + started[index].set() + await writer.drain() + + tasks = [asyncio.create_task(drain(i)) for i in range(3)] + for event in started: + await event.wait() + assert all(not task.done() for task in tasks) + tasks[1].cancel() + with pytest.raises(asyncio.CancelledError): + await tasks[1] + getattr(writer.transport, finish)() + results = await asyncio.gather(tasks[0], tasks[2], return_exceptions=True) + # Graceful connection_lost(None) releases existing drains; a new + # drain observes the closed connection and must fail. + assert results == [None, None] + await asyncio.wait_for(writer.wait_closed(), 5) + with pytest.raises(ConnectionResetError): + await writer.drain() + finally: + for task in tasks: + if not task.done(): + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + peer.close() + writer.close() + await asyncio.wait_for(writer.wait_closed(), 5) run(loop, main()) From abc25636ac4adf48d0204f4f4da2e4163f04e3c2 Mon Sep 17 00:00:00 2001 From: Yehor Smoliakov Date: Sat, 10 Oct 2026 03:50:04 -0700 Subject: [PATCH 3/4] test: use the portable asyncio timeout exception --- tests/test_stream_writer_lifecycle.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_stream_writer_lifecycle.py b/tests/test_stream_writer_lifecycle.py index f49d9ac..041f057 100644 --- a/tests/test_stream_writer_lifecycle.py +++ b/tests/test_stream_writer_lifecycle.py @@ -68,7 +68,7 @@ async def drain(): with pytest.raises(asyncio.CancelledError): await task else: - with pytest.raises(TimeoutError): + with pytest.raises(asyncio.TimeoutError): await asyncio.wait_for(task, 0) return reference From 25f3d3e61b4330e9fb96509c3659832834e5a6c7 Mon Sep 17 00:00:00 2001 From: Yehor Smoliakov Date: Sat, 10 Oct 2026 03:55:09 -0700 Subject: [PATCH 4/4] test: allow legitimate IOCP drain progress during cancellation --- tests/test_stream_writer_lifecycle.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/test_stream_writer_lifecycle.py b/tests/test_stream_writer_lifecycle.py index 041f057..57561b4 100644 --- a/tests/test_stream_writer_lifecycle.py +++ b/tests/test_stream_writer_lifecycle.py @@ -167,7 +167,10 @@ async def drain(index): tasks[1].cancel() with pytest.raises(asyncio.CancelledError): await tasks[1] - assert not tasks[0].done() and not tasks[2].done() + # The kernel may have accepted queued bytes during cancellation, + # even though the peer application has not read yet (notably IOCP). + # Other drains may already succeed, but must never be cancelled. + assert not tasks[0].cancelled() and not tasks[2].cancelled() async def receive(): received = bytearray() @@ -237,6 +240,9 @@ async def drain(index): tasks[1].cancel() with pytest.raises(asyncio.CancelledError): await tasks[1] + # Here pause is explicit: no kernel write completion can resume + # the protocol. Cancelling one drain must leave the others blocked. + assert not tasks[0].done() and not tasks[2].done() getattr(writer.transport, finish)() results = await asyncio.gather(tasks[0], tasks[2], return_exceptions=True) # Graceful connection_lost(None) releases existing drains; a new