diff --git a/src/main.py b/src/main.py index 6620edf..1d6f861 100644 --- a/src/main.py +++ b/src/main.py @@ -53,6 +53,8 @@ EVENT_STREAM_TIMEOUT_SECONDS = 5.0 READINESS_TIMEOUT_SECONDS = 5.0 REVIEW_DETAIL_TIMEOUT_SECONDS = 5.0 NOTIFICATION_MUTATION_TIMEOUT_SECONDS = 5.0 +BULK_NOTIFICATION_CONCURRENCY = 5 +BULK_NOTIFICATION_DEADLINE_SECONDS = 6.0 LIVE_SNAPSHOT_FRESHNESS_SECONDS = 8.0 FRONTEND_DIR = Path(__file__).resolve().parent.parent / "frontend" _live_snapshot_task: asyncio.Task | None = None @@ -444,13 +446,37 @@ async def _mark_notification_read_result(thread_id: int) -> tuple[int, bool]: @app.patch("/api/v1/notifications/read") async def read_notifications(batch: NotificationReadBatch) -> JSONResponse: thread_ids = list(dict.fromkeys(batch.ids)) - results = await asyncio.gather( - *(_mark_notification_read_result(thread_id) for thread_id in thread_ids) + + semaphore = asyncio.Semaphore(BULK_NOTIFICATION_CONCURRENCY) + + async def mark_within_limit(thread_id: int) -> tuple[int, bool]: + async with semaphore: + return await _mark_notification_read_result(thread_id) + + tasks = [asyncio.create_task(mark_within_limit(thread_id)) for thread_id in thread_ids] + done, pending = await asyncio.wait( + tasks, timeout=BULK_NOTIFICATION_DEADLINE_SECONDS + ) + for task in pending: + task.cancel() + if pending: + await asyncio.gather(*pending, return_exceptions=True) + + succeeded_ids = { + thread_id + for task in done + if not task.cancelled() and task.exception() is None + for thread_id, succeeded in [task.result()] + if succeeded + } + failed = [thread_id for thread_id in thread_ids if thread_id not in succeeded_ids] + return JSONResponse( + { + "marked": [thread_id for thread_id in thread_ids if thread_id in succeeded_ids], + "failed": failed, + }, + headers={"Retry-After": "1"} if failed else None, ) - return JSONResponse({ - "marked": [thread_id for thread_id, succeeded in results if succeeded], - "failed": [thread_id for thread_id, succeeded in results if not succeeded], - }) @app.patch("/api/v1/notifications/{thread_id}/read") diff --git a/tests/test_notification_read.py b/tests/test_notification_read.py index 7ca397b..78532b6 100644 --- a/tests/test_notification_read.py +++ b/tests/test_notification_read.py @@ -1,4 +1,5 @@ import asyncio +import time import httpx import pytest @@ -77,6 +78,83 @@ async def test_bulk_mark_read_reports_partial_progress_and_retains_only_failures assert main._live_snapshot_value == {"notifications": [{"id": 43}]} +@pytest.mark.anyio +async def test_bulk_mark_read_limits_upstream_concurrency(monkeypatch): + active = 0 + peak_active = 0 + started = asyncio.Event() + release = asyncio.Event() + calls = [] + + monkeypatch.setattr(main, "_read_notification_ids", set()) + monkeypatch.setattr(main, "_live_snapshot_value", None) + + async def mark(thread_id): + nonlocal active, peak_active + calls.append(thread_id) + active += 1 + peak_active = max(peak_active, active) + if active == 5: + started.set() + try: + await release.wait() + finally: + active -= 1 + + monkeypatch.setattr(main, "mark_notification_read", mark) + transport = httpx.ASGITransport(app=main.app) + async with httpx.AsyncClient(transport=transport, base_url="http://test") as client: + request = asyncio.create_task( + client.patch("/api/v1/notifications/read", json={"ids": list(range(1, 51))}) + ) + await asyncio.wait_for(started.wait(), timeout=1) + await asyncio.sleep(0) + assert peak_active == 5 + assert len(calls) == 5 + release.set() + response = await request + + assert response.status_code == 200 + assert response.json() == {"marked": list(range(1, 51)), "failed": []} + assert sorted(calls) == list(range(1, 51)) + + +@pytest.mark.anyio +async def test_bulk_mark_read_cancels_unfinished_work_at_batch_deadline(monkeypatch): + active = 0 + cancelled = [] + + async def mark(thread_id): + nonlocal active + active += 1 + try: + await asyncio.sleep(0.08) + except asyncio.CancelledError: + cancelled.append(thread_id) + raise + finally: + active -= 1 + + monkeypatch.setattr(main, "mark_notification_read", mark) + monkeypatch.setattr(main, "BULK_NOTIFICATION_DEADLINE_SECONDS", 0.01, raising=False) + monkeypatch.setattr(main, "_read_notification_ids", set()) + monkeypatch.setattr(main, "_live_snapshot_value", None) + transport = httpx.ASGITransport(app=main.app) + started_at = time.monotonic() + async with httpx.AsyncClient(transport=transport, base_url="http://test") as client: + response = await client.patch( + "/api/v1/notifications/read", json={"ids": list(range(1, 11))} + ) + elapsed = time.monotonic() - started_at + + assert elapsed < 0.06 + assert response.status_code == 200 + assert response.json() == {"marked": [], "failed": list(range(1, 11))} + assert response.headers["retry-after"] == "1" + assert active == 0 + assert sorted(cancelled) == list(range(1, 6)) + + @pytest.mark.anyio async def test_bulk_mark_read_rejects_empty_invalid_and_oversized_batches(monkeypatch): marked = []