diff --git a/worker/lyra_worker/claim.py b/worker/lyra_worker/claim.py index 0c2f162..10e21f6 100644 --- a/worker/lyra_worker/claim.py +++ b/worker/lyra_worker/claim.py @@ -1,5 +1,28 @@ import psycopg +# States a job passes through while a worker is actively processing it. On startup this +# fresh worker owns none of them, so any job left here was orphaned by a crash/restart. +_IN_FLIGHT_STATES = ("matching", "matched", "downloading", "tagging") + + +def reclaim_stuck_jobs(conn: psycopg.Connection) -> int: + """Reset jobs left mid-pipeline by a prior worker crash/restart back to 'requested' so + they get re-claimed. Safe to call once at startup: nothing is processing them, staging is + cleared beforehand, and the pipeline redoes intake→import idempotently (atomic import). + Returns how many were reclaimed. Terminal jobs (imported/needs_attention) are untouched.""" + with conn.cursor() as cur: + cur.execute( + """ + UPDATE "Job" + SET state = 'requested', "currentStage" = 'intake', "claimedAt" = NULL, + error = NULL, "downloadProgress" = 0, "updatedAt" = now() + WHERE state IN ('matching', 'matched', 'downloading', 'tagging') + """ + ) + n = cur.rowcount + conn.commit() + return n + def claim_next(conn: psycopg.Connection) -> str | None: """Atomically claim the oldest queued job. Returns its id or None.""" diff --git a/worker/lyra_worker/main.py b/worker/lyra_worker/main.py index ca261e1..59a07e9 100644 --- a/worker/lyra_worker/main.py +++ b/worker/lyra_worker/main.py @@ -5,7 +5,7 @@ import uuid import psycopg -from lyra_worker.claim import claim_next +from lyra_worker.claim import claim_next, reclaim_stuck_jobs from lyra_worker.config import get_config from lyra_worker.crypto import secret_key_problem from lyra_worker.db import wait_for_db @@ -207,6 +207,9 @@ def run_forever() -> None: _require_secret_key() conn = wait_for_db() clear_staging_root(STAGING_ROOT) # sweep any staging dirs orphaned by a prior crash + reclaimed = reclaim_stuck_jobs(conn) # requeue jobs a prior crash left mid-pipeline + if reclaimed: + print(f"worker: reclaimed {reclaimed} in-flight job(s) from a prior crash", flush=True) adapters = build_adapters(get_config(conn)) resolver = build_resolver() tagger = build_tagger() diff --git a/worker/tests/test_reclaim.py b/worker/tests/test_reclaim.py new file mode 100644 index 0000000..efa6a00 --- /dev/null +++ b/worker/tests/test_reclaim.py @@ -0,0 +1,55 @@ +from lyra_worker.claim import claim_next, reclaim_stuck_jobs +from tests.conftest import insert_request + + +def _set_state(conn, job_id, state, stage="download", claimed=True, progress=0.5): + with conn.cursor() as cur: + cur.execute( + 'UPDATE "Job" SET state = %s, "currentStage" = %s, ' + '"claimedAt" = ' + ("now()" if claimed else "NULL") + ', ' + '"downloadProgress" = %s, error = %s WHERE id = %s', + (state, stage, progress, "boom" if state == "needs_attention" else None, job_id), + ) + conn.commit() + + +def _job(conn, job_id): + with conn.cursor() as cur: + cur.execute('SELECT state, "currentStage", "claimedAt", "downloadProgress", error ' + 'FROM "Job" WHERE id = %s', (job_id,)) + return cur.fetchone() + + +def test_reclaims_in_flight_states(conn): + ids = {} + for state in ("matching", "matched", "downloading", "tagging"): + jid = insert_request(conn, album=state) + _set_state(conn, jid, state) + ids[state] = jid + + assert reclaim_stuck_jobs(conn) == 4 + + for state, jid in ids.items(): + st, stage, claimed_at, progress, err = _job(conn, jid) + assert st == "requested", state + assert stage == "intake" + assert claimed_at is None + assert progress == 0 + assert err is None + # a reclaimed job is immediately claimable again + assert claim_next(conn) is not None + + +def test_leaves_terminal_and_queued_jobs_untouched(conn): + imported = insert_request(conn, album="imported") + _set_state(conn, imported, "imported", stage="import") + attention = insert_request(conn, album="attention") + _set_state(conn, attention, "needs_attention", stage="download") + queued = insert_request(conn, album="queued") # stays 'requested'/'intake' from insert + + assert reclaim_stuck_jobs(conn) == 0 + + assert _job(conn, imported)[0] == "imported" + assert _job(conn, attention)[0] == "needs_attention" + assert _job(conn, attention)[4] == "boom" # its error is preserved + assert _job(conn, queued)[0] == "requested"