394 lines
16 KiB
Python
394 lines
16 KiB
Python
"""JobService — durable job coordinator (no execution; the worker lands in US02-03).
|
|
|
|
Responsibilities: enqueue with idempotency keys and one-mutating-job locks, hand a
|
|
queued job to a worker atomically (compare-and-set claim with a fencing token),
|
|
track heartbeats/leases, validate every state transition, derive progress from
|
|
durable item state, request cancellation, and recover jobs whose worker vanished.
|
|
|
|
State machine (concept §3):
|
|
|
|
queued ─▶ running ─▶ succeeded
|
|
│ │ ╲──▶ failed ─▶ retry_queued ─▶ running
|
|
│ ╲──▶ cancelling ─▶ cancelled
|
|
╰──▶ cancelled
|
|
|
|
Terminal states (succeeded, cancelled) have no outgoing transition, so a duplicate
|
|
terminal transition is rejected.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import uuid
|
|
from collections import Counter
|
|
from datetime import datetime, timedelta, timezone
|
|
from collections.abc import Iterable, Sequence
|
|
|
|
from sqlalchemy import select, text, update
|
|
from sqlalchemy.exc import IntegrityError
|
|
from sqlalchemy.orm import sessionmaker
|
|
|
|
from photo_pipeline.models import Job, JobEvent, JobItem
|
|
|
|
|
|
class JobState:
|
|
QUEUED = "queued"
|
|
RUNNING = "running"
|
|
CANCELLING = "cancelling"
|
|
CANCELLED = "cancelled"
|
|
SUCCEEDED = "succeeded"
|
|
FAILED = "failed"
|
|
RETRY_QUEUED = "retry_queued"
|
|
|
|
|
|
class ItemState:
|
|
QUEUED = "queued"
|
|
RUNNING = "running"
|
|
SUCCEEDED = "succeeded"
|
|
FAILED = "failed"
|
|
CANCELLED = "cancelled"
|
|
|
|
|
|
# Jobs that occupy a lock or are eligible to run are "active" for lock/one-job checks.
|
|
ACTIVE_STATES = frozenset(
|
|
{JobState.QUEUED, JobState.RUNNING, JobState.CANCELLING, JobState.RETRY_QUEUED}
|
|
)
|
|
TERMINAL_STATES = frozenset({JobState.SUCCEEDED, JobState.CANCELLED})
|
|
CLAIMABLE_STATES = frozenset({JobState.QUEUED, JobState.RETRY_QUEUED})
|
|
|
|
ALLOWED_TRANSITIONS = {
|
|
JobState.QUEUED: {JobState.RUNNING, JobState.CANCELLED, JobState.CANCELLING},
|
|
JobState.RUNNING: {JobState.SUCCEEDED, JobState.FAILED, JobState.CANCELLING},
|
|
JobState.CANCELLING: {JobState.CANCELLED, JobState.SUCCEEDED, JobState.FAILED},
|
|
JobState.FAILED: {JobState.RETRY_QUEUED},
|
|
JobState.RETRY_QUEUED: {JobState.RUNNING, JobState.CANCELLED, JobState.CANCELLING},
|
|
JobState.SUCCEEDED: set(),
|
|
JobState.CANCELLED: set(),
|
|
}
|
|
|
|
ITEM_TRANSITIONS = {
|
|
# running -> queued is an interrupted item being made resumable again.
|
|
ItemState.QUEUED: {ItemState.RUNNING, ItemState.CANCELLED},
|
|
ItemState.RUNNING: {ItemState.SUCCEEDED, ItemState.FAILED, ItemState.CANCELLED, ItemState.QUEUED},
|
|
ItemState.FAILED: {ItemState.QUEUED, ItemState.RUNNING},
|
|
ItemState.SUCCEEDED: set(),
|
|
ItemState.CANCELLED: set(),
|
|
}
|
|
|
|
|
|
def can_transition(current: str, target: str) -> bool:
|
|
return target in ALLOWED_TRANSITIONS.get(current, set())
|
|
|
|
|
|
class JobError(RuntimeError):
|
|
pass
|
|
|
|
|
|
class JobBlocked(JobError):
|
|
"""Enqueue refused by a prerequisite or the one-mutating-job lock policy."""
|
|
|
|
def __init__(self, code: str, detail: str = "") -> None:
|
|
super().__init__(f"{code}: {detail}" if detail else code)
|
|
self.code = code
|
|
self.detail = detail
|
|
|
|
|
|
class InvalidTransition(JobError):
|
|
pass
|
|
|
|
|
|
class JobConflict(JobError):
|
|
"""Optimistic/fencing check failed; another owner or version won the race."""
|
|
|
|
|
|
def _now() -> datetime:
|
|
return datetime.now(timezone.utc)
|
|
|
|
|
|
class JobService:
|
|
def __init__(self, session_factory: sessionmaker) -> None:
|
|
self._session_factory = session_factory
|
|
|
|
# ── enqueue ────────────────────────────────────────────────────────────────
|
|
def enqueue(
|
|
self,
|
|
job_type: str,
|
|
*,
|
|
lock: str | None = None,
|
|
idempotency_key: str | None = None,
|
|
items: Iterable[str] = (),
|
|
) -> dict:
|
|
with self._session_factory() as session:
|
|
if idempotency_key:
|
|
existing = session.scalar(
|
|
select(Job).where(Job.idempotency_key == idempotency_key)
|
|
)
|
|
if existing:
|
|
return self._snapshot(session, existing.id)
|
|
if lock:
|
|
blocker = session.scalar(
|
|
select(Job).where(Job.lock_key == lock, Job.state.in_(ACTIVE_STATES))
|
|
)
|
|
if blocker:
|
|
raise JobBlocked("lock_held", f"{lock} held by job {blocker.id}")
|
|
|
|
job = Job(
|
|
id=str(uuid.uuid4()),
|
|
job_type=job_type,
|
|
state=JobState.QUEUED,
|
|
lock_key=lock,
|
|
idempotency_key=idempotency_key,
|
|
fencing_token=0,
|
|
version=1,
|
|
)
|
|
session.add(job)
|
|
for item_key in items:
|
|
session.add(
|
|
JobItem(job_id=job.id, item_key=item_key, state=ItemState.QUEUED)
|
|
)
|
|
self._event(session, job.id, "queued")
|
|
try:
|
|
session.commit()
|
|
except IntegrityError:
|
|
# Concurrent enqueue with the same idempotency key won the insert.
|
|
session.rollback()
|
|
existing = session.scalar(
|
|
select(Job).where(Job.idempotency_key == idempotency_key)
|
|
)
|
|
if existing:
|
|
return self._snapshot(session, existing.id)
|
|
raise
|
|
return self._snapshot(session, job.id)
|
|
|
|
# ── claim / lease ──────────────────────────────────────────────────────────
|
|
def claim(
|
|
self, job_types: Sequence[str], worker_id: str, *, lease_seconds: int = 60
|
|
) -> dict | None:
|
|
now = _now()
|
|
expires = now + timedelta(seconds=lease_seconds)
|
|
with self._session_factory() as session:
|
|
candidates = (
|
|
session.execute(
|
|
select(Job)
|
|
.where(Job.state.in_(CLAIMABLE_STATES), Job.job_type.in_(list(job_types)))
|
|
.order_by(Job.created_at)
|
|
)
|
|
.scalars()
|
|
.all()
|
|
)
|
|
for job in candidates:
|
|
# Compare-and-set: only the worker whose UPDATE matches the observed
|
|
# state wins; concurrent claimers see rowcount 0 and try the next job.
|
|
result = session.execute(
|
|
update(Job)
|
|
.where(Job.id == job.id, Job.state == job.state)
|
|
.values(
|
|
state=JobState.RUNNING,
|
|
lease_owner=worker_id,
|
|
lease_expires_at=expires,
|
|
heartbeat_at=now,
|
|
started_at=job.started_at or now,
|
|
fencing_token=Job.fencing_token + 1,
|
|
version=Job.version + 1,
|
|
updated_at=now,
|
|
)
|
|
)
|
|
if result.rowcount == 1:
|
|
self._event(session, job.id, "claimed", worker_id)
|
|
session.commit()
|
|
return self._snapshot(session, job.id)
|
|
return None
|
|
|
|
def heartbeat(self, job_id: str, worker_id: str, *, lease_seconds: int = 60) -> bool:
|
|
now = _now()
|
|
with self._session_factory() as session:
|
|
result = session.execute(
|
|
update(Job)
|
|
.where(
|
|
Job.id == job_id,
|
|
Job.lease_owner == worker_id,
|
|
Job.state == JobState.RUNNING,
|
|
)
|
|
.values(heartbeat_at=now, lease_expires_at=now + timedelta(seconds=lease_seconds))
|
|
)
|
|
session.commit()
|
|
return result.rowcount == 1
|
|
|
|
# ── transitions ────────────────────────────────────────────────────────────
|
|
def transition(
|
|
self,
|
|
job_id: str,
|
|
to_state: str,
|
|
*,
|
|
worker_id: str | None = None,
|
|
fencing_token: int | None = None,
|
|
error: tuple[str, str] | None = None,
|
|
) -> dict:
|
|
now = _now()
|
|
with self._session_factory() as session:
|
|
job = session.get(Job, job_id)
|
|
if job is None:
|
|
raise JobError(f"unknown job {job_id}")
|
|
if fencing_token is not None and job.fencing_token != fencing_token:
|
|
raise JobConflict(
|
|
f"stale fencing token {fencing_token} (current {job.fencing_token})"
|
|
)
|
|
if not can_transition(job.state, to_state):
|
|
raise InvalidTransition(f"{job.state} -> {to_state}")
|
|
if worker_id is not None and job.lease_owner not in (None, worker_id):
|
|
raise JobConflict(f"job {job_id} owned by {job.lease_owner}, not {worker_id}")
|
|
job.state = to_state
|
|
job.version += 1
|
|
job.updated_at = now
|
|
if error:
|
|
job.error_code, job.error_message = error
|
|
if to_state in TERMINAL_STATES:
|
|
job.finished_at = now
|
|
job.lease_owner = None
|
|
job.lease_expires_at = None
|
|
self._event(session, job_id, f"state:{to_state}", error[1] if error else None)
|
|
session.commit()
|
|
return self._snapshot(session, job_id)
|
|
|
|
def cancel(self, job_id: str) -> dict:
|
|
with self._session_factory() as session:
|
|
job = session.get(Job, job_id)
|
|
if job is None:
|
|
raise JobError(f"unknown job {job_id}")
|
|
# queued work can be cancelled outright; running work drains via cancelling.
|
|
target = (
|
|
JobState.CANCELLED
|
|
if job.state in (JobState.QUEUED, JobState.RETRY_QUEUED)
|
|
else JobState.CANCELLING
|
|
)
|
|
return self.transition(job_id, target)
|
|
|
|
# ── items / progress ───────────────────────────────────────────────────────
|
|
def set_item(
|
|
self,
|
|
job_id: str,
|
|
item_key: str,
|
|
state: str,
|
|
*,
|
|
error: tuple[str, str] | None = None,
|
|
fencing_token: int | None = None,
|
|
) -> None:
|
|
now = _now()
|
|
with self._session_factory() as session:
|
|
if fencing_token is not None:
|
|
job = session.get(Job, job_id)
|
|
if job is None:
|
|
raise JobError(f"unknown job {job_id}")
|
|
if job.fencing_token != fencing_token:
|
|
raise JobConflict(
|
|
f"stale fencing token {fencing_token} (current {job.fencing_token})"
|
|
)
|
|
item = session.get(JobItem, {"job_id": job_id, "item_key": item_key})
|
|
if item is None:
|
|
raise JobError(f"unknown item {job_id}/{item_key}")
|
|
if state != item.state and state not in ITEM_TRANSITIONS.get(item.state, set()):
|
|
raise InvalidTransition(f"item {item.state} -> {state}")
|
|
if state == ItemState.RUNNING and item.state != ItemState.RUNNING:
|
|
item.attempt_count += 1
|
|
item.state = state
|
|
item.updated_at = now
|
|
if error:
|
|
item.error_code, item.error_message = error
|
|
session.commit()
|
|
|
|
def progress(self, job_id: str) -> dict:
|
|
with self._session_factory() as session:
|
|
states = session.execute(
|
|
select(JobItem.state).where(JobItem.job_id == job_id)
|
|
).scalars().all()
|
|
counts = Counter(states)
|
|
total = len(states)
|
|
done = counts.get(ItemState.SUCCEEDED, 0) + counts.get(ItemState.FAILED, 0) + counts.get(
|
|
ItemState.CANCELLED, 0
|
|
)
|
|
return {"total": total, "done": done, "by_state": dict(counts)}
|
|
|
|
# ── recovery ───────────────────────────────────────────────────────────────
|
|
def recover_stale(self, *, now: datetime | None = None) -> list[str]:
|
|
"""Return running jobs whose lease expired to ``retry_queued`` for safe resume."""
|
|
now = now or _now()
|
|
recovered: list[str] = []
|
|
with self._session_factory() as session:
|
|
stale = (
|
|
session.execute(
|
|
select(Job).where(
|
|
Job.state == JobState.RUNNING, Job.lease_expires_at < now
|
|
)
|
|
)
|
|
.scalars()
|
|
.all()
|
|
)
|
|
for job in stale:
|
|
job.state = JobState.RETRY_QUEUED
|
|
job.lease_owner = None
|
|
job.lease_expires_at = None
|
|
job.version += 1
|
|
job.updated_at = now
|
|
self._event(session, job.id, "interrupted", "lease expired")
|
|
recovered.append(job.id)
|
|
session.commit()
|
|
return recovered
|
|
|
|
def blockers(self, lock: str) -> list[dict]:
|
|
with self._session_factory() as session:
|
|
held = (
|
|
session.execute(
|
|
select(Job).where(Job.lock_key == lock, Job.state.in_(ACTIVE_STATES))
|
|
)
|
|
.scalars()
|
|
.all()
|
|
)
|
|
return [{"id": j.id, "state": j.state, "job_type": j.job_type} for j in held]
|
|
|
|
def get(self, job_id: str) -> dict | None:
|
|
with self._session_factory() as session:
|
|
if session.get(Job, job_id) is None:
|
|
return None
|
|
return self._snapshot(session, job_id)
|
|
|
|
def events(self, job_id: str) -> list[dict]:
|
|
with self._session_factory() as session:
|
|
rows = session.execute(
|
|
select(JobEvent).where(JobEvent.job_id == job_id).order_by(JobEvent.created_at)
|
|
).scalars()
|
|
return [{"type": e.event_type, "message": e.message} for e in rows]
|
|
|
|
def events_after(self, job_id: str, after: int = 0) -> list[dict]:
|
|
"""Durable events with ``seq`` > ``after``, ordered — the resumable cursor for
|
|
SSE and polling. SQLite's monotonic rowid is the sequence (job_events is a
|
|
rowid table), so no extra column is needed."""
|
|
with self._session_factory() as session:
|
|
rows = session.execute(
|
|
text(
|
|
"SELECT rowid AS seq, event_type, message FROM job_events "
|
|
"WHERE job_id = :jid AND rowid > :after ORDER BY rowid"
|
|
),
|
|
{"jid": job_id, "after": after},
|
|
).mappings().all()
|
|
return [{"seq": r["seq"], "type": r["event_type"], "message": r["message"]} for r in rows]
|
|
|
|
# ── helpers ──────────────────────────────────────────────────────────────────
|
|
@staticmethod
|
|
def _event(session, job_id: str, event_type: str, message: str | None = None) -> None:
|
|
session.add(
|
|
JobEvent(id=str(uuid.uuid4()), job_id=job_id, event_type=event_type, message=message)
|
|
)
|
|
|
|
def _snapshot(self, session, job_id: str) -> dict:
|
|
job = session.get(Job, job_id)
|
|
return {
|
|
"id": job.id,
|
|
"job_type": job.job_type,
|
|
"state": job.state,
|
|
"lock_key": job.lock_key,
|
|
"idempotency_key": job.idempotency_key,
|
|
"lease_owner": job.lease_owner,
|
|
"fencing_token": job.fencing_token,
|
|
"version": job.version,
|
|
"error_code": job.error_code,
|
|
}
|