Merge pull request 'Keep pooled live snapshots cancellation-safe' (#138) from timmy/137-cancellation-safe-live-snapshot into main

Merge pull request 'Keep pooled live snapshots cancellation-safe' (#138) from timmy/137-cancellation-safe-live-snapshot into main

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(

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)"]

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