photoanalyzer/photo_pipeline/jobs/worker.py

"""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._session_factory)
        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