US02-03: Execute Jobs with Leases, Locks, and Recovery (#56)

photo_pipeline/jobs/worker.py

@@ -0,0 +1,138 @@
+"""Durable worker: claims jobs and runs their items safely.
+
+Each worker serves a set of job types (its lanes), claims one job at a time (a
+bounded single lane), and processes items sequentially with heartbeats. It holds
+the fencing token from its claim and stamps every state write with it, so a worker
+that was superseded after a lease expiry cannot commit stale results. Cancellation
+is cooperative (checked between items and offered to handlers), leaving items
+resumable. On resume, items left ``running`` by a dead worker are reset to
+``queued`` and re-run (handlers are idempotent).
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+
+from sqlalchemy import select
+from sqlalchemy.orm import sessionmaker
+
+from photo_pipeline.jobs.handlers import REGISTRY, Cancelled, Handler, JobContext
+from photo_pipeline.models import JobItem
+from photo_pipeline.services.jobs import ItemState, JobConflict, JobService, JobState
+
+
+class Worker:
+    def __init__(
+        self,
+        session_factory: sessionmaker,
+        handlers: Mapping[str, Handler] | None = None,
+        worker_id: str = "worker",
+        *,
+        job_types: Sequence[str] | None = None,
+        lease_seconds: int = 60,
+    ) -> None:
+        self._session_factory = session_factory
+        self.service = JobService(session_factory)
+        self.handlers: Mapping[str, Handler] = handlers if handlers is not None else REGISTRY
+        self.worker_id = worker_id
+        self.job_types = list(job_types if job_types is not None else self.handlers.keys())
+        self.lease_seconds = lease_seconds
+
+    def run_once(self) -> str | None:
+        """Recover stragglers, then claim and fully process one job. Returns its id."""
+        self.service.recover_stale()
+        claimed = self.service.claim(
+            self.job_types, self.worker_id, lease_seconds=self.lease_seconds
+        )
+        if claimed is None:
+            return None
+        self._process(claimed["id"], claimed["job_type"], claimed["fencing_token"])
+        return claimed["id"]
+
+    def _process(self, job_id: str, job_type: str, token: int) -> None:
+        handler = self.handlers[job_type]
+        ctx = JobContext(job_id, self.worker_id, token, self.service)
+        self._reset_interrupted_items(job_id, token)
+
+        cancelled = False
+        any_failed = False
+        for item_key in self._queued_items(job_id):
+            if ctx.cancelled():
+                cancelled = True
+                break
+            try:
+                self.service.set_item(job_id, item_key, ItemState.RUNNING, fencing_token=token)
+                handler(item_key, ctx)
+            except Cancelled:
+                # Cooperative stop: leave the item resumable.
+                self.service.set_item(job_id, item_key, ItemState.QUEUED, fencing_token=token)
+                cancelled = True
+                break
+            except JobConflict:
+                # Superseded mid-item; stop and let the new owner finish.
+                return
+            except Exception as error:  # handler failure for this item
+                any_failed = True
+                self.service.set_item(
+                    job_id,
+                    item_key,
+                    ItemState.FAILED,
+                    error=("handler_error", str(error)[:200]),
+                    fencing_token=token,
+                )
+            else:
+                self.service.set_item(job_id, item_key, ItemState.SUCCEEDED, fencing_token=token)
+            self.service.heartbeat(job_id, self.worker_id, lease_seconds=self.lease_seconds)
+
+        self._finalize(job_id, token, cancelled=cancelled, any_failed=any_failed)
+
+    def _finalize(self, job_id: str, token: int, *, cancelled: bool, any_failed: bool) -> None:
+        snapshot = self.service.get(job_id)
+        if snapshot is None:
+            return
+        try:
+            if cancelled or snapshot["state"] == JobState.CANCELLING:
+                self.service.transition(
+                    job_id, JobState.CANCELLED, worker_id=self.worker_id, fencing_token=token
+                )
+            elif any_failed:
+                self.service.transition(
+                    job_id,
+                    JobState.FAILED,
+                    worker_id=self.worker_id,
+                    fencing_token=token,
+                    error=("items_failed", "one or more items failed"),
+                )
+            else:
+                self.service.transition(
+                    job_id, JobState.SUCCEEDED, worker_id=self.worker_id, fencing_token=token
+                )
+        except JobConflict:
+            # A newer worker owns the job now; do not overwrite its outcome.
+            return
+
+    def _reset_interrupted_items(self, job_id: str, token: int) -> None:
+        for item_key in self._items_in_state(job_id, ItemState.RUNNING):
+            self.service.set_item(job_id, item_key, ItemState.QUEUED, fencing_token=token)
+
+    def _queued_items(self, job_id: str) -> list[str]:
+        return self._items_in_state(job_id, ItemState.QUEUED)
+
+    def _items_in_state(self, job_id: str, state: str) -> list[str]:
+        with self._session_factory() as session:
+            return list(
+                session.execute(
+                    select(JobItem.item_key)
+                    .where(JobItem.job_id == job_id, JobItem.state == state)
+                    .order_by(JobItem.item_key)
+                ).scalars()
+            )
+
+    def run_forever(self, *, idle_sleep: float = 1.0, iterations: int | None = None) -> None:
+        import time
+
+        count = 0
+        while iterations is None or count < iterations:
+            if self.run_once() is None:
+                time.sleep(idle_sleep)
+            count += 1