Build Relay Lab reliability dashboard

2026-11-07 21:12:55 -07:00
commit c7298662a1
7 changed files with 597 additions and 0 deletions
@@ -0,0 +1,133 @@
+"""Durable job state and atomic worker claims."""
+
+import sqlite3
+from contextlib import contextmanager
+from uuid import uuid4
+
+
+class Conflict(Exception):
+    """An idempotency key was reused for a different request."""
+
+
+class QueueStore:
+    def __init__(self, path):
+        self.path = path
+        with self.connection() as db:
+            db.execute("PRAGMA journal_mode=WAL")
+            db.executescript("""
+                CREATE TABLE IF NOT EXISTS jobs (
+                    id TEXT PRIMARY KEY,
+                    request_key TEXT UNIQUE NOT NULL,
+                    scenario TEXT NOT NULL,
+                    state TEXT NOT NULL,
+                    attempts INTEGER NOT NULL DEFAULT 0,
+                    available_at REAL NOT NULL,
+                    lease_token TEXT,
+                    lease_until REAL
+                );
+                CREATE TABLE IF NOT EXISTS events (
+                    id INTEGER PRIMARY KEY,
+                    job_id TEXT NOT NULL REFERENCES jobs(id),
+                    at REAL NOT NULL,
+                    message TEXT NOT NULL
+                );
+                CREATE INDEX IF NOT EXISTS jobs_ready ON jobs(state, available_at);
+                CREATE INDEX IF NOT EXISTS jobs_leased ON jobs(state, lease_until);
+            """)
+
+    @contextmanager
+    def connection(self, write=False):
+        db = sqlite3.connect(self.path, timeout=5)
+        db.row_factory = sqlite3.Row
+        try:
+            if write:
+                db.execute("BEGIN IMMEDIATE")
+            with db:
+                yield db
+        finally:
+            db.close()
+
+    @staticmethod
+    def record(db, job_id, now, message):
+        db.execute(
+            "INSERT INTO events(job_id, at, message) VALUES (?, ?, ?)",
+            (job_id, now, message),
+        )
+
+    def enqueue(self, key, scenario, now):
+        with self.connection(write=True) as db:
+            prior = db.execute("SELECT * FROM jobs WHERE request_key = ?", (key,)).fetchone()
+            if prior:
+                if prior["scenario"] != scenario:
+                    raise Conflict("This key already belongs to a different scenario.")
+                self.record(db, prior["id"], now, "Duplicate request returned the existing job")
+                return prior["id"], False
+            job_id = uuid4().hex
+            db.execute(
+                "INSERT INTO jobs(id, request_key, scenario, state, available_at) VALUES (?, ?, ?, 'queued', ?)",
+                (job_id, key, scenario, now),
+            )
+            self.record(db, job_id, now, "Job queued")
+            return job_id, True
+
+    def claim(self, now, lease_seconds):
+        with self.connection(write=True) as db:
+            expired = db.execute(
+                "SELECT id, attempts FROM jobs WHERE state = 'running' AND lease_until <= ?",
+                (now,),
+            ).fetchall()
+            for row in expired:
+                state = "dead" if row["attempts"] >= 3 else "queued"
+                db.execute("""
+                    UPDATE jobs SET state = ?, available_at = ?, lease_token = NULL,
+                        lease_until = NULL WHERE id = ?
+                """, (state, now, row["id"]))
+                message = "Lease expired · attempts exhausted" if state == "dead" else "Lease expired · job recovered"
+                self.record(db, row["id"], now, message)
+            row = db.execute("""
+                SELECT * FROM jobs WHERE state = 'queued' AND available_at <= ?
+                ORDER BY available_at, rowid LIMIT 1
+            """, (now,)).fetchone()
+            if row is None:
+                return None
+            token = uuid4().hex
+            attempt = row["attempts"] + 1
+            db.execute("""
+                UPDATE jobs SET state = 'running', attempts = ?, lease_token = ?, lease_until = ?
+                WHERE id = ?
+            """, (attempt, token, now + lease_seconds, row["id"]))
+            self.record(db, row["id"], now, f"Worker claimed attempt {attempt}")
+            return {**dict(row), "attempts": attempt, "lease_token": token}
+
+    def finish(self, job, success, now):
+        attempt = job["attempts"]
+        delay = 2 ** attempt
+        if success:
+            state, message = "delivered", "Consumer acknowledged the event"
+        elif attempt >= 3:
+            state, message = "dead", "Consumer failed · attempts exhausted"
+        else:
+            state, message = "queued", f"Consumer failed · retry in {delay}s"
+        with self.connection(write=True) as db:
+            updated = db.execute("""
+                UPDATE jobs SET state = ?, available_at = ?, lease_token = NULL, lease_until = NULL
+                WHERE id = ? AND state = 'running' AND lease_token = ? AND lease_until > ?
+            """, (state, now + (delay if state == "queued" else 0), job["id"], job["lease_token"], now))
+            if updated.rowcount:
+                self.record(db, job["id"], now, message)
+            return updated.rowcount == 1
+
+    def snapshot(self):
+        with self.connection() as db:
+            db.execute("BEGIN")
+            counts = db.execute("SELECT state, COUNT(*) AS total FROM jobs GROUP BY state").fetchall()
+            jobs = db.execute("""
+                SELECT id, request_key, scenario, state, attempts, available_at
+                FROM jobs ORDER BY rowid DESC LIMIT 24
+            """).fetchall()
+            events = db.execute("SELECT * FROM events ORDER BY id DESC LIMIT 32").fetchall()
+        return {
+            "counts": {row["state"]: row["total"] for row in counts},
+            "jobs": [dict(row) for row in jobs],
+            "events": [dict(row) for row in events],
+        }