136 lines
5.6 KiB
Python
136 lines
5.6 KiB
Python
"""Durable job state and atomic worker claims."""
|
|
|
|
import sqlite3
|
|
from contextlib import contextmanager
|
|
from uuid import uuid4
|
|
|
|
MAX_ATTEMPTS = 3
|
|
|
|
|
|
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"] >= MAX_ATTEMPTS 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 >= MAX_ATTEMPTS:
|
|
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],
|
|
}
|