157 lines
6.7 KiB
Python
157 lines
6.7 KiB
Python
"""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.faults import JOB_ITEM_DONE, maybe_fault
|
|
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,
|
|
config: object | None = None,
|
|
) -> 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
|
|
self.config = config
|
|
|
|
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.config
|
|
)
|
|
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)
|
|
# One item is durably done, the job outcome is not: the control point
|
|
# for a crash mid-batch (US07-04). Recovery must not re-run this item.
|
|
maybe_fault(JOB_ITEM_DONE)
|
|
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:
|
|
# A handler may stop on its own — an upload batch cancelled through
|
|
# its own API never touches the job — so the job can still be
|
|
# ``running`` here. Record the request before the outcome: a stop is
|
|
# always observable as cancelling → cancelled, and ``running ->
|
|
# cancelled`` is not a legal jump. Without this hop the transition
|
|
# is rejected and the job keeps its lock forever.
|
|
if snapshot["state"] == JobState.RUNNING:
|
|
self.service.transition(
|
|
job_id, JobState.CANCELLING, worker_id=self.worker_id, fencing_token=token
|
|
)
|
|
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
|