photoanalyzer/photo_pipeline/services/jobs.py

"""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, 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 = {
    ItemState.QUEUED: {ItemState.RUNNING, ItemState.CANCELLED},
    ItemState.RUNNING: {ItemState.SUCCEEDED, ItemState.FAILED, ItemState.CANCELLED},
    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,
        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 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
    ) -> None:
        now = _now()
        with self._session_factory() as session:
            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]

    # ── 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,
        }