From 0dc5c2c8314c1fe4855d552dbb6bfc0da932f498 Mon Sep 17 00:00:00 2001 From: timmy Date: Thu, 6 Aug 2026 20:51:53 +0000 Subject: [PATCH] fix: keep pooled live snapshots cancellation-safe (#137) --- src/main.py | 25 ++++++++++-- tests/test_gitea_transport.py | 38 ++++++++++++++++++ tests/test_live_snapshot.py | 75 +++++++++++++++++++++++++++++++++++ 3 files changed, 135 insertions(+), 3 deletions(-) diff --git a/src/main.py b/src/main.py index d540f5d..4804676 100644 --- a/src/main.py +++ b/src/main.py @@ -26,11 +26,23 @@ from src.views import router as frontend_router @asynccontextmanager async def lifespan(_app: FastAPI): + global _live_snapshot_task gitea_proxy.start_client() try: yield finally: - await gitea_proxy.stop_client() + task = _live_snapshot_task + if task is not None and not task.done(): + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + try: + await gitea_proxy.stop_client() + finally: + if _live_snapshot_task is task: + _live_snapshot_task = None app = FastAPI(title="Stackchain Dashboard", lifespan=lifespan) @@ -251,15 +263,22 @@ async def _build_live_snapshot() -> dict: } +async def _build_live_snapshot_before_deadline() -> dict: + async with asyncio.timeout(CONTEXT_TIMEOUT_SECONDS): + return await _build_live_snapshot() + + @app.get("/api/v1/live") async def live_snapshot() -> JSONResponse: """Return a fresh, section-aware snapshot; join only an active identical load.""" global _live_snapshot_task if _live_snapshot_task is None or _live_snapshot_task.done(): - _live_snapshot_task = asyncio.create_task(_build_live_snapshot()) + _live_snapshot_task = asyncio.create_task( + _build_live_snapshot_before_deadline() + ) task = _live_snapshot_task try: - result = await asyncio.wait_for(task, timeout=CONTEXT_TIMEOUT_SECONDS) + result = await asyncio.shield(task) return JSONResponse(result) except TimeoutError: return JSONResponse( diff --git a/tests/test_gitea_transport.py b/tests/test_gitea_transport.py index 4985cab..f7fc3cb 100644 --- a/tests/test_gitea_transport.py +++ b/tests/test_gitea_transport.py @@ -1,3 +1,5 @@ +import asyncio + import httpx import pytest @@ -41,3 +43,39 @@ async def test_application_lifespan_opens_and_closes_gitea_transport(monkeypatch assert calls == ["start"] assert calls == ["start", "stop"] + + +@pytest.mark.anyio +async def test_application_shutdown_finishes_snapshot_before_closing_transport(monkeypatch): + calls = [] + started = asyncio.Event() + + async def active_snapshot(): + started.set() + try: + await asyncio.Event().wait() + finally: + calls.append("snapshot cancelled") + + monkeypatch.setattr(gitea_proxy, "start_client", lambda: calls.append("start")) + + async def stop_client(): + assert main._live_snapshot_task is not None + state = "done" if main._live_snapshot_task.done() else "active" + calls.append(f"stop ({state})") + + monkeypatch.setattr(gitea_proxy, "stop_client", stop_client) + + try: + async with main.app.router.lifespan_context(main.app): + main._live_snapshot_task = asyncio.create_task(active_snapshot()) + await started.wait() + finally: + task = main._live_snapshot_task + if task is not None and not task.done(): + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + main._live_snapshot_task = None + + assert calls == ["start", "snapshot cancelled", "stop (done)"] diff --git a/tests/test_live_snapshot.py b/tests/test_live_snapshot.py index ae6874a..da6c76a 100644 --- a/tests/test_live_snapshot.py +++ b/tests/test_live_snapshot.py @@ -184,3 +184,78 @@ async def test_live_snapshot_coalesces_only_simultaneous_requests(monkeypatch): await main.live_snapshot() assert user_calls == 2 + + +@pytest.mark.anyio +async def test_cancelling_one_waiter_does_not_cancel_the_shared_snapshot(monkeypatch): + started = asyncio.Event() + release = asyncio.Event() + snapshot_calls = 0 + + async def blocked_snapshot(): + nonlocal snapshot_calls + snapshot_calls += 1 + started.set() + await release.wait() + return { + "context": {}, + "events": [], + "notifications": [], + "sections": {}, + } + + monkeypatch.setattr(main, "_build_live_snapshot", blocked_snapshot) + + disconnected = asyncio.create_task(main.live_snapshot()) + await started.wait() + survivor = asyncio.create_task(main.live_snapshot()) + await asyncio.sleep(0) + + disconnected.cancel() + with pytest.raises(asyncio.CancelledError): + await disconnected + + assert main._live_snapshot_task is not None + assert not main._live_snapshot_task.done() + + release.set() + response = await survivor + + assert response.status_code == 200 + assert snapshot_calls == 1 + + +@pytest.mark.anyio +async def test_shared_snapshot_deadline_cancels_upstream_work_for_all_waiters(monkeypatch): + started = asyncio.Event() + cancelled = asyncio.Event() + + async def snapshot_that_exceeds_deadline(): + started.set() + try: + await asyncio.Event().wait() + finally: + cancelled.set() + + monkeypatch.setattr(main, "_build_live_snapshot", snapshot_that_exceeds_deadline) + monkeypatch.setattr(main, "CONTEXT_TIMEOUT_SECONDS", 0.01) + + first = asyncio.create_task(main.live_snapshot()) + await started.wait() + second = asyncio.create_task(main.live_snapshot()) + + try: + responses = await asyncio.gather(first, second) + await asyncio.wait_for(cancelled.wait(), timeout=0.1) + + assert [response.status_code for response in responses] == [503, 503] + assert [payload(response)["error"] for response in responses] == [ + "Gitea live snapshot timed out after 0.01s", + "Gitea live snapshot timed out after 0.01s", + ] + finally: + task = main._live_snapshot_task + if task is not None and not task.done(): + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task -- 2.43.0