Add game data capture vertical slice
This commit is contained in:
@@ -0,0 +1,157 @@
|
||||
"""Bounded durable leases that follow the frozen job state machine."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sqlite3
|
||||
from datetime import UTC, datetime, timedelta
|
||||
|
||||
from game_data_capture.domain.state_machines import require_transition
|
||||
from game_data_capture.storage.database import StorageDatabase
|
||||
from game_data_capture.storage.repository import StorageError, _id
|
||||
|
||||
_KINDS = {"classify", "ocr", "vision_fallback", "frame_extract", "export"}
|
||||
|
||||
|
||||
def _now() -> datetime:
|
||||
return datetime.now(UTC)
|
||||
|
||||
|
||||
def _lease(seconds: int) -> None:
|
||||
if isinstance(seconds, bool) or not isinstance(seconds, int) or not 1 <= seconds <= 300:
|
||||
raise ValueError("lease duration must be from 1 to 300 seconds")
|
||||
|
||||
|
||||
class JobLeases:
|
||||
def __init__(self, database: StorageDatabase) -> None:
|
||||
self.database = database
|
||||
|
||||
def create(
|
||||
self, job_id: str, kind: str, capture_id: str | None = None, max_attempts: int = 3
|
||||
) -> None:
|
||||
_id(job_id)
|
||||
if capture_id is not None:
|
||||
_id(capture_id)
|
||||
if (
|
||||
kind not in _KINDS
|
||||
or isinstance(max_attempts, bool)
|
||||
or not isinstance(max_attempts, int)
|
||||
or not 1 <= max_attempts <= 20
|
||||
):
|
||||
raise ValueError("invalid job definition")
|
||||
with self.database.connect() as connection:
|
||||
connection.execute(
|
||||
"INSERT INTO jobs(job_id, capture_id, kind, state, max_attempts) "
|
||||
"VALUES (?, ?, ?, 'queued', ?)",
|
||||
(job_id, capture_id, kind, max_attempts),
|
||||
)
|
||||
|
||||
def _recover_expired(self, connection: sqlite3.Connection, now: datetime) -> None:
|
||||
# State changes use the contract path: running -> retry_wait -> queued/failed.
|
||||
rows = list(
|
||||
connection.execute(
|
||||
"SELECT job_id, attempt, max_attempts FROM jobs "
|
||||
"WHERE state = 'running' AND lease_expires_at < ?",
|
||||
(now.isoformat(),),
|
||||
)
|
||||
)
|
||||
for row in rows:
|
||||
require_transition("job", "running", "retry_wait")
|
||||
changed = connection.execute(
|
||||
"UPDATE jobs SET state = 'retry_wait', lease_owner = NULL, "
|
||||
"lease_expires_at = NULL, revision = revision + 1 "
|
||||
"WHERE job_id = ? AND state = 'running'",
|
||||
(row["job_id"],),
|
||||
).rowcount
|
||||
if changed != 1:
|
||||
raise StorageError("expired job changed concurrently")
|
||||
target = "failed" if row["attempt"] >= row["max_attempts"] else "queued"
|
||||
require_transition("job", "retry_wait", target)
|
||||
changed = connection.execute(
|
||||
"UPDATE jobs SET state = ?, revision = revision + 1 "
|
||||
"WHERE job_id = ? AND state = 'retry_wait'",
|
||||
(target, row["job_id"]),
|
||||
).rowcount
|
||||
if changed != 1:
|
||||
raise StorageError("expired job changed concurrently")
|
||||
|
||||
def acquire(self, owner: str, lease_seconds: int) -> str | None:
|
||||
_id(owner)
|
||||
_lease(lease_seconds)
|
||||
now = _now()
|
||||
with self.database.connect() as connection:
|
||||
connection.execute("BEGIN IMMEDIATE")
|
||||
try:
|
||||
self._recover_expired(connection, now)
|
||||
row = connection.execute(
|
||||
"SELECT job_id FROM jobs WHERE state = 'queued' "
|
||||
"AND attempt < max_attempts ORDER BY job_id LIMIT 1"
|
||||
).fetchone()
|
||||
if row is None:
|
||||
connection.execute("COMMIT")
|
||||
return None
|
||||
result = connection.execute(
|
||||
"UPDATE jobs SET state = 'running', lease_owner = ?, lease_expires_at = ?, "
|
||||
"attempt = attempt + 1, revision = revision + 1 WHERE job_id = ? "
|
||||
"AND state = 'queued' AND attempt < max_attempts",
|
||||
(owner, (now + timedelta(seconds=lease_seconds)).isoformat(), row["job_id"]),
|
||||
)
|
||||
connection.execute("COMMIT")
|
||||
return str(row["job_id"]) if result.rowcount == 1 else None
|
||||
except Exception:
|
||||
if connection.in_transaction:
|
||||
connection.execute("ROLLBACK")
|
||||
raise
|
||||
|
||||
def renew(self, job_id: str, owner: str, lease_seconds: int) -> None:
|
||||
_id(job_id)
|
||||
_id(owner)
|
||||
_lease(lease_seconds)
|
||||
now = _now()
|
||||
with self.database.connect() as connection:
|
||||
result = connection.execute(
|
||||
"UPDATE jobs SET lease_expires_at = ?, revision = revision + 1 "
|
||||
"WHERE job_id = ? AND state = 'running' AND lease_owner = ? "
|
||||
"AND lease_expires_at >= ?",
|
||||
(
|
||||
(now + timedelta(seconds=lease_seconds)).isoformat(),
|
||||
job_id,
|
||||
owner,
|
||||
now.isoformat(),
|
||||
),
|
||||
)
|
||||
if result.rowcount != 1:
|
||||
raise StorageError("lease is not owned or has expired")
|
||||
|
||||
def finish(self, job_id: str, owner: str, target: str) -> None:
|
||||
_id(job_id)
|
||||
_id(owner)
|
||||
if target not in {"succeeded", "failed", "cancelled"}:
|
||||
raise ValueError("job finish target must be terminal")
|
||||
require_transition("job", "running", target)
|
||||
now = _now()
|
||||
with self.database.connect() as connection:
|
||||
result = connection.execute(
|
||||
"UPDATE jobs SET state = ?, lease_owner = NULL, lease_expires_at = NULL, "
|
||||
"revision = revision + 1 WHERE job_id = ? AND state = 'running' "
|
||||
"AND lease_owner = ? AND lease_expires_at >= ?",
|
||||
(target, job_id, owner, now.isoformat()),
|
||||
)
|
||||
if result.rowcount != 1:
|
||||
raise StorageError("lease is not owned or has expired")
|
||||
|
||||
def recover(self) -> int:
|
||||
now = _now()
|
||||
with self.database.connect() as connection:
|
||||
connection.execute("BEGIN IMMEDIATE")
|
||||
try:
|
||||
before = connection.execute(
|
||||
"SELECT count(*) FROM jobs WHERE state = 'running' AND lease_expires_at < ?",
|
||||
(now.isoformat(),),
|
||||
).fetchone()[0]
|
||||
self._recover_expired(connection, now)
|
||||
connection.execute("COMMIT")
|
||||
return int(before)
|
||||
except Exception:
|
||||
if connection.in_transaction:
|
||||
connection.execute("ROLLBACK")
|
||||
raise
|
||||
Reference in New Issue
Block a user