From 791f777b91255dadd2bdada0d1b942c84c368242 Mon Sep 17 00:00:00 2001 From: Adrian Ehrsam Date: Sat, 10 Oct 2026 12:58:17 +0000 Subject: [PATCH] Fix disconnect watcher rescheduling a finished asyncio.Timeout (#59) With BaseHTTPMiddleware in front (as in our apps) and a server whose send() takes a moment (Granian), PostgresJsonResponse's disconnect watcher woke up after almost every response and failed with "Cannot change state of finished Timeout" ("Task exception was never retrieved"): the middleware's receive() waits in an anyio task group that it cancels once the response has been sent, the watcher's own cancellation arrives with it and is swallowed there, and receive() then returns a synthetic http.disconnect after the stream's cancel scope has been exited. The watcher now only reschedules (and only keeps receiving) while the stream is still running. CancelOnDisconnectRoute's watcher waits on our own Event, which cannot swallow the cancellation, but gets the same guard defensively. Bump to 0.16.1. Co-Authored-By: Claude Opus 5.5 --- pgdevkit/fastapi/cancel_on_disconnect.py | 6 ++++- pgdevkit/fastapi/json_response.py | 10 +++++-- pyproject.toml | 2 +- tests/_asgi.py | 23 ++++++++++++++++- tests/_granian_app.py | 5 ++++ tests/test_cancel_on_disconnect.py | 33 ++++++++++++++++++++++-- tests/test_pg_json_cancel.py | 26 +++++++++++++++++-- tests/test_pg_json_granian.py | 12 +++++++++ uv.lock | 2 +- 9 files changed, 109 insertions(+), 10 deletions(-) diff --git a/pgdevkit/fastapi/cancel_on_disconnect.py b/pgdevkit/fastapi/cancel_on_disconnect.py index 0b32815..113887e 100644 --- a/pgdevkit/fastapi/cancel_on_disconnect.py +++ b/pgdevkit/fastapi/cancel_on_disconnect.py @@ -74,9 +74,12 @@ 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: @@ -84,6 +87,7 @@ async def cancel_when_gone(cancel_scope: asyncio.Timeout) -> None: try: return await handler(request) finally: + handled = True watcher.cancel() except TimeoutError: if not cancel_scope.expired(): # the handler's own timeout diff --git a/pgdevkit/fastapi/json_response.py b/pgdevkit/fastapi/json_response.py index 318491f..d78a43b 100644 --- a/pgdevkit/fastapi/json_response.py +++ b/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(): diff --git a/pyproject.toml b/pyproject.toml index 01369e1..44e127c 100644 --- a/pyproject.toml +++ b/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" diff --git a/tests/_asgi.py b/tests/_asgi.py index f117c29..619c4a4 100644 --- a/tests/_asgi.py +++ b/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) diff --git a/tests/_granian_app.py b/tests/_granian_app.py index 3dc0b9b..10dd9c8 100644 --- a/tests/_granian_app.py +++ b/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( diff --git a/tests/test_cancel_on_disconnect.py b/tests/test_cancel_on_disconnect.py index 94afeaa..88528c6 100644 --- a/tests/test_cancel_on_disconnect.py +++ b/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 diff --git a/tests/test_pg_json_cancel.py b/tests/test_pg_json_cancel.py index 130bf17..60783fd 100644 --- a/tests/test_pg_json_cancel.py +++ b/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: diff --git a/tests/test_pg_json_granian.py b/tests/test_pg_json_granian.py index def2e99..35fd625 100644 --- a/tests/test_pg_json_granian.py +++ b/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: diff --git a/uv.lock b/uv.lock index 0c8e19d..efe15c5 100644 --- a/uv.lock +++ b/uv.lock @@ -451,7 +451,7 @@ wheels = [ [[package]] name = "pgdevkit" -version = "0.16.0" +version = "0.16.1" source = { editable = "." } dependencies = [ { name = "docker" },