"""Process-shared live snapshot state and refresh coordination.""" from __future__ import annotations import json import os import secrets import sqlite3 import time from dataclasses import dataclass from pathlib import Path from typing import Callable, Iterable from src.private_state import connect_private_sqlite SECTIONS = ("context", "events", "notifications") class RefreshLeaseLost(RuntimeError): """The refresh is no longer authorized to publish shared state.""" @dataclass(frozen=True) class LiveSnapshotState: value: dict | None created_at: dict[str, float | None] failure_count: dict[str, int] retry_at: dict[str, float | None] revisions: dict[str, int] generation: str refreshing_sections: set[str] lease_expires_at: float | None @dataclass(frozen=True) class LiveSnapshotMetadata: """Shared snapshot coordination fields without the potentially large value.""" created_at: dict[str, float | None] failure_count: dict[str, int] retry_at: dict[str, float | None] revisions: dict[str, int] generation: str refreshing_sections: set[str] lease_expires_at: float | None class LiveSnapshotStore: """A private SQLite snapshot with an expiring, cross-process refresh lease.""" def __init__( self, path: str | os.PathLike[str], *, clock: Callable[[], float] | None = None, ): self.path = Path(path) self.clock = clock or time.time self._initialize() def _connect(self) -> sqlite3.Connection: connection = connect_private_sqlite(self.path, timeout=1.0, isolation_level=None) connection.row_factory = sqlite3.Row connection.execute("PRAGMA busy_timeout = 1000") return connection def _initialize(self) -> None: with self._connect() as connection: connection.executescript( """ CREATE TABLE IF NOT EXISTS live_snapshot ( singleton INTEGER PRIMARY KEY CHECK (singleton = 1), value_json TEXT, created_at_json TEXT NOT NULL, failure_count_json TEXT NOT NULL, retry_at_json TEXT NOT NULL, revisions_json TEXT NOT NULL, generation TEXT NOT NULL ); CREATE TABLE IF NOT EXISTS live_refresh_lease ( singleton INTEGER PRIMARY KEY CHECK (singleton = 1), owner TEXT NOT NULL, sections_json TEXT NOT NULL, expires_at REAL NOT NULL ); CREATE TABLE IF NOT EXISTS live_read_notification ( notification_id INTEGER PRIMARY KEY ); """ ) connection.execute( """INSERT OR IGNORE INTO live_snapshot VALUES (1, NULL, ?, ?, ?, ?, ?)""", ( json.dumps({section: None for section in SECTIONS}), json.dumps({section: 0 for section in SECTIONS}), json.dumps({section: None for section in SECTIONS}), json.dumps({section: 0 for section in SECTIONS}), secrets.token_hex(8), ), ) def try_acquire_refresh( self, sections: Iterable[str], *, lease_seconds: float ) -> str | None: requested = sorted(set(sections)) if not requested or any(section not in SECTIONS for section in requested): raise ValueError("refresh lease requires known sections") 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 live_refresh_lease WHERE singleton = 1" ).fetchone() if active is not None and active["expires_at"] > now: connection.rollback() return None connection.execute("DELETE FROM live_refresh_lease WHERE singleton = 1") connection.execute( "INSERT INTO live_refresh_lease VALUES (1, ?, ?, ?)", (owner, json.dumps(requested, separators=(",", ":")), now + lease_seconds), ) connection.commit() return owner def load(self) -> LiveSnapshotState: now = self.clock() with self._connect() as connection: row = connection.execute( "SELECT * FROM live_snapshot WHERE singleton = 1" ).fetchone() lease = connection.execute( "SELECT sections_json, expires_at FROM live_refresh_lease " "WHERE singleton = 1 AND expires_at > ?", (now,), ).fetchone() assert row is not None return LiveSnapshotState( value=json.loads(row["value_json"]) if row["value_json"] is not None else None, created_at=json.loads(row["created_at_json"]), failure_count=json.loads(row["failure_count_json"]), retry_at=json.loads(row["retry_at_json"]), revisions=json.loads(row["revisions_json"]), generation=row["generation"], refreshing_sections=set(json.loads(lease["sections_json"])) if lease else set(), lease_expires_at=lease["expires_at"] if lease else None, ) def load_metadata(self) -> LiveSnapshotMetadata: """Load freshness and coordination state without reading the snapshot value.""" now = self.clock() with self._connect() as connection: row = connection.execute( """SELECT created_at_json, failure_count_json, retry_at_json, revisions_json, generation FROM live_snapshot WHERE singleton = 1""" ).fetchone() lease = connection.execute( "SELECT sections_json, expires_at FROM live_refresh_lease " "WHERE singleton = 1 AND expires_at > ?", (now,), ).fetchone() assert row is not None return LiveSnapshotMetadata( created_at=json.loads(row["created_at_json"]), failure_count=json.loads(row["failure_count_json"]), retry_at=json.loads(row["retry_at_json"]), revisions=json.loads(row["revisions_json"]), generation=row["generation"], refreshing_sections=set(json.loads(lease["sections_json"])) if lease else set(), lease_expires_at=lease["expires_at"] if lease else None, ) def publish_refresh( self, owner: str, *, value: dict, created_at: dict[str, float | None], failure_count: dict[str, int], retry_at: dict[str, float | None], changed_sections: Iterable[str], ) -> LiveSnapshotState: changed = set(changed_sections) if changed - set(SECTIONS): raise ValueError("publication contains unknown sections") if any(set(mapping) != set(SECTIONS) for mapping in (created_at, failure_count, retry_at)): raise ValueError("publication metadata must contain every section") now = self.clock() with self._connect() as connection: connection.execute("BEGIN IMMEDIATE") lease = connection.execute( "SELECT owner, expires_at FROM live_refresh_lease WHERE singleton = 1" ).fetchone() if lease is None or lease["owner"] != owner or lease["expires_at"] <= now: connection.rollback() raise RefreshLeaseLost("live refresh lease expired or changed owner") row = connection.execute( "SELECT revisions_json FROM live_snapshot WHERE singleton = 1" ).fetchone() revisions = json.loads(row["revisions_json"]) notifications = value.get("notifications") if isinstance(notifications, list): returned_ids = { item.get("id") for item in notifications if isinstance(item, dict) } read_ids = { item["notification_id"] for item in connection.execute( "SELECT notification_id FROM live_read_notification" ) } value = dict(value) value["notifications"] = [ item for item in notifications if not isinstance(item, dict) or item.get("id") not in read_ids ] if returned_ids: placeholders = ",".join("?" for _ in returned_ids) connection.execute( f"DELETE FROM live_read_notification WHERE notification_id NOT IN ({placeholders})", tuple(returned_ids), ) else: connection.execute("DELETE FROM live_read_notification") for section in changed: revisions[section] += 1 connection.execute( """UPDATE live_snapshot SET value_json = ?, created_at_json = ?, failure_count_json = ?, retry_at_json = ?, revisions_json = ? WHERE singleton = 1""", ( json.dumps(value, separators=(",", ":")), json.dumps(created_at, separators=(",", ":")), json.dumps(failure_count, separators=(",", ":")), json.dumps(retry_at, separators=(",", ":")), json.dumps(revisions, separators=(",", ":")), ), ) connection.execute( "DELETE FROM live_refresh_lease WHERE singleton = 1 AND owner = ?", (owner,) ) connection.commit() return self.load() def remove_notifications(self, notification_ids: Iterable[int]) -> LiveSnapshotState: """Filter confirmed reads from every worker and subsequent stale refreshes.""" read_ids = set(notification_ids) if not read_ids: return self.load() with self._connect() as connection: connection.execute("BEGIN IMMEDIATE") connection.executemany( "INSERT OR IGNORE INTO live_read_notification VALUES (?)", ((notification_id,) for notification_id in read_ids), ) row = connection.execute( "SELECT value_json, revisions_json FROM live_snapshot WHERE singleton = 1" ).fetchone() value = json.loads(row["value_json"]) if row["value_json"] else None revisions = json.loads(row["revisions_json"]) if value is not None and isinstance(value.get("notifications"), list): previous = value["notifications"] retained = [ item for item in previous if not isinstance(item, dict) or item.get("id") not in read_ids ] if retained != previous: value["notifications"] = retained revisions["notifications"] += 1 connection.execute( "UPDATE live_snapshot SET value_json = ?, revisions_json = ? WHERE singleton = 1", ( json.dumps(value, separators=(",", ":")), json.dumps(revisions, separators=(",", ":")), ), ) connection.commit() return self.load() def remove_notification(self, notification_id: int) -> LiveSnapshotState: """Filter one confirmed read from every worker and subsequent stale refresh.""" return self.remove_notifications([notification_id]) def release_refresh(self, owner: str) -> None: """Relinquish a lease that became unnecessary after a coherent recheck.""" with self._connect() as connection: connection.execute( "DELETE FROM live_refresh_lease WHERE singleton = 1 AND owner = ?", (owner,), ) def fail_refresh( self, owner: str, *, failure_count: dict[str, int], retry_at: dict[str, float | None], ) -> LiveSnapshotState: """Publish backoff metadata and relinquish a failed refresh atomically.""" now = self.clock() with self._connect() as connection: connection.execute("BEGIN IMMEDIATE") lease = connection.execute( "SELECT owner, expires_at FROM live_refresh_lease WHERE singleton = 1" ).fetchone() if lease is None or lease["owner"] != owner or lease["expires_at"] <= now: connection.rollback() raise RefreshLeaseLost("live refresh lease expired or changed owner") connection.execute( """UPDATE live_snapshot SET failure_count_json = ?, retry_at_json = ? WHERE singleton = 1""", ( json.dumps(failure_count, separators=(",", ":")), json.dumps(retry_at, separators=(",", ":")), ), ) connection.execute("DELETE FROM live_refresh_lease WHERE singleton = 1") connection.commit() return self.load()