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" },