333 lines
14 KiB
Python
333 lines
14 KiB
Python
"""Process-shared live snapshot state and refresh coordination."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import secrets
|
|
import sqlite3
|
|
import stat
|
|
import time
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Callable, Iterable
|
|
|
|
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 = sqlite3.connect(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:
|
|
self.path.parent.mkdir(mode=0o700, parents=True, exist_ok=True)
|
|
try:
|
|
os.chmod(self.path.parent, stat.S_IRWXU)
|
|
except OSError:
|
|
pass
|
|
old_umask = os.umask(0o077)
|
|
try:
|
|
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),
|
|
),
|
|
)
|
|
finally:
|
|
os.umask(old_umask)
|
|
os.chmod(self.path, stat.S_IRUSR | stat.S_IWUSR)
|
|
|
|
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()
|