209 lines
8.2 KiB
Python
209 lines
8.2 KiB
Python
"""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 release_refresh(self, owner: str | None) -> None:
|
|
if owner is None:
|
|
return
|
|
with self._connect() as connection:
|
|
connection.execute(
|
|
"DELETE FROM available_issue_refresh_lease "
|
|
"WHERE singleton = 1 AND owner = ?",
|
|
(owner,),
|
|
)
|
|
|
|
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()
|