Test writer backpressure and fix drain cancellation and shutdown

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`.

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()

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<PyWeakrefReference>,
+}
+
+#[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::<PyFastStreamProtocol>()?
+                .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<PyAny>,
+    ) -> PyResult<Py<PyAny>> {
+        // 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<Py<PyAny>> {
-        self.build_drain_future(py, None)
+    fn _drain_helper(mut slf: PyRefMut<'_, Self>, py: Python<'_>) -> PyResult<Py<PyAny>> {
+        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::<PyStreamTransport>() {
+                    transport.borrow().core.is_closing()
+                } else {
+                    self.transport
+                        .call_method0(py, "is_closing")?
+                        .extract::<bool>(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)
     }
 }
 

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())