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 new file mode 100644 index 0000000..57561b4 --- /dev/null +++ b/tests/test_stream_writer_lifecycle.py @@ -0,0 +1,263 @@ +"""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("drain_api", ["writer", "protocol"]) +@pytest.mark.parametrize("cancel_mode", ["cancel", "timeout"]) +def test_cancelled_drains_release_waiters_while_still_paused( + loop, cancel_mode, drain_api +): + async def main(): + _, writer, peer = await connect(loop) + 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() + paused = True + + async def cancelled_drain(): + entered = asyncio.Event() + + async def drain(): + entered.set() + await drain_call() + + 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(asyncio.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 drain_call() + + live = asyncio.create_task(live_drain()) + await entered.wait() + assert not live.done() + protocol.resume_writing() + paused = False + await live + await drain_call() + + finally: + if paused: + protocol.resume_writing() + peer.close() + writer.close() + await asyncio.wait_for(writer.wait_closed(), 5) + + 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()) + if isinstance(loop, rsloop.Loop): + assert type(writer).__name__ == "PyFastStreamWriter" + 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 asyncio.wait_for(writer.wait_closed(), 5) + + 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] + # 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() + 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 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] + # 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 + # 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())