Cancel blocked server-stream handlers on client disconnect

src/connectrpc/_server_async.py

@@ -5,7 +5,7 @@
 import functools
 import inspect
 from abc import ABC, abstractmethod
-from asyncio import CancelledError, Event, create_task, sleep
+from asyncio import CancelledError, create_task, get_running_loop, sleep, timeout
 from dataclasses import replace
 from http import HTTPStatus
 from typing import TYPE_CHECKING, Generic, TypeVar, cast
@@ -423,8 +423,9 @@ async def _handle_stream(
                 self._read_max_bytes,
             )
 
-            disconnect_detected: Event | None = None
-            monitor_task = None
+            watch: contextlib.AbstractAsyncContextManager[None] = (
+                contextlib.nullcontext()
+            )
 
             match endpoint:
                 case EndpointUnary():
@@ -446,42 +447,32 @@ async def _handle_stream(
 
                     # The request has been fully consumed; monitor receive() for a
                     # client disconnect so we can stop streaming promptly.
-                    disconnect_detected = Event()
-
-                    async def _watch_for_disconnect() -> None:
-                        while True:
-                            msg = await receive()
-                            if msg["type"] == "http.disconnect":
-                                disconnect_detected.set()
-                                return
-
-                    monitor_task = create_task(_watch_for_disconnect())
+                    watch = _cancel_on_disconnect(receive)
                 case EndpointBidiStream():
                     response_stream = endpoint.function(request_stream, ctx)
 
             try:
-                async for message in response_stream:
-                    if disconnect_detected is not None and disconnect_detected.is_set():
-                        raise ConnectError(Code.CANCELED, "Client disconnected")
-                    # Don't send headers until the first message to allow logic a chance to add
-                    # response headers.
-                    if not sent_headers:
-                        await _send_stream_response_headers(
-                            send, protocol, codec, resp_compression.name(), ctx
+                # Exit the watch before aclose() below, so an error from the
+                # generator's cleanup can't leak the watch task.
+                async with watch:
+                    async for message in response_stream:
+                        # Don't send headers until the first message to allow logic a
+                        # chance to add response headers.
+                        if not sent_headers:
+                            await _send_stream_response_headers(
+                                send, protocol, codec, resp_compression.name(), ctx
+                            )
+                            sent_headers = True
+
+                        body = writer.write(message)
+                        await send(
+                            {
+                                "type": "http.response.body",
+                                "body": body,
+                                "more_body": True,
+                            }
                         )
-                        sent_headers = True
-
-                    body = writer.write(message)
-                    await send(
-                        {"type": "http.response.body", "body": body, "more_body": True}
-                    )
             finally:
-                # Cancel the monitor first so a throwing generator finally-block
-                # doesn't leak the task.
-                if monitor_task is not None:
-                    monitor_task.cancel()
-                    with contextlib.suppress(CancelledError):
-                        await monitor_task
                 # Explicitly close the stream so that any generator finally-blocks
                 # run promptly (Python defers async-generator cleanup to GC otherwise).
                 aclose = getattr(response_stream, "aclose", None)
@@ -590,6 +581,36 @@ def _resolve_endpoints(self, service: _SVC) -> Mapping[str, Endpoint]:
         return resolved_endpoints
 
 
+@contextlib.asynccontextmanager
+async def _cancel_on_disconnect(receive: ASGIReceiveCallable) -> AsyncGenerator[None]:
+    """Cancel the enclosed block with CANCELED when the client disconnects.
+
+    This reads receive(), so only use it once the request body has been read.
+    """
+    scope = timeout(None)
+
+    async def watch() -> None:
+        while True:
+            msg = await receive()
+            if msg["type"] == "http.disconnect":
+                scope.reschedule(get_running_loop().time())
+                return
+
+    try:
+        async with scope:
+            task = create_task(watch())
+            try:
+                yield
+            finally:
+                task.cancel()
+                with contextlib.suppress(CancelledError):
+                    await task
+    except TimeoutError as e:
+        if scope.expired():
+            raise ConnectError(Code.CANCELED, "Client disconnected") from e
+        raise
+
+
 async def _send_stream_response_headers(
     send: ASGISendCallable,
     protocol: ServerProtocol,

test/test_roundtrip.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import asyncio
+import json
 import random
 import struct
 from typing import TYPE_CHECKING
@@ -604,7 +605,83 @@ async def send(message):
             if response_count >= 3:
                 disconnect_trigger.set()
 
-    scope: HTTPScope = {
+    scope = _make_similar_hats_scope()
+
+    await asyncio.wait_for(app(scope, receive, send), timeout=5.0)
+
+    assert generator_closed.is_set(), (
+        "generator should be closed after client disconnect"
+    )
+
+
+@pytest.mark.asyncio
+async def test_server_stream_client_disconnect_while_blocked() -> None:
+    """A handler blocked between messages is cancelled when the client disconnects."""
+    events: list[str] = []
+
+    class BlockedHaberdasher(Haberdasher):
+        async def make_similar_hats(self, request, _ctx):
+            try:
+                yield Hat(size=request.inches)
+                await asyncio.Event().wait()
+            finally:
+                events.append("handler closed")
+
+    class WrappingInterceptor:
+        async def intercept_server_stream(self, call_next, request, ctx):
+            try:
+                async for response in call_next(request, ctx):
+                    yield response
+            finally:
+                events.append("interceptor closed")
+
+    class Recorder:
+        async def on_start(self, ctx):
+            pass
+
+        async def on_end(self, _token, _ctx, error):
+            code = error.code.value if isinstance(error, ConnectError) else error
+            events.append(f"on_end {code}")
+
+    app = HaberdasherASGIApplication(
+        BlockedHaberdasher(), interceptors=[Recorder(), WrappingInterceptor()]
+    )
+
+    request_bytes = Size(inches=10).to_binary()
+    request_body = struct.pack(">BI", 0, len(request_bytes)) + request_bytes
+
+    first_message_sent = asyncio.Event()
+    sent_request = False
+    sent: list[bytes] = []
+
+    async def receive() -> HTTPRequestEvent | HTTPDisconnectEvent:
+        nonlocal sent_request
+        if not sent_request:
+            sent_request = True
+            return {"type": "http.request", "body": request_body, "more_body": False}
+        await first_message_sent.wait()
+        return {"type": "http.disconnect"}
+
+    async def send(message):
+        if message["type"] == "http.response.body":
+            sent.append(message["body"])
+            first_message_sent.set()
+
+    scope = _make_similar_hats_scope()
+
+    await asyncio.wait_for(app(scope, receive, send), timeout=1.0)
+
+    assert events == ["handler closed", "interceptor closed", "on_end canceled"]
+    assert len(sent) == 2
+    end = sent[-1]
+    assert end[0] & 0b10, "last body is the end-of-stream message"
+    assert json.loads(end[5:]) == {
+        "error": {"code": "canceled", "message": "Client disconnected"}
+    }
+
+
+def _make_similar_hats_scope() -> HTTPScope:
+    return {
         "type": "http",
         "asgi": {"spec_version": "2.0", "version": "3.0"},
         "http_version": "1.1",
@@ -619,11 +696,3 @@ async def send(message):
         "server": None,
         "extensions": None,
     }
-
-    # Without the fix the app hangs forever (generator never stopped), causing a
-    # TimeoutError here.  With the fix it terminates promptly after the disconnect.
-    await asyncio.wait_for(app(scope, receive, send), timeout=5.0)
-
-    assert generator_closed.is_set(), (
-        "generator should be closed after client disconnect"
-    )