perf: share live refreshes across workers (#433)

tests/test_live_snapshot.py

@@ -1,14 +1,21 @@
 import asyncio
 import json
+import sqlite3
 
 import httpx
 import pytest
 
 from src import main
+from src.live_snapshot_store import LiveSnapshotStore
 
 
 @pytest.fixture(autouse=True)
-def reset_live_snapshot_task():
+def reset_live_snapshot_task(tmp_path, monkeypatch):
+    monkeypatch.setattr(
+        main,
+        "_live_snapshot_store",
+        LiveSnapshotStore(tmp_path / "live.sqlite3", clock=lambda: main._live_snapshot_clock()),
+    )
     main._live_snapshot_task = None
     main._live_snapshot_value = None
     main._live_snapshot_created_at = None
@@ -131,7 +138,7 @@ async def test_live_snapshot_omits_section_bodies_at_known_revisions(monkeypatch
 
 
 @pytest.mark.anyio
-async def test_live_snapshot_token_from_another_worker_never_suppresses_content(monkeypatch):
+async def test_live_snapshot_token_from_another_worker_never_suppresses_content(monkeypatch, tmp_path):
     async def user():
         return {"id": 1, "login": "timmy"}
 
@@ -144,8 +151,6 @@ async def test_live_snapshot_token_from_another_worker_never_suppresses_content(
     monkeypatch.setattr(main, "pull_requests", empty)
     monkeypatch.setattr(main, "activity_events", lambda _user: empty())
     monkeypatch.setattr(main, "notifications", empty)
-    monkeypatch.setattr(main, "_live_revision_generation", "worker-a")
-
     worker_a = payload(await main.live_snapshot())
     worker_a_token = worker_a["revisions"]["context"]
 
@@ -157,12 +162,15 @@ async def test_live_snapshot_token_from_another_worker_never_suppresses_content(
     main._live_section_revisions = {
         section: 0 for section in main.LIVE_SNAPSHOT_SECTIONS
     }
-    monkeypatch.setattr(main, "_live_revision_generation", "worker-b")
+    main._live_snapshot_store = LiveSnapshotStore(
+        tmp_path / "other-worker.sqlite3", clock=lambda: main._live_snapshot_clock()
+    )
 
     worker_b = payload(await main.live_snapshot(context_revision=worker_a_token))
 
-    assert worker_a_token == "worker-a.1"
-    assert worker_b["revisions"]["context"] == "worker-b.1"
+    assert worker_a_token.endswith(".1")
+    assert worker_b["revisions"]["context"].endswith(".1")
+    assert worker_b["revisions"]["context"] != worker_a_token
     assert worker_b["context"]["user"]["login"] == "timmy"
 
 
@@ -302,7 +310,10 @@ async def test_live_snapshot_reuses_completed_snapshot_inside_freshness_window(m
     first = asyncio.create_task(main.live_snapshot())
     await asyncio.sleep(0)
     second = asyncio.create_task(main.live_snapshot())
-    await asyncio.sleep(0)
+    for _ in range(100):
+        if user_calls:
+            break
+        await asyncio.to_thread(lambda: None)
     assert user_calls == 1
 
     release.set()
@@ -332,7 +343,7 @@ async def test_stale_snapshot_returns_immediately_while_one_refresh_revalidates(
             "sections": {},
         }
 
-    monkeypatch.setattr(main.time, "monotonic", lambda: now)
+    monkeypatch.setattr(main, "_live_snapshot_clock", lambda: now)
     monkeypatch.setattr(main, "_build_live_snapshot", snapshot)
 
     first = payload(await main.live_snapshot())
@@ -373,13 +384,14 @@ async def test_failed_revalidation_enters_cooldown_and_keeps_last_snapshot(monke
             },
         }
 
-    monkeypatch.setattr(main.time, "monotonic", lambda: now)
+    monkeypatch.setattr(main, "_live_snapshot_clock", lambda: now)
     monkeypatch.setattr(main, "_build_live_snapshot", snapshot)
 
     await main.live_snapshot()
     now += main.LIVE_SNAPSHOT_FRESHNESS_SECONDS + 1
     stale = payload(await main.live_snapshot())
-    await asyncio.sleep(0)
+    assert main._live_snapshot_task is not None
+    await asyncio.gather(main._live_snapshot_task, return_exceptions=True)
     degraded = payload(await main.live_snapshot())
 
     assert stale["freshness"]["revalidating"] is True
@@ -421,13 +433,14 @@ async def test_partial_refresh_updates_fresh_sections_and_retains_failed_section
             },
         }
 
-    monkeypatch.setattr(main.time, "monotonic", lambda: now)
+    monkeypatch.setattr(main, "_live_snapshot_clock", lambda: now)
     monkeypatch.setattr(main, "_build_live_snapshot", snapshot)
 
     await main.live_snapshot()
     now += main.LIVE_SNAPSHOT_FRESHNESS_SECONDS + 1
     await main.live_snapshot()
-    await asyncio.sleep(0)
+    assert main._live_snapshot_task is not None
+    await main._live_snapshot_task
     result = payload(await main.live_snapshot())
 
     assert result["context"] == {"generation": 2}
@@ -470,7 +483,7 @@ async def test_notification_cooldown_does_not_stop_due_work_and_activity_refresh
             raise ConnectionError("notifications unavailable")
         return []
 
-    monkeypatch.setattr(main.time, "monotonic", lambda: now)
+    monkeypatch.setattr(main, "_live_snapshot_clock", lambda: now)
     monkeypatch.setattr(main, "current_user", user)
     monkeypatch.setattr(main, "repos", work)
     monkeypatch.setattr(main, "issues", empty_work)
@@ -537,20 +550,19 @@ async def test_retry_window_starts_one_refresh_and_success_clears_degraded_state
             "sections": {"context": "fresh", "events": "fresh", "notifications": "fresh"},
         }
 
-    monkeypatch.setattr(main.time, "monotonic", lambda: now)
+    monkeypatch.setattr(main, "_live_snapshot_clock", lambda: now)
     monkeypatch.setattr(main, "_build_live_snapshot", snapshot)
 
     await main.live_snapshot()
     now += main.LIVE_SNAPSHOT_FRESHNESS_SECONDS + 1
     await main.live_snapshot()
-    await asyncio.sleep(0)
+    assert main._live_snapshot_task is not None
+    await main._live_snapshot_task
     now += main.LIVE_SNAPSHOT_RETRY_BASE_SECONDS
 
     first = asyncio.create_task(main.live_snapshot())
     second = asyncio.create_task(main.live_snapshot())
-    await asyncio.sleep(0)
-    await asyncio.sleep(0)
-    assert retry_started.is_set()
+    await asyncio.wait_for(retry_started.wait(), timeout=1.0)
     assert builds == 3
     assert payload(await first)["freshness"]["revalidating"] is True
     assert payload(await second)["freshness"]["revalidating"] is True
@@ -640,3 +652,123 @@ async def test_shared_snapshot_deadline_cancels_upstream_work_for_all_waiters(mo
             task.cancel()
             with pytest.raises(asyncio.CancelledError):
                 await task
+
+
+@pytest.mark.anyio
+async def test_cold_refresh_cooldown_remains_a_retryable_503(monkeypatch):
+    now = 100.0
+
+    async def unavailable(_sections=None):
+        raise ConnectionError("private upstream failure")
+
+    monkeypatch.setattr(main, "_live_snapshot_clock", lambda: now)
+    monkeypatch.setattr(main, "_build_live_snapshot", unavailable)
+
+    first = await main.live_snapshot()
+    second = await main.live_snapshot()
+
+    assert first.status_code == 503
+    assert second.status_code == 503
+    assert payload(second) == {
+        "error": "Gitea live snapshot is temporarily unavailable"
+    }
+
+
+@pytest.mark.anyio
+async def test_cancelled_refresh_releases_shared_lease(monkeypatch):
+    started = asyncio.Event()
+
+    async def blocked(_sections=None):
+        started.set()
+        await asyncio.Event().wait()
+
+    monkeypatch.setattr(main, "_build_live_snapshot", blocked)
+    request = asyncio.create_task(main.live_snapshot())
+    await started.wait()
+    assert main._live_snapshot_task is not None
+
+    main._live_snapshot_task.cancel()
+    await asyncio.gather(request, main._live_snapshot_task, return_exceptions=True)
+
+    replacement = main._live_snapshot_store.try_acquire_refresh(
+        {"context"}, lease_seconds=1
+    )
+    assert replacement is not None
+
+
+@pytest.mark.anyio
+async def test_shared_store_outage_returns_retryable_503(monkeypatch):
+    def unavailable():
+        raise sqlite3.OperationalError("private database detail")
+
+    monkeypatch.setattr(main._live_snapshot_store, "load", unavailable)
+
+    response = await main.live_snapshot()
+
+    assert response.status_code == 503
+    assert response.headers["retry-after"] == "1"
+    assert payload(response) == {
+        "error": "Gitea live snapshot state is temporarily unavailable"
+    }
+
+
+@pytest.mark.anyio
+async def test_independent_workers_reuse_shared_live_snapshot(monkeypatch, tmp_path):
+    calls = 0
+
+    async def user():
+        nonlocal calls
+        calls += 1
+        return {"id": 1, "login": "timmy"}
+
+    async def empty():
+        return []
+
+    monkeypatch.setattr(main, "current_user", user)
+    monkeypatch.setattr(main, "repos", empty)
+    monkeypatch.setattr(main, "issues", empty)
+    monkeypatch.setattr(main, "pull_requests", empty)
+    monkeypatch.setattr(main, "activity_events", lambda _user: empty())
+    monkeypatch.setattr(main, "notifications", empty)
+    path = tmp_path / "worker-shared.sqlite3"
+    monkeypatch.setattr(main, "_live_snapshot_store", LiveSnapshotStore(path, clock=lambda: 100.0))
+
+    first = payload(await main.live_snapshot())
+    main._live_snapshot_value = None
+    main._live_snapshot_created_at = None
+    main._live_section_created_at = {section: None for section in main.LIVE_SNAPSHOT_SECTIONS}
+    main._live_snapshot_store = LiveSnapshotStore(path, clock=lambda: 100.0)
+    second = payload(await main.live_snapshot())
+
+    assert calls == 1
+    assert second["context"] == first["context"]
+    assert second["revisions"] == first["revisions"]
+
+
+@pytest.mark.anyio
+async def test_live_snapshot_persists_reboot_stable_wall_timestamps(monkeypatch, tmp_path):
+    epoch = 1_700_000_000.0
+
+    async def user():
+        return {"id": 1, "login": "timmy"}
+
+    async def empty():
+        return []
+
+    monkeypatch.setattr(main, "_live_snapshot_clock", lambda: epoch, raising=False)
+    monkeypatch.setattr(
+        main,
+        "_live_snapshot_store",
+        LiveSnapshotStore(tmp_path / "wall.sqlite3", clock=lambda: main._live_snapshot_clock()),
+    )
+    monkeypatch.setattr(main, "current_user", user)
+    monkeypatch.setattr(main, "repos", empty)
+    monkeypatch.setattr(main, "issues", empty)
+    monkeypatch.setattr(main, "pull_requests", empty)
+    monkeypatch.setattr(main, "activity_events", lambda _user: empty())
+    monkeypatch.setattr(main, "notifications", empty)
+
+    response = await main.live_snapshot()
+
+    assert response.status_code == 200
+    assert set(main._live_snapshot_store.load().created_at.values()) == {epoch}

tests/test_live_snapshot_store.py

@@ -0,0 +1,134 @@
+import threading
+import os
+
+import pytest
+
+from src import live_snapshot_store
+from src.live_snapshot_store import LiveSnapshotStore, RefreshLeaseLost
+
+
+def test_independent_stores_racing_for_refresh_have_one_lease_winner(tmp_path):
+    path = tmp_path / "live.sqlite3"
+    first = LiveSnapshotStore(path, clock=lambda: 100.0)
+    second = LiveSnapshotStore(path, clock=lambda: 100.0)
+    barrier = threading.Barrier(2)
+    results = []
+
+    def acquire(store):
+        barrier.wait()
+        results.append(store.try_acquire_refresh({"context", "events"}, lease_seconds=5))
+
+    threads = [threading.Thread(target=acquire, args=(store,)) for store in (first, second)]
+    for thread in threads:
+        thread.start()
+    for thread in threads:
+        thread.join()
+
+    assert sum(result is not None for result in results) == 1
+
+
+def test_snapshot_and_shared_revisions_publish_atomically(tmp_path):
+    path = tmp_path / "live.sqlite3"
+    writer = LiveSnapshotStore(path, clock=lambda: 100.0)
+    reader = LiveSnapshotStore(path, clock=lambda: 100.0)
+    lease = writer.try_acquire_refresh(set(("context", "events", "notifications")), lease_seconds=5)
+    assert lease is not None
+
+    state = writer.publish_refresh(
+        lease,
+        value={
+            "context": {"generation": 1},
+            "events": [],
+            "notifications": [],
+            "sections": {section: "fresh" for section in ("context", "events", "notifications")},
+        },
+        created_at={section: 100.0 for section in ("context", "events", "notifications")},
+        failure_count={section: 0 for section in ("context", "events", "notifications")},
+        retry_at={section: None for section in ("context", "events", "notifications")},
+        changed_sections=set(("context", "events", "notifications")),
+    )
+
+    observed = reader.load()
+    assert observed.value == state.value
+    assert observed.revisions == {"context": 1, "events": 1, "notifications": 1}
+    assert observed.refreshing_sections == set()
+    with pytest.raises(RefreshLeaseLost):
+        writer.publish_refresh(
+            lease,
+            value={"context": {"generation": 2}},
+            created_at=observed.created_at,
+            failure_count=observed.failure_count,
+            retry_at=observed.retry_at,
+            changed_sections={"context"},
+        )
+
+
+def test_expired_refresh_lease_can_be_recovered(tmp_path):
+    now = 100.0
+    path = tmp_path / "live.sqlite3"
+    abandoned = LiveSnapshotStore(path, clock=lambda: now)
+    recovery = LiveSnapshotStore(path, clock=lambda: now)
+
+    first = abandoned.try_acquire_refresh({"context"}, lease_seconds=5)
+    assert first is not None
+    assert recovery.try_acquire_refresh({"context"}, lease_seconds=5) is None
+    now = 106.0
+    second = recovery.try_acquire_refresh({"context"}, lease_seconds=5)
+
+    assert second is not None
+    assert second != first
+
+
+def test_default_store_clock_uses_reboot_stable_wall_time(tmp_path, monkeypatch):
+    epoch = 1_700_000_000.0
+    monkeypatch.setattr(live_snapshot_store.time, "time", lambda: epoch)
+    store = LiveSnapshotStore(tmp_path / "live.sqlite3")
+
+    lease = store.try_acquire_refresh({"context"}, lease_seconds=5)
+
+    assert lease is not None
+    assert store.load().lease_expires_at == epoch + 5
+
+
+def test_store_is_private_and_does_not_persist_upstream_token(tmp_path, monkeypatch):
+    token = "gitea-super-secret-token-material"
+    monkeypatch.setenv("GITEA_TOKEN", token)
+    path = tmp_path / "private-state" / "live.sqlite3"
+
+    LiveSnapshotStore(path)
+
+    assert os.stat(path).st_mode & 0o777 == 0o600
+    assert os.stat(path.parent).st_mode & 0o777 == 0o700
+    assert token.encode() not in path.read_bytes()
+
+
+def test_read_notification_filter_is_shared_with_future_publications(tmp_path):
+    path = tmp_path / "live.sqlite3"
+    first = LiveSnapshotStore(path, clock=lambda: 100.0)
+    second = LiveSnapshotStore(path, clock=lambda: 100.0)
+    lease = first.try_acquire_refresh({"notifications"}, lease_seconds=5)
+    assert lease is not None
+    metadata = {section: None for section in ("context", "events", "notifications")}
+    first.publish_refresh(
+        lease,
+        value={"notifications": [{"id": 7}, {"id": 8}], "sections": {"notifications": "fresh"}},
+        created_at=metadata,
+        failure_count={section: 0 for section in metadata},
+        retry_at=metadata,
+        changed_sections={"notifications"},
+    )
+
+    removed = second.remove_notification(7)
+    lease = first.try_acquire_refresh({"notifications"}, lease_seconds=5)
+    assert lease is not None
+    republished = first.publish_refresh(
+        lease,
+        value={"notifications": [{"id": 7}, {"id": 8}], "sections": {"notifications": "fresh"}},
+        created_at=metadata,
+        failure_count={section: 0 for section in metadata},
+        retry_at=metadata,
+        changed_sections={"notifications"},
+    )
+
+    assert removed.value["notifications"] == [{"id": 8}]
+    assert republished.value["notifications"] == [{"id": 8}]

tests/test_notification_read.py

@@ -5,6 +5,17 @@ import httpx
 import pytest
 
 from src import gitea_proxy, main
+from src.live_snapshot_store import LiveSnapshotStore
+
+
+@pytest.fixture(autouse=True)
+def isolated_live_snapshot_store(tmp_path, monkeypatch):
+    monkeypatch.setattr(
+        main,
+        "_live_snapshot_store",
+        LiveSnapshotStore(tmp_path / "live.sqlite3", clock=lambda: main._live_snapshot_clock()),
+    )
+    monkeypatch.setattr(main, "_read_notification_ids", set())
 
 
 @pytest.mark.anyio