Fix disconnect watcher rescheduling a finished asyncio.Timeout (#59)

pgdevkit/fastapi/cancel_on_disconnect.py

@@ -74,16 +74,20 @@ async def cancel_on_disconnect(request: Request) -> Response:
             request._receive = receiver.receive  # the same Request object, so a cached body is seen by everyone
             reader = asyncio.create_task(receiver.read())
 
+            handled = False  # as in PostgresJsonResponse: once set, the cancel scope may be gone
+
             async def cancel_when_gone(cancel_scope: asyncio.Timeout) -> None:
                 await receiver.disconnected.wait()
-                cancel_scope.reschedule(loop.time())
+                if not handled:  # (the Event can't swallow the cancellation; this keeps reschedule() safe anyway)
+                    cancel_scope.reschedule(loop.time())
 
             try:
                 async with asyncio.timeout(None) as cancel_scope:
                     watcher = asyncio.create_task(cancel_when_gone(cancel_scope))
                     try:
                         return await handler(request)
                     finally:
+                        handled = True
                         watcher.cancel()
             except TimeoutError:
                 if not cancel_scope.expired():  # the handler's own timeout

pgdevkit/fastapi/json_response.py

@@ -112,11 +112,16 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
         # Starlette only watches for http.disconnect on ASGI spec < 2.4, so do it ourselves. An asyncio timeout is
         # the cancel scope: unlike an anyio task group, it doesn't leave Granian hanging when the body fails.
         loop = asyncio.get_running_loop()
+        streamed = False  # set before the watcher is cancelled: from then on the cancel scope may be gone
 
         async def watch_for_disconnect(cancel_scope: asyncio.Timeout) -> None:
-            while (await receive())["type"] != "http.disconnect":
+            # Checked again after every receive(): a wrapped receive() can swallow the watcher's cancellation and
+            # return instead (BaseHTTPMiddleware's turns it into an http.disconnect once the response has been sent),
+            # and reschedule() raises once the scope has been exited.
+            while not streamed and (await receive())["type"] != "http.disconnect":
                 pass
-            cancel_scope.reschedule(loop.time())  # i.e. cancel the stream now
+            if not streamed:
+                cancel_scope.reschedule(loop.time())  # i.e. cancel the stream now
 
         started = False
 
@@ -131,6 +136,7 @@ async def tracking_send(message: Message) -> None:
                 try:
                     await self.stream_response(tracking_send)
                 finally:
+                    streamed = True
                     watcher.cancel()
         except TimeoutError:
             if not cancel_scope.expired():

pyproject.toml

@@ -11,7 +11,7 @@ packages = ["pgdevkit"]
 
 [project]
 name = "pgdevkit"
-version = "0.16.0"
+version = "0.16.1"
 description = "A helper for developing with Postgres"
 readme = "README.md"
 requires-python = ">=3.14"

tests/_asgi.py

@@ -1,8 +1,10 @@
 """Drive an ASGI response by hand, with a client that can disconnect at will (httpx's ASGITransport can't)."""
 
 import asyncio
+import contextlib
+import gc
 import time
-from collections.abc import Awaitable, Callable
+from collections.abc import AsyncIterator, Awaitable, Callable
 
 import psycopg
 
@@ -13,6 +15,7 @@ def __init__(self) -> None:
         self.sent: list[dict] = []
         self.fail_sends_after: int | None = None  # simulate ASGI spec >= 2.4: send() raises OSError once gone
         self.disconnect_when_send_fails = False  # ... and http.disconnect arrives at that very moment
+        self.send_iterations = 0  # event loop iterations each send() takes, as a write to a real socket does
 
     async def receive(self) -> dict:
         await self.disconnect.wait()
@@ -24,6 +27,8 @@ async def send(self, message: dict) -> None:
                 self.disconnect.set()
             raise OSError("client went away")
         self.sent.append(message)
+        for _ in range(self.send_iterations):
+            await asyncio.sleep(0)
 
     async def call(self, app: Callable[..., Awaitable[None]]) -> None:
         await app({"type": "http", "asgi": {"spec_version": "2.3"}}, self.receive, self.send)
@@ -51,3 +56,19 @@ async def wait_until(condition: Callable[[], Awaitable[bool] | bool], *, timeout
             return True
         await asyncio.sleep(0.05)
     return False
+
+
+@contextlib.asynccontextmanager
+async def unhandled_loop_errors() -> AsyncIterator[list[dict]]:
+    """Collect what the event loop would log, e.g. "Task exception was never retrieved" for a failed watcher task."""
+    loop = asyncio.get_running_loop()
+    errors: list[dict] = []
+    previous_handler = loop.get_exception_handler()
+    loop.set_exception_handler(lambda _, context: errors.append(context))
+    try:
+        yield errors
+        for _ in range(5):  # let leftover tasks finish, then have the failed ones reported
+            await asyncio.sleep(0)
+        gc.collect()
+    finally:
+        loop.set_exception_handler(previous_handler)

tests/_granian_app.py

@@ -58,6 +58,11 @@ async def sleep():
     return PostgresJsonResponse("select pg_sleep(30) as slept", pool=pool)
 
 
+@app.get("/fails-at-once")  # before the first byte: a 500 JSON error
+async def fails_at_once():
+    return PostgresJsonResponse("select 1 / 0 as boom", pool=pool)
+
+
 @app.get("/fails-midway")
 async def fails_midway():
     return PostgresJsonResponse(

tests/test_cancel_on_disconnect.py

@@ -5,6 +5,7 @@
 import json
 import uuid
 from collections.abc import AsyncIterator
+from typing import Literal
 
 import pytest
 from fastapi import APIRouter, Depends, FastAPI, HTTPException, Request
@@ -14,8 +15,8 @@
 from pydantic import BaseModel
 
 from pgdevkit.db import fetch_all
-from pgdevkit.fastapi import CancelOnDisconnectRoute
-from tests._asgi import active_queries, wait_until
+from pgdevkit.fastapi import CancelOnDisconnectRoute, cancel_on_disconnect
+from tests._asgi import active_queries, unhandled_loop_errors, wait_until
 
 
 class Filter(BaseModel):
@@ -319,6 +320,34 @@ async def test_a_handler_that_finished_wins_over_a_disconnect_that_arrives_with_
     assert (client.status, client.body) == (200, b'{"done":true}')
 
 
+class _SwallowingEvent(asyncio.Event):
+    async def wait(self) -> Literal[True]:
+        try:
+            return await super().wait()
+        except asyncio.CancelledError:  # what BaseHTTPMiddleware's receive() does to PostgresJsonResponse's watcher
+            return True
+
+
+async def test_a_watcher_whose_cancellation_is_swallowed_leaves_the_finished_scope_alone(
+    app: FastAPI, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """Issue #59 hit PostgresJsonResponse's watcher; this one waits on our own Event, which can't swallow its
+    cancellation. If it ever does, the watcher wakes up after the handler's cancel scope has been exited."""
+
+    class Receiver(cancel_on_disconnect._Receiver):
+        def __init__(self, *args: object) -> None:
+            super().__init__(*args)
+            self.disconnected = _SwallowingEvent()
+
+    monkeypatch.setattr(cancel_on_disconnect, "_Receiver", Receiver)
+    async with unhandled_loop_errors() as errors:  # "Task exception was never retrieved: RuntimeError(...)"
+        client = HttpClient(path="/ok")  # (not /instant: the watcher must be waiting by the time the handler is done)
+        await client.call(app)
+
+    assert (client.status, client.body) == (200, b'[{"one":1}]')
+    assert [(e["message"], e.get("exception")) for e in errors] == []
+
+
 async def test_context_variables_and_yield_dependencies_work_as_without_the_route(app: FastAPI) -> None:
     client = HttpClient(path="/context")
     await client.call(app)  # reset() in another context than set() would raise ValueError: 500

tests/test_pg_json_cancel.py

@@ -13,11 +13,13 @@
 from psycopg import AsyncConnection
 from psycopg.pq import TransactionStatus
 from psycopg_pool import AsyncConnectionPool
-from starlette.requests import ClientDisconnect
+from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
+from starlette.requests import ClientDisconnect, Request
+from starlette.responses import Response
 
 from pgdevkit.fastapi import PostgresJsonResponse, json_response
 from pgdevkit.fastapi.json_response import _cancel_if_running
-from tests._asgi import FakeClient, active_queries, wait_until
+from tests._asgi import FakeClient, active_queries, unhandled_loop_errors, wait_until
 
 SLEEP_30S = "select pg_sleep(30) as slept"
 # padded rows: Postgres only flushes its 8 kB output buffer, so tiny rows would all arrive at the very end
@@ -156,6 +158,26 @@ async def test_a_second_cancellation_during_cleanup_does_not_skip_it(postgres_ds
         assert await wait_until(lambda: query_stopped(postgres_dsn, app_name), timeout=3)
 
 
+@pytest.mark.parametrize("query", ["select 1 / 0 as boom", "select 1 as one"])
+async def test_a_receive_that_swallows_the_watchers_cancellation_does_not_break_it(
+    pool: AsyncConnectionPool, query: str
+) -> None:
+    """Issue #59 (Granian, BaseHTTPMiddleware in front): that middleware's receive() waits in an anyio task group,
+    which it cancels once the response has been sent. The watcher's own cancellation arrives with it and is swallowed
+    there, and receive() returns an http.disconnect after the stream has finished, its cancel scope exited already."""
+
+    async def passthrough(request: Request, call_next: RequestResponseEndpoint) -> Response:
+        return await call_next(request)
+
+    async with unhandled_loop_errors() as errors:  # "Task exception was never retrieved: RuntimeError(...)"
+        client = FakeClient()
+        client.send_iterations = 2  # the race needs a server whose send() takes a moment, as Granian's does
+        await client.call(BaseHTTPMiddleware(PostgresJsonResponse(query, pool=pool), dispatch=passthrough))
+
+    assert client.sent[0]["status"] == (500 if "boom" in query else 200)
+    assert [(e["message"], e.get("exception")) for e in errors] == []
+
+
 async def test_a_finished_query_is_left_alone(pool: AsyncConnectionPool) -> None:
     """A cancel sent after the query finished could hit the next statement of the pooled connection."""
     async with pool.connection() as conn:

tests/test_pg_json_granian.py

@@ -70,6 +70,18 @@ async def test_normal_response(server: str) -> None:
         assert (await client.get(f"{server}/ok")).json() == [{"i": 1}, {"i": 2}, {"i": 3}]
 
 
+@pytest.mark.parametrize("path", ["/ok", "/fails-at-once"])
+async def test_finished_responses_leave_no_failed_watcher_behind(server: str, server_log: Path, path: str) -> None:
+    """Issue #59: with BaseHTTPMiddleware in front, the disconnect watcher used to wake up after almost every response
+    and fail with "Cannot change state of finished Timeout" (logged as "Task exception was never retrieved")."""
+    log_start = len(server_log.read_text())
+    async with httpx.AsyncClient() as client:
+        for _ in range(20):
+            assert (await client.get(f"{server}{path}")).status_code == (200 if path == "/ok" else 500)
+    await asyncio.sleep(0.2)  # the server logs it when the failed task is collected
+    assert "never retrieved" not in server_log.read_text()[log_start:]
+
+
 async def test_aborted_client_cancels_the_query(
     server: str, server_log: Path, postgres_dsn: str, app_name: str
 ) -> None: