From 20d517671388be0adf511cad5013ead2ee91b3dd Mon Sep 17 00:00:00 2001 From: timmy Date: Mon, 10 Aug 2026 04:40:15 +0000 Subject: [PATCH] perf: share Find Work catalog across workers (#453) --- README.md | 8 +- src/available_issue_snapshot_store.py | 198 +++++++++++++++++++ src/main.py | 105 ++++++++-- tests/test_available_issue_snapshot_store.py | 90 +++++++++ tests/test_gitea_work_search.py | 35 +++- 5 files changed, 417 insertions(+), 19 deletions(-) create mode 100644 src/available_issue_snapshot_store.py create mode 100644 tests/test_available_issue_snapshot_store.py diff --git a/README.md b/README.md index 52febef..79a3a0a 100644 --- a/README.md +++ b/README.md @@ -241,7 +241,13 @@ same stale snapshot while that refresh runs, and can recover an abandoned lease Set `STACKCHAIN_LIVE_SNAPSHOT_DB` to override the default `STACKCHAIN_STATE_DIR/live-snapshot.sqlite3`; keep the containing directory on private, worker-shared writable storage. The database and directory are restricted to the service -account and never contain the Gitea token. Each bounded opaque revision token includes the +account and never contain the Gitea token. Find Work uses the same worker-shared pattern: +`STACKCHAIN_AVAILABLE_ISSUE_SNAPSHOT_DB` overrides +`STACKCHAIN_STATE_DIR/available-issue-snapshot.sqlite3`. One expiring lease bounds each +catalog scan across the deployment, shared retry metadata prevents worker-by-worker retry +bursts, and confirmed claims are removed from every worker's retained catalog. Releasing +an assignment invalidates the shared catalog so the newly available issue can be discovered +by the next authoritative scan. Each bounded opaque revision token includes the store generation, so a token from a different deployment or before replacement of the store cannot suppress different content. The browser sends its known tokens on later polls, so `/api/v1/live` can omit unchanged section bodies while still returning current freshness and diff --git a/src/available_issue_snapshot_store.py b/src/available_issue_snapshot_store.py new file mode 100644 index 0000000..a40102a --- /dev/null +++ b/src/available_issue_snapshot_store.py @@ -0,0 +1,198 @@ +"""Worker-shared Find Work catalog state.""" + +from __future__ import annotations + +import json +import os +import secrets +import stat +import sqlite3 +import time +from dataclasses import dataclass +from pathlib import Path + + +class RefreshLeaseLost(RuntimeError): + pass + + +@dataclass(frozen=True) +class AvailableIssueSnapshotState: + items: list[dict] | None + created_at: float | None + retry_at: float | None + refreshing: bool + lease_expires_at: float | None + + +class AvailableIssueSnapshotStore: + def __init__(self, path, *, clock=None): + self.path = Path(path) + self.clock = clock or time.time + self.path.parent.mkdir(mode=0o700, parents=True, exist_ok=True) + os.chmod(self.path.parent, stat.S_IRWXU) + old_umask = os.umask(0o077) + try: + with self._connect() as connection: + connection.executescript( + """ + CREATE TABLE IF NOT EXISTS available_issue_snapshot ( + singleton INTEGER PRIMARY KEY CHECK (singleton = 1), + items_json TEXT, + created_at REAL, + retry_at REAL + ); + CREATE TABLE IF NOT EXISTS available_issue_refresh_lease ( + singleton INTEGER PRIMARY KEY CHECK (singleton = 1), + owner TEXT NOT NULL, + expires_at REAL NOT NULL + ); + CREATE TABLE IF NOT EXISTS available_issue_claim ( + repository TEXT NOT NULL, + number INTEGER NOT NULL, + PRIMARY KEY (repository, number) + ); + """ + ) + connection.execute( + "INSERT OR IGNORE INTO available_issue_snapshot VALUES (1, NULL, NULL, NULL)" + ) + finally: + os.umask(old_umask) + os.chmod(self.path, stat.S_IRUSR | stat.S_IWUSR) + + def _connect(self): + connection = sqlite3.connect(self.path, timeout=1.0, isolation_level=None) + connection.row_factory = sqlite3.Row + connection.execute("PRAGMA busy_timeout = 1000") + return connection + + def try_acquire_refresh(self, *, lease_seconds: float) -> str | None: + if lease_seconds <= 0: + raise ValueError("refresh lease duration must be positive") + now = self.clock() + owner = secrets.token_hex(16) + with self._connect() as connection: + connection.execute("BEGIN IMMEDIATE") + active = connection.execute( + "SELECT expires_at FROM available_issue_refresh_lease WHERE singleton = 1" + ).fetchone() + if active is not None and active["expires_at"] > now: + connection.rollback() + return None + connection.execute("DELETE FROM available_issue_refresh_lease WHERE singleton = 1") + connection.execute( + "INSERT INTO available_issue_refresh_lease VALUES (1, ?, ?)", + (owner, now + lease_seconds), + ) + connection.commit() + return owner + + def load(self) -> AvailableIssueSnapshotState: + now = self.clock() + with self._connect() as connection: + row = connection.execute( + "SELECT * FROM available_issue_snapshot WHERE singleton = 1" + ).fetchone() + lease = connection.execute( + "SELECT expires_at FROM available_issue_refresh_lease " + "WHERE singleton = 1 AND expires_at > ?", (now,) + ).fetchone() + return AvailableIssueSnapshotState( + items=json.loads(row["items_json"]) if row["items_json"] else None, + created_at=row["created_at"], + retry_at=row["retry_at"], + refreshing=lease is not None, + lease_expires_at=lease["expires_at"] if lease else None, + ) + + def publish(self, owner: str | None, *, items: list[dict]) -> AvailableIssueSnapshotState: + now = self.clock() + with self._connect() as connection: + connection.execute("BEGIN IMMEDIATE") + lease = connection.execute( + "SELECT owner, expires_at FROM available_issue_refresh_lease WHERE singleton = 1" + ).fetchone() + if ( + owner is None + or lease is None + or lease["owner"] != owner + or lease["expires_at"] <= now + ): + connection.rollback() + raise RefreshLeaseLost("available issue refresh lease expired or changed owner") + claimed = { + (row["repository"], row["number"]) + for row in connection.execute( + "SELECT repository, number FROM available_issue_claim" + ) + } + items = [ + item for item in items + if (item.get("repository"), item.get("number")) not in claimed + ] + connection.execute( + "UPDATE available_issue_snapshot SET items_json = ?, created_at = ?, " + "retry_at = NULL WHERE singleton = 1", + (json.dumps(items, separators=(",", ":")), now), + ) + connection.execute("DELETE FROM available_issue_refresh_lease WHERE singleton = 1") + connection.commit() + return self.load() + + def remove_claimed(self, repository: str, number: int) -> AvailableIssueSnapshotState: + with self._connect() as connection: + connection.execute("BEGIN IMMEDIATE") + connection.execute( + "INSERT OR IGNORE INTO available_issue_claim VALUES (?, ?)", + (repository, number), + ) + row = connection.execute( + "SELECT items_json FROM available_issue_snapshot WHERE singleton = 1" + ).fetchone() + items = json.loads(row["items_json"]) if row["items_json"] else None + if items is not None: + items = [ + item for item in items + if item.get("repository") != repository or item.get("number") != number + ] + connection.execute( + "UPDATE available_issue_snapshot SET items_json = ? WHERE singleton = 1", + (json.dumps(items, separators=(",", ":")),), + ) + connection.commit() + return self.load() + + def invalidate(self) -> AvailableIssueSnapshotState: + with self._connect() as connection: + connection.execute("BEGIN IMMEDIATE") + connection.execute( + "UPDATE available_issue_snapshot SET items_json = NULL, " + "created_at = NULL, retry_at = NULL WHERE singleton = 1" + ) + connection.execute("DELETE FROM available_issue_claim") + connection.commit() + return self.load() + + def fail_refresh(self, owner: str | None, *, retry_at: float) -> AvailableIssueSnapshotState: + now = self.clock() + with self._connect() as connection: + connection.execute("BEGIN IMMEDIATE") + lease = connection.execute( + "SELECT owner, expires_at FROM available_issue_refresh_lease WHERE singleton = 1" + ).fetchone() + if ( + owner is None + or lease is None + or lease["owner"] != owner + or lease["expires_at"] <= now + ): + connection.rollback() + raise RefreshLeaseLost("available issue refresh lease expired or changed owner") + connection.execute( + "UPDATE available_issue_snapshot SET retry_at = ? WHERE singleton = 1", + (retry_at,), + ) + connection.execute("DELETE FROM available_issue_refresh_lease WHERE singleton = 1") + connection.commit() + return self.load() diff --git a/src/main.py b/src/main.py index f5a25e6..71a6b2d 100644 --- a/src/main.py +++ b/src/main.py @@ -19,6 +19,7 @@ from fastapi.staticfiles import StaticFiles from pydantic import BaseModel, Field, PositiveInt, field_validator, model_validator from src import dashboard_auth, gitea_proxy +from src.available_issue_snapshot_store import AvailableIssueSnapshotStore from src.compression import NegotiatedGZipMiddleware from src.gitea_proxy import ( activity_events, @@ -94,6 +95,7 @@ LIVE_SNAPSHOT_RETRY_BASE_SECONDS = 5.0 LIVE_SNAPSHOT_RETRY_MAX_SECONDS = 60.0 AVAILABLE_ISSUE_SNAPSHOT_FRESHNESS_SECONDS = 15.0 AVAILABLE_ISSUE_SNAPSHOT_RETRY_SECONDS = 5.0 +AVAILABLE_ISSUE_SNAPSHOT_LEASE_SECONDS = WORK_PAGE_TIMEOUT_SECONDS + 1.0 FRONTEND_DIR = Path(__file__).resolve().parent.parent / "frontend" _live_snapshot_task: asyncio.Task | None = None _live_snapshot_value: dict | None = None @@ -133,9 +135,16 @@ _idempotency_ledger = IdempotencyLedger( lock_timeout_seconds=AUTHORED_ACTION_LEDGER_LOCK_TIMEOUT_SECONDS, ) _available_issue_snapshot_task: asyncio.Task | None = None +_available_issue_snapshot_lock = asyncio.Lock() _available_issue_snapshot_value: list[dict] | None = None _available_issue_snapshot_created_at: float | None = None _available_issue_snapshot_retry_at: float | None = None +_available_issue_snapshot_store = AvailableIssueSnapshotStore( + os.getenv( + "STACKCHAIN_AVAILABLE_ISSUE_SNAPSHOT_DB", + str(_state_dir / "available-issue-snapshot.sqlite3"), + ) +) class ContextPayloadError(ValueError): @@ -1365,14 +1374,30 @@ async def paged_work( }) -async def _refresh_available_issue_snapshot() -> list[dict]: +async def _refresh_available_issue_snapshot(lease_owner: str) -> list[dict]: global _available_issue_snapshot_value, _available_issue_snapshot_created_at global _available_issue_snapshot_retry_at - result = await gitea_proxy.available_issue_snapshot() - _available_issue_snapshot_value = result - _available_issue_snapshot_created_at = time.monotonic() - _available_issue_snapshot_retry_at = None - return result + try: + result = await gitea_proxy.available_issue_snapshot() + shared = await asyncio.to_thread( + _available_issue_snapshot_store.publish, lease_owner, items=result + ) + _available_issue_snapshot_value = shared.items + _available_issue_snapshot_created_at = time.monotonic() + _available_issue_snapshot_retry_at = None + return shared.items or [] + except asyncio.CancelledError: + raise + except Exception: + try: + await asyncio.to_thread( + _available_issue_snapshot_store.fail_refresh, + lease_owner, + retry_at=time.time() + AVAILABLE_ISSUE_SNAPSHOT_RETRY_SECONDS, + ) + except Exception: + pass + raise def _observe_available_issue_refresh(task: asyncio.Task) -> None: @@ -1384,8 +1409,43 @@ def _observe_available_issue_refresh(task: asyncio.Task) -> None: ) +async def _start_available_issue_refresh() -> bool: + """Atomically start this worker's refresh when the shared lease is available.""" + global _available_issue_snapshot_task + async with _available_issue_snapshot_lock: + if ( + _available_issue_snapshot_task is not None + and not _available_issue_snapshot_task.done() + ): + return True + lease_owner = await asyncio.to_thread( + _available_issue_snapshot_store.try_acquire_refresh, + lease_seconds=AVAILABLE_ISSUE_SNAPSHOT_LEASE_SECONDS, + ) + if lease_owner is None: + return False + _available_issue_snapshot_task = asyncio.create_task( + _refresh_available_issue_snapshot(lease_owner) + ) + _available_issue_snapshot_task.add_done_callback(_observe_available_issue_refresh) + return True + + async def _available_issue_snapshot() -> tuple[list[dict], bool, bool, bool]: global _available_issue_snapshot_task + global _available_issue_snapshot_value, _available_issue_snapshot_created_at + shared = await asyncio.to_thread(_available_issue_snapshot_store.load) + wall_now = time.time() + if shared.items is not None: + _available_issue_snapshot_value = shared.items + if ( + shared.created_at is not None + and wall_now - shared.created_at < AVAILABLE_ISSUE_SNAPSHOT_FRESHNESS_SECONDS + ): + _available_issue_snapshot_created_at = time.monotonic() + return shared.items, False, False, False + if shared.retry_at is not None and wall_now < shared.retry_at: + return shared.items, True, False, True now = time.monotonic() if ( _available_issue_snapshot_value is not None @@ -1400,13 +1460,18 @@ async def _available_issue_snapshot() -> tuple[list[dict], bool, bool, bool]: and now < _available_issue_snapshot_retry_at ): return _available_issue_snapshot_value, True, False, True - if _available_issue_snapshot_task is None or _available_issue_snapshot_task.done(): - _available_issue_snapshot_task = asyncio.create_task( - _refresh_available_issue_snapshot() - ) - _available_issue_snapshot_task.add_done_callback(_observe_available_issue_refresh) + local_refresh = await _start_available_issue_refresh() if _available_issue_snapshot_value is not None: return _available_issue_snapshot_value, True, True, False + if not local_refresh: + for _ in range(50): + await asyncio.sleep(0.02) + shared = await asyncio.to_thread(_available_issue_snapshot_store.load) + if shared.items is not None: + return shared.items, False, False, False + if not shared.refreshing: + break + raise RuntimeError("available issue catalog refresh is owned by another worker") try: return await asyncio.shield(_available_issue_snapshot_task), False, False, False except Exception: @@ -2154,6 +2219,7 @@ async def reply_to_notification( async def claim_available_issue( owner: str, repo: str, number: int = PathParam(gt=0) ) -> JSONResponse: + global _available_issue_snapshot_value repository = f"{owner}/{repo}" try: result = await asyncio.wait_for( @@ -2171,12 +2237,16 @@ async def claim_available_issue( status_code=503, headers={"Retry-After": "1"}, ) - if _available_issue_snapshot_value is not None: - _available_issue_snapshot_value[:] = [ - item for item in _available_issue_snapshot_value - if not ( - item.get("repository") == repository and item.get("number") == number - ) + retained = _available_issue_snapshot_value + shared = await asyncio.to_thread( + _available_issue_snapshot_store.remove_claimed, repository, number + ) + if shared.items is not None: + _available_issue_snapshot_value = shared.items + elif retained is not None: + _available_issue_snapshot_value = [ + item for item in retained + if item.get("repository") != repository or item.get("number") != number ] return JSONResponse(result) @@ -2201,6 +2271,7 @@ async def release_assigned_issue( headers={"Retry-After": "1"}, ) if result.get("available"): + await asyncio.to_thread(_available_issue_snapshot_store.invalidate) _available_issue_snapshot_value = None _available_issue_snapshot_created_at = None return JSONResponse(result) diff --git a/tests/test_available_issue_snapshot_store.py b/tests/test_available_issue_snapshot_store.py new file mode 100644 index 0000000..7828151 --- /dev/null +++ b/tests/test_available_issue_snapshot_store.py @@ -0,0 +1,90 @@ +import threading +import os + +from src.available_issue_snapshot_store import AvailableIssueSnapshotStore + + +def test_independent_workers_allow_only_one_catalog_refresh(tmp_path): + path = tmp_path / "available.sqlite3" + first = AvailableIssueSnapshotStore(path, clock=lambda: 100.0) + second = AvailableIssueSnapshotStore(path, clock=lambda: 100.0) + barrier = threading.Barrier(2) + results = [] + + def acquire(store): + barrier.wait() + results.append(store.try_acquire_refresh(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(owner is not None for owner in results) == 1 + + +def test_published_catalog_and_metadata_are_visible_to_another_worker(tmp_path): + path = tmp_path / "available.sqlite3" + writer = AvailableIssueSnapshotStore(path, clock=lambda: 100.0) + reader = AvailableIssueSnapshotStore(path, clock=lambda: 100.0) + owner = writer.try_acquire_refresh(lease_seconds=5) + + writer.publish(owner, items=[{"repository": "stackchain/api", "number": 7}]) + + state = reader.load() + assert state.items == [{"repository": "stackchain/api", "number": 7}] + assert state.created_at == 100.0 + assert state.retry_at is None + assert state.refreshing is False + + +def test_failed_refresh_backoff_is_shared_and_lease_is_released(tmp_path): + path = tmp_path / "available.sqlite3" + writer = AvailableIssueSnapshotStore(path, clock=lambda: 100.0) + reader = AvailableIssueSnapshotStore(path, clock=lambda: 100.0) + owner = writer.try_acquire_refresh(lease_seconds=5) + + writer.fail_refresh(owner, retry_at=105.0) + + state = reader.load() + assert state.retry_at == 105.0 + assert state.refreshing is False + assert reader.try_acquire_refresh(lease_seconds=5) is not None + + +def test_confirmed_claim_is_removed_for_every_worker(tmp_path): + path = tmp_path / "available.sqlite3" + writer = AvailableIssueSnapshotStore(path, clock=lambda: 100.0) + reader = AvailableIssueSnapshotStore(path, clock=lambda: 100.0) + owner = writer.try_acquire_refresh(lease_seconds=5) + writer.publish(owner, items=[ + {"repository": "stackchain/api", "number": 7}, + {"repository": "stackchain/web", "number": 8}, + ]) + + writer.remove_claimed("stackchain/api", 7) + + assert reader.load().items == [{"repository": "stackchain/web", "number": 8}] + + +def test_invalidation_clears_catalog_and_claim_filters(tmp_path): + store = AvailableIssueSnapshotStore(tmp_path / "available.sqlite3", clock=lambda: 100.0) + owner = store.try_acquire_refresh(lease_seconds=5) + store.publish(owner, items=[{"repository": "stackchain/api", "number": 7}]) + store.remove_claimed("stackchain/api", 7) + + store.invalidate() + owner = store.try_acquire_refresh(lease_seconds=5) + store.publish(owner, items=[{"repository": "stackchain/api", "number": 7}]) + + assert store.load().items == [{"repository": "stackchain/api", "number": 7}] + + +def test_catalog_store_is_private(tmp_path): + path = tmp_path / "state" / "available.sqlite3" + + AvailableIssueSnapshotStore(path) + + assert os.stat(path.parent).st_mode & 0o777 == 0o700 + assert os.stat(path).st_mode & 0o777 == 0o600 diff --git a/tests/test_gitea_work_search.py b/tests/test_gitea_work_search.py index 1642781..2bbc00f 100644 --- a/tests/test_gitea_work_search.py +++ b/tests/test_gitea_work_search.py @@ -6,11 +6,18 @@ import httpx from src import gitea_proxy from src import main +from src.available_issue_snapshot_store import AvailableIssueSnapshotStore @pytest.fixture(autouse=True) -def reset_available_issue_snapshot(): +def reset_available_issue_snapshot(tmp_path, monkeypatch): + monkeypatch.setattr( + main, + "_available_issue_snapshot_store", + AvailableIssueSnapshotStore(tmp_path / "available-issues.sqlite3"), + ) main._available_issue_snapshot_task = None + main._available_issue_snapshot_lock = asyncio.Lock() main._available_issue_snapshot_value = None main._available_issue_snapshot_created_at = None main._available_issue_snapshot_retry_at = None @@ -19,6 +26,7 @@ def reset_available_issue_snapshot(): if task is not None and not task.done(): task.cancel() main._available_issue_snapshot_task = None + main._available_issue_snapshot_lock = asyncio.Lock() main._available_issue_snapshot_value = None main._available_issue_snapshot_created_at = None main._available_issue_snapshot_retry_at = None @@ -207,6 +215,31 @@ async def test_available_issue_endpoint_is_bounded_retryable_and_no_store(monkey assert calls == [True] +@pytest.mark.anyio +async def test_available_issue_endpoint_reuses_catalog_published_by_another_worker( + monkeypatch, tmp_path +): + store = AvailableIssueSnapshotStore(tmp_path / "available.sqlite3", clock=lambda: 100.0) + owner = store.try_acquire_refresh(lease_seconds=5) + store.publish(owner, items=[{"repository": "stackchain/api", "number": 7}]) + monkeypatch.setattr(main, "_available_issue_snapshot_store", store, raising=False) + monkeypatch.setattr(main.time, "time", lambda: 100.0) + + async def must_not_scan(): + raise AssertionError("fresh shared catalog should avoid an upstream scan") + + monkeypatch.setattr(main.gitea_proxy, "available_issue_snapshot", must_not_scan) + transport = httpx.ASGITransport(app=main.app) + async with httpx.AsyncClient(transport=transport, base_url="http://test") as client: + response = await client.get("/api/v1/available-issues?page=1") + + assert response.status_code == 200 + assert response.json() == { + "items": [{"repository": "stackchain/api", "number": 7}], + "page": 1, "total": 1, "has_more": False, + } + + @pytest.mark.anyio async def test_available_issue_endpoint_coalesces_cold_scan_and_reuses_it_for_pages(monkeypatch): calls = 0