From f0f163158dfc90cc063cdc3fd085819e1e3e8431 Mon Sep 17 00:00:00 2001 From: domverse Date: Wed, 15 Jul 2026 23:12:51 +0200 Subject: [PATCH] US02-03: Execute Jobs with Leases, Locks, and Recovery --- photo_pipeline/__main__.py | 11 +++ photo_pipeline/jobs/__init__.py | 1 + photo_pipeline/jobs/handlers.py | 54 +++++++++++ photo_pipeline/jobs/locks.py | 44 +++++++++ photo_pipeline/jobs/worker.py | 138 +++++++++++++++++++++++++++ photo_pipeline/services/jobs.py | 24 ++++- tests/e2e/test_worker_kill.py | 84 +++++++++++++++++ tests/integration/test_worker.py | 156 +++++++++++++++++++++++++++++++ tests/story_traceability.json | 5 + tests/unit/test_lock_order.py | 28 ++++++ 10 files changed, 543 insertions(+), 2 deletions(-) create mode 100644 photo_pipeline/jobs/__init__.py create mode 100644 photo_pipeline/jobs/handlers.py create mode 100644 photo_pipeline/jobs/locks.py create mode 100644 photo_pipeline/jobs/worker.py create mode 100644 tests/e2e/test_worker_kill.py create mode 100644 tests/integration/test_worker.py create mode 100644 tests/unit/test_lock_order.py diff --git a/photo_pipeline/__main__.py b/photo_pipeline/__main__.py index 4486d6e..06d4a41 100644 --- a/photo_pipeline/__main__.py +++ b/photo_pipeline/__main__.py @@ -14,6 +14,8 @@ def main(argv: Sequence[str] | None = None) -> int: commands = parser.add_subparsers(dest="command", required=True) commands.add_parser("serve", help="Run the API server") commands.add_parser("migrate", help="Upgrade the database to the latest revision") + worker_cmd = commands.add_parser("worker", help="Run a durable-job worker") + worker_cmd.add_argument("--id", default="worker-1", help="Worker id (lease owner)") args = parser.parse_args(argv) config = Config.from_env() @@ -23,6 +25,15 @@ def main(argv: Sequence[str] | None = None) -> int: run_migrations(config.database_url) return 0 + if args.command == "worker": + from photo_pipeline.db import create_db_engine, create_session_factory + from photo_pipeline.jobs.worker import Worker + + run_migrations(config.database_url) + engine = create_db_engine(config.database_url) + Worker(create_session_factory(engine), worker_id=args.id).run_forever() + return 0 + import uvicorn from photo_pipeline.api.app import create_app diff --git a/photo_pipeline/jobs/__init__.py b/photo_pipeline/jobs/__init__.py new file mode 100644 index 0000000..15c3f08 --- /dev/null +++ b/photo_pipeline/jobs/__init__.py @@ -0,0 +1 @@ +"""Durable job execution: worker loop, handler dispatch, and lock ordering.""" diff --git a/photo_pipeline/jobs/handlers.py b/photo_pipeline/jobs/handlers.py new file mode 100644 index 0000000..57f9509 --- /dev/null +++ b/photo_pipeline/jobs/handlers.py @@ -0,0 +1,54 @@ +"""Job-type handler dispatch. + +A handler processes one item of a job: ``handler(item_key, ctx)`` and either +returns (success), raises ``Cancelled`` (cooperative stop, resumable), or raises +any other exception (item failure). Handlers must be idempotent — at-least-once +delivery means the same item can run again after an interrupted attempt. + +Domain handlers (safety scoring, analysis, …) register here in later stories; the +registry ships empty so the worker engine can be built and tested independently. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Callable + +if TYPE_CHECKING: + from photo_pipeline.services.jobs import JobService + +Handler = Callable[[str, "JobContext"], None] + + +class Cancelled(Exception): + """Raised by a handler that observed cancellation and stopped cleanly.""" + + +@dataclass +class JobContext: + """What a handler needs from the coordinator: cancellation checks and heartbeats.""" + + job_id: str + worker_id: str + fencing_token: int + service: "JobService" + + def cancelled(self) -> bool: + from photo_pipeline.services.jobs import JobState + + snapshot = self.service.get(self.job_id) + return snapshot is None or snapshot["state"] in ( + JobState.CANCELLING, + JobState.CANCELLED, + ) + + def heartbeat(self, *, lease_seconds: int = 60) -> bool: + return self.service.heartbeat(self.job_id, self.worker_id, lease_seconds=lease_seconds) + + +# Populated by domain stories via register(); empty for now. +REGISTRY: dict[str, Handler] = {} + + +def register(job_type: str, handler: Handler) -> None: + REGISTRY[job_type] = handler diff --git a/photo_pipeline/jobs/locks.py b/photo_pipeline/jobs/locks.py new file mode 100644 index 0000000..b49acf6 --- /dev/null +++ b/photo_pipeline/jobs/locks.py @@ -0,0 +1,44 @@ +"""Lock hierarchy and acquisition ordering. + +Locks are acquired broad → narrow to prevent deadlocks: never take a broader lock +while holding a narrower one (concept §16). The coordinator already enforces +one active job per lock key; this guards multi-lock operations. + + library lease → stage/job lease → album/folder lease → asset lease +""" + +from __future__ import annotations + +from collections.abc import Iterable + +# Lower rank = broader scope. Mutating lanes map onto these tiers. +LOCK_RANK = { + "library": 0, + "library_write": 0, + "rename": 1, + "upload": 1, + "archive": 1, + "album": 2, + "asset": 3, + "exif": 3, +} + + +class LockOrderError(RuntimeError): + pass + + +def rank(lock: str) -> int: + if lock not in LOCK_RANK: + raise LockOrderError(f"unknown lock {lock!r}") + return LOCK_RANK[lock] + + +def validate_acquisition(held: Iterable[str], acquiring: str) -> None: + """Reject acquiring a broader lock than one already held (deadlock risk).""" + new_rank = rank(acquiring) + for lock in held: + if rank(lock) > new_rank: + raise LockOrderError( + f"cannot acquire broader lock {acquiring!r} while holding narrower {lock!r}" + ) diff --git a/photo_pipeline/jobs/worker.py b/photo_pipeline/jobs/worker.py new file mode 100644 index 0000000..f7b660c --- /dev/null +++ b/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 diff --git a/photo_pipeline/services/jobs.py b/photo_pipeline/services/jobs.py index f750ba9..fbd552d 100644 --- a/photo_pipeline/services/jobs.py +++ b/photo_pipeline/services/jobs.py @@ -66,8 +66,9 @@ ALLOWED_TRANSITIONS = { } 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.RUNNING: {ItemState.SUCCEEDED, ItemState.FAILED, ItemState.CANCELLED, ItemState.QUEUED}, ItemState.FAILED: {ItemState.QUEUED, ItemState.RUNNING}, ItemState.SUCCEEDED: set(), ItemState.CANCELLED: set(), @@ -219,6 +220,7 @@ class JobService: to_state: str, *, worker_id: str | None = None, + fencing_token: int | None = None, error: tuple[str, str] | None = None, ) -> dict: now = _now() @@ -226,6 +228,10 @@ class JobService: 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): @@ -258,10 +264,24 @@ class JobService: # ── items / progress ─────────────────────────────────────────────────────── def set_item( - self, job_id: str, item_key: str, state: str, *, error: tuple[str, str] | None = None + 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}") diff --git a/tests/e2e/test_worker_kill.py b/tests/e2e/test_worker_kill.py new file mode 100644 index 0000000..00ff9a7 --- /dev/null +++ b/tests/e2e/test_worker_kill.py @@ -0,0 +1,84 @@ +"""Process test: kill a worker mid-item, then prove the job recovers and a fresh +worker finishes it. The killed worker's lease expires and its ownership is fenced.""" + +import subprocess +import sys +import time +from pathlib import Path + +import pytest + +from photo_pipeline.db import create_db_engine, create_session_factory, run_migrations +from photo_pipeline.jobs.worker import Worker +from photo_pipeline.services.jobs import ItemState, JobService, JobState + +REPO = Path(__file__).resolve().parents[2] + +# A worker that claims the job, marks its item running, signals it started, then +# blocks — so the test can kill it mid-item. +WORKER_SCRIPT = """ +import sys, time +from pathlib import Path +sys.path.insert(0, {repo!r}) +from photo_pipeline.db import create_db_engine, create_session_factory +from photo_pipeline.jobs.worker import Worker + +db_url, started = sys.argv[1], Path(sys.argv[2]) + +def blocking_handler(item_key, ctx): + started.write_text("started") + time.sleep(30) + +sf = create_session_factory(create_db_engine(db_url)) +Worker(sf, {{"slow": blocking_handler}}, "killable", lease_seconds=1).run_once() +""" + + +@pytest.fixture +def db_url(tmp_path): + url = f"sqlite:///{tmp_path / 'kill.db'}" + run_migrations(url) + return url + + +def test_killed_worker_job_recovers_and_completes(tmp_path, db_url): + engine = create_db_engine(db_url) + service = JobService(create_session_factory(engine)) + job = service.enqueue("slow", items=["only"]) + + script = tmp_path / "worker_script.py" + script.write_text(WORKER_SCRIPT.format(repo=str(REPO))) + started = tmp_path / "started.flag" + + proc = subprocess.Popen( + [sys.executable, str(script), db_url, str(started)], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + try: + deadline = time.monotonic() + 20 + while time.monotonic() < deadline and not started.exists(): + if proc.poll() is not None: + _, err = proc.communicate() + pytest.fail(f"worker exited early: {err.decode(errors='replace')}") + time.sleep(0.1) + assert started.exists(), "worker never began the item" + assert service.get(job["id"])["state"] == JobState.RUNNING + finally: + proc.kill() + proc.wait(timeout=10) + + # Let the 1s lease expire, then a fresh worker recovers and finishes the job. + time.sleep(1.3) + completed = [] + fresh = Worker( + create_session_factory(engine), + {"slow": lambda item, ctx: completed.append(item)}, + "recovery-worker", + ) + fresh.run_once() + + assert service.get(job["id"])["state"] == JobState.SUCCEEDED + assert completed == ["only"] # the interrupted item was reprocessed + assert service.progress(job["id"])["by_state"] == {ItemState.SUCCEEDED: 1} + engine.dispose() diff --git a/tests/integration/test_worker.py b/tests/integration/test_worker.py new file mode 100644 index 0000000..f3e160c --- /dev/null +++ b/tests/integration/test_worker.py @@ -0,0 +1,156 @@ +"""Worker execution: success, failure, cooperative cancel, resume, fencing, and a +randomized concurrency stress with no duplicate completion or lost events.""" + +import threading +from collections import Counter +from concurrent.futures import ThreadPoolExecutor +from datetime import timedelta + +import pytest + +from photo_pipeline.db import create_db_engine, create_session_factory, run_migrations +from photo_pipeline.jobs.worker import Worker +from photo_pipeline.services.jobs import ( + ItemState, + JobConflict, + JobService, + JobState, + _now, +) + + +@pytest.fixture +def db_url(tmp_path): + url = f"sqlite:///{tmp_path / 'worker.db'}" + run_migrations(url) + return url + + +@pytest.fixture +def make_factory(): + engines = [] + + def _make(url): + engine = create_db_engine(url) + engines.append(engine) + return create_session_factory(engine) + + yield _make + for engine in engines: + engine.dispose() + + +@pytest.fixture +def sf(db_url, make_factory): + return make_factory(db_url) + + +@pytest.fixture +def jobs(sf): + return JobService(sf) + + +def test_worker_processes_all_items(sf, jobs): + handled = [] + worker = Worker(sf, {"scan": lambda item, ctx: handled.append(item)}, "w1") + job = jobs.enqueue("scan", items=["a", "b", "c"]) + assert worker.run_once() == job["id"] + assert jobs.get(job["id"])["state"] == JobState.SUCCEEDED + assert sorted(handled) == ["a", "b", "c"] + assert jobs.progress(job["id"])["done"] == 3 + + +def test_handler_failure_fails_the_job(sf, jobs): + def handler(item, ctx): + if item == "bad": + raise ValueError("nope") + + worker = Worker(sf, {"scan": handler}, "w1") + job = jobs.enqueue("scan", items=["ok", "bad"]) + worker.run_once() + assert jobs.get(job["id"])["state"] == JobState.FAILED + assert jobs.progress(job["id"])["by_state"] == { + ItemState.SUCCEEDED: 1, + ItemState.FAILED: 1, + } + + +def test_cooperative_cancellation_leaves_items_resumable(sf, jobs): + # Cancel arrives while the first item runs; the worker stops before the next. + def handler(item, ctx): + if item == "a": + ctx.service.cancel(ctx.job_id) + + worker = Worker(sf, {"scan": handler}, "w1") + job = jobs.enqueue("scan", items=["a", "b"]) + worker.run_once() + assert jobs.get(job["id"])["state"] == JobState.CANCELLED + by_state = jobs.progress(job["id"])["by_state"] + assert by_state.get(ItemState.SUCCEEDED) == 1 # "a" finished + assert by_state.get(ItemState.QUEUED) == 1 # "b" left resumable + + +def test_fencing_rejects_superseded_worker(sf, jobs): + job = jobs.enqueue("scan", items=["a"]) + stale = jobs.claim(["scan"], "old") + jobs.recover_stale(now=_now() + timedelta(hours=1)) + fresh = jobs.claim(["scan"], "new") + assert fresh["fencing_token"] > stale["fencing_token"] + with pytest.raises(JobConflict): + jobs.set_item(job["id"], "a", ItemState.SUCCEEDED, fencing_token=stale["fencing_token"]) + with pytest.raises(JobConflict): + jobs.transition( + job["id"], JobState.SUCCEEDED, fencing_token=stale["fencing_token"] + ) + + +def test_recovered_job_resumes_remaining_items(sf, jobs): + job = jobs.enqueue("scan", items=["a", "b"]) + # Simulate a crashed attempt: "a" done, "b" left running, no finalize. + crashed = jobs.claim(["scan"], "dead") + token = crashed["fencing_token"] + jobs.set_item(job["id"], "a", ItemState.RUNNING, fencing_token=token) + jobs.set_item(job["id"], "a", ItemState.SUCCEEDED, fencing_token=token) + jobs.set_item(job["id"], "b", ItemState.RUNNING, fencing_token=token) + jobs.recover_stale(now=_now() + timedelta(hours=1)) + + handled = [] + fresh_worker = Worker(sf, {"scan": lambda item, ctx: handled.append(item)}, "alive") + fresh_worker.run_once() + assert jobs.get(job["id"])["state"] == JobState.SUCCEEDED + assert handled == ["b"] # "a" already succeeded, only "b" reprocessed + assert jobs.progress(job["id"])["by_state"][ItemState.SUCCEEDED] == 2 + + +def test_concurrent_workers_no_duplicate_completion_or_lost_events(db_url, make_factory): + service = JobService(make_factory(db_url)) + job_ids = [service.enqueue("scan", items=["x", "y"])["id"] for _ in range(12)] + + calls = Counter() + lock = threading.Lock() + + def handler(item, ctx): + with lock: + calls[(ctx.job_id, item)] += 1 + + def drain(worker_index): + worker = Worker(make_factory(db_url), {"scan": handler}, f"w{worker_index}") + empties = 0 + for _ in range(200): + if worker.run_once() is None: + empties += 1 + if empties >= 2: + break + else: + empties = 0 + + with ThreadPoolExecutor(max_workers=4) as pool: + list(pool.map(drain, range(4))) + + # Every job completed exactly once; every item handled exactly once. + for job_id in job_ids: + assert service.get(job_id)["state"] == JobState.SUCCEEDED + types = [e["type"] for e in service.events(job_id)] + assert "claimed" in types and "state:succeeded" in types # no lost events + assert all(count == 1 for count in calls.values()) + assert len(calls) == len(job_ids) * 2 diff --git a/tests/story_traceability.json b/tests/story_traceability.json index c75c468..68b879f 100644 --- a/tests/story_traceability.json +++ b/tests/story_traceability.json @@ -48,6 +48,11 @@ "US02-02": [ "tests/unit/test_job_state_machine.py", "tests/integration/test_jobs.py" + ], + "US02-03": [ + "tests/unit/test_lock_order.py", + "tests/integration/test_worker.py", + "tests/e2e/test_worker_kill.py" ] } } diff --git a/tests/unit/test_lock_order.py b/tests/unit/test_lock_order.py new file mode 100644 index 0000000..4cd23da --- /dev/null +++ b/tests/unit/test_lock_order.py @@ -0,0 +1,28 @@ +"""Lock acquisition ordering (deadlock prevention).""" + +import pytest + +from photo_pipeline.jobs import locks + + +def test_broad_to_narrow_ranks(): + assert locks.rank("library") < locks.rank("rename") + assert locks.rank("rename") < locks.rank("album") + assert locks.rank("album") < locks.rank("asset") + + +def test_acquiring_narrower_while_holding_broader_is_allowed(): + locks.validate_acquisition(["library"], "asset") # no raise + locks.validate_acquisition(["library", "album"], "asset") + + +def test_acquiring_broader_while_holding_narrower_is_rejected(): + with pytest.raises(locks.LockOrderError): + locks.validate_acquisition(["asset"], "library") + with pytest.raises(locks.LockOrderError): + locks.validate_acquisition(["album"], "rename") + + +def test_unknown_lock_is_rejected(): + with pytest.raises(locks.LockOrderError): + locks.rank("nonsense")