From 8bbc1db32f3401c2c2b14b717864dd1cb87bfa5d Mon Sep 17 00:00:00 2001 From: Jonathan Date: Sun, 12 Jul 2026 12:37:06 +0200 Subject: [PATCH] feat(worker): chunk the library scan + reconnect on DB loss Two worker-loop robustness fixes: - Scan chunking: maybe_run_scan advances the library scan by at most one scan_chunk (default scan.chunkSize=25 albums) per loop iteration, so a large library no longer blocks job-claim/monitor for minutes-to-hours. Progress persists in scan.inProgress/scan.cursor/scan.progress and is resumable; scan.result (unchanged format) is written on completion. - DB reconnect: wrap the loop body in `except psycopg.OperationalError` (AdminShutdown subclasses it) to close the dead handle and reconnect via wait_for_db() in-process, instead of crashing and relying on the container restart policy. Verified via a backend-termination test. Co-Authored-By: Claude Opus 4.8 (1M context) --- worker/lyra_worker/main.py | 149 +++++++++++++++++++++--------- worker/tests/test_db_reconnect.py | 35 +++++++ worker/tests/test_scan_trigger.py | 54 +++++++++-- 3 files changed, 186 insertions(+), 52 deletions(-) create mode 100644 worker/tests/test_db_reconnect.py diff --git a/worker/lyra_worker/main.py b/worker/lyra_worker/main.py index fc228b4..64b45c8 100644 --- a/worker/lyra_worker/main.py +++ b/worker/lyra_worker/main.py @@ -1,5 +1,7 @@ import time +import psycopg + from lyra_worker.claim import claim_next from lyra_worker.config import get_config from lyra_worker.db import wait_for_db @@ -11,12 +13,13 @@ from lyra_worker.registry import ( build_adapters, build_browser, build_probe, build_resolver, build_similarity_sources, build_tagger, ) -from lyra_worker.scan import scan_library +from lyra_worker.scan import scan_chunk IDLE_SLEEP = 2.0 MONITOR_TICK_SECONDS = 60.0 DISCOVER_TICK_SECONDS = 300.0 # check the discovery interval every 5 min (interval itself is hours) DEST_ROOT = "/music" # must match run_pipeline's default library root +DEFAULT_SCAN_CHUNK = 25 # albums resolved per loop iteration (~25-50s at MB ~1 req/s) _TRUE = {"1", "true", "yes", "on"} @@ -30,11 +33,61 @@ def _set_config(conn, key: str, value: str) -> None: conn.commit() -def _run_scan(conn, resolver, probe, browser, dest_root: str = DEST_ROOT) -> None: - result = scan_library(conn, resolver, probe, browser, dest_root) - _set_config(conn, "scan.result", f"imported {result.imported}, skipped {result.skipped}") - _set_config(conn, "scan.requested", "false") - print(f"worker: library scan done — {result.imported} imported, {result.skipped} skipped", flush=True) +def _scan_chunk_size(config) -> int: + try: + return max(1, int(config.get("scan.chunkSize", ""))) + except (TypeError, ValueError): + return DEFAULT_SCAN_CHUNK + + +def _parse_progress(s: str) -> tuple[int, int]: + """Parse the internal 'imported/skipped' running-total marker; ('', junk) -> (0, 0).""" + try: + i, _, k = (s or "").partition("/") + return int(i), int(k) + except (ValueError, AttributeError): + return 0, 0 + + +def maybe_run_scan(conn, resolver, probe, browser, config, dest_root: str = DEST_ROOT) -> None: + """Advance a chunked library scan by at most one chunk per call, so the worker loop + keeps claiming jobs and ticking the monitor between chunks. `scan.requested` starts a + fresh scan; `scan.inProgress`/`scan.cursor`/`scan.progress` persist state until the walk + is exhausted, at which point `scan.result` is written (matching the old one-shot output).""" + requested = str(config.get("scan.requested", "")).strip().lower() in _TRUE + in_progress = str(config.get("scan.inProgress", "")).strip().lower() in _TRUE + if not requested and not in_progress: + return + if requested and not in_progress: # start a fresh scan + _set_config(conn, "scan.inProgress", "true") + _set_config(conn, "scan.requested", "false") + _set_config(conn, "scan.cursor", "") + _set_config(conn, "scan.progress", "0/0") + cursor, imported_so_far, skipped_so_far = "", 0, 0 + else: # resume the in-progress scan from its persisted cursor + cursor = config.get("scan.cursor", "") or "" + imported_so_far, skipped_so_far = _parse_progress(config.get("scan.progress", "")) + try: + new_cursor, imp, skp, done = scan_chunk( + conn, resolver, probe, browser, dest_root, cursor, _scan_chunk_size(config) + ) + except Exception as e: # a scan error must never kill the worker + print(f"worker: library scan chunk failed: {e}", flush=True) + conn.rollback() + _set_config(conn, "scan.inProgress", "false") # abandon the scan; a new request restarts it + _set_config(conn, "scan.requested", "false") + return + imported_total = imported_so_far + imp + skipped_total = skipped_so_far + skp + if done: + _set_config(conn, "scan.result", f"imported {imported_total}, skipped {skipped_total}") + _set_config(conn, "scan.progress", "") + _set_config(conn, "scan.cursor", "") + _set_config(conn, "scan.inProgress", "false") + print(f"worker: library scan done — {imported_total} imported, {skipped_total} skipped", flush=True) + else: + _set_config(conn, "scan.cursor", new_cursor) + _set_config(conn, "scan.progress", f"{imported_total}/{skipped_total}") def _run_discovery(conn, sources, browser, config=None) -> None: @@ -80,48 +133,54 @@ def run_forever() -> None: last_discover_tick = 0.0 try: while True: - config = get_config(conn) - mcfg = MonitorConfig.from_config(config) - now = time.monotonic() - if mcfg.enabled and now - last_tick >= MONITOR_TICK_SECONDS: - try: - sweep(conn, browser, mcfg) - except Exception as e: # a monitor error must never kill the worker - print(f"worker: monitor sweep failed: {e}", flush=True) - conn.rollback() - last_tick = now - - dcfg = DiscoveryConfig.from_config(config) - last_discover_tick = maybe_run_discovery( - conn, similarity_sources, browser, config, dcfg, now, last_discover_tick - ) - - if str(config.get("scan.requested", "")).strip().lower() in _TRUE: - try: - _run_scan(conn, resolver, probe, browser) - except Exception as e: # a scan error must never kill the worker - print(f"worker: library scan failed: {e}", flush=True) - conn.rollback() - _set_config(conn, "scan.requested", "false") - - job_id = claim_next(conn) - if job_id is None: - time.sleep(IDLE_SLEEP) - continue - print(f"worker: claimed job {job_id}", flush=True) + # A transient DB connection loss (e.g. AdminShutdown when Postgres restarts) raises + # OperationalError from any DB op below; reconnect in-process rather than crashing and + # relying on the container restart policy. Inner rollbacks on a dead connection also + # raise OperationalError and compose up to here. try: - run_pipeline(conn, job_id, adapters, resolver=resolver, tagger=tagger, - upgrade_cutoff=mcfg.quality_cutoff) - except Exception as e: # one bad job must not take down the worker loop - print(f"worker: pipeline failed for job {job_id}: {e}", flush=True) - conn.rollback() - if mcfg.enabled: + config = get_config(conn) + mcfg = MonitorConfig.from_config(config) + now = time.monotonic() + if mcfg.enabled and now - last_tick >= MONITOR_TICK_SECONDS: + try: + sweep(conn, browser, mcfg) + except Exception as e: # a monitor error must never kill the worker + print(f"worker: monitor sweep failed: {e}", flush=True) + conn.rollback() + last_tick = now + + dcfg = DiscoveryConfig.from_config(config) + last_discover_tick = maybe_run_discovery( + conn, similarity_sources, browser, config, dcfg, now, last_discover_tick + ) + + maybe_run_scan(conn, resolver, probe, browser, config) + + job_id = claim_next(conn) + if job_id is None: + time.sleep(IDLE_SLEEP) + continue + print(f"worker: claimed job {job_id}", flush=True) try: - reconcile(conn, job_id, mcfg) - except Exception as e: - print(f"worker: reconcile failed for job {job_id}: {e}", flush=True) + run_pipeline(conn, job_id, adapters, resolver=resolver, tagger=tagger, + upgrade_cutoff=mcfg.quality_cutoff) + except Exception as e: # one bad job must not take down the worker loop + print(f"worker: pipeline failed for job {job_id}: {e}", flush=True) conn.rollback() - print(f"worker: finished job {job_id}", flush=True) + if mcfg.enabled: + try: + reconcile(conn, job_id, mcfg) + except Exception as e: + print(f"worker: reconcile failed for job {job_id}: {e}", flush=True) + conn.rollback() + print(f"worker: finished job {job_id}", flush=True) + except psycopg.OperationalError as e: + print(f"worker: db connection lost ({e}); reconnecting...", flush=True) + try: + conn.close() + except Exception: + pass + conn = wait_for_db() finally: conn.close() diff --git a/worker/tests/test_db_reconnect.py b/worker/tests/test_db_reconnect.py new file mode 100644 index 0000000..4bfc96e --- /dev/null +++ b/worker/tests/test_db_reconnect.py @@ -0,0 +1,35 @@ +import psycopg +import pytest + +from lyra_worker.db import connect, wait_for_db + + +def test_terminated_backend_raises_operational_error_then_reconnects(conn): + """The run_forever reconnect boundary catches psycopg.OperationalError and re-runs + wait_for_db(). This exercises both primitives: a server-side backend termination (the + real AdminShutdown case) surfaces as OperationalError, and wait_for_db reconnects.""" + victim = connect() + try: + with victim.cursor() as cur: + cur.execute("SELECT pg_backend_pid()") + pid = cur.fetchone()[0] + with conn.cursor() as cur: # terminate the victim's backend, as a DB restart would + cur.execute("SELECT pg_terminate_backend(%s)", (pid,)) + conn.commit() + with pytest.raises(psycopg.OperationalError): # the exact type the loop catches + with victim.cursor() as cur: + cur.execute("SELECT 1") + cur.fetchone() + finally: + try: + victim.close() + except Exception: + pass + + fresh = wait_for_db() # reconnect exactly as run_forever does + try: + with fresh.cursor() as cur: + cur.execute("SELECT 1") + assert cur.fetchone()[0] == 1 + finally: + fresh.close() diff --git a/worker/tests/test_scan_trigger.py b/worker/tests/test_scan_trigger.py index 88f9639..e13361c 100644 --- a/worker/tests/test_scan_trigger.py +++ b/worker/tests/test_scan_trigger.py @@ -1,5 +1,5 @@ from lyra_worker.adapters.fakes import FakeMbBrowser -from lyra_worker.main import _run_scan +from lyra_worker.main import maybe_run_scan from lyra_worker.types import MBTarget @@ -11,7 +11,7 @@ class FakeProbe: class FakeResolver: def resolve(self, artist, album): - return MBTarget(artist=artist, album=album, rg_mbid="rg1", artist_mbid="a1") + return MBTarget(artist=artist, album=album, rg_mbid=f"rg-{album}", artist_mbid="a1") def _set(conn, key, value): @@ -31,15 +31,55 @@ def _get(conn, key): return row[0] if row else None -def test_run_scan_writes_result_and_clears_flag(conn, tmp_path): - (tmp_path / "John Mayer" / "Continuum (2006)").mkdir(parents=True) - (tmp_path / "John Mayer" / "Continuum (2006)" / "01.flac").write_bytes(b"") +def _album(tmp_path, artist, folder): + d = tmp_path / artist / folder + d.mkdir(parents=True) + (d / "01.flac").write_bytes(b"") + + +def _cfg(conn): + from lyra_worker.config import get_config + return get_config(conn) + + +def test_scan_completes_in_one_chunk_writes_result_and_clears_flag(conn, tmp_path): + _album(tmp_path, "John Mayer", "Continuum (2006)") _set(conn, "scan.requested", "true") - _run_scan(conn, FakeResolver(), FakeProbe(), FakeMbBrowser(), dest_root=str(tmp_path)) + # one album, default chunk size 25 -> starts + finishes in a single call + maybe_run_scan(conn, FakeResolver(), FakeProbe(), FakeMbBrowser(), _cfg(conn), dest_root=str(tmp_path)) - assert _get(conn, "scan.requested") == "false" # flag cleared + assert _get(conn, "scan.requested") == "false" # flag cleared + assert _get(conn, "scan.inProgress") == "false" # scan finished assert "imported" in (_get(conn, "scan.result") or "") # summary written with conn.cursor() as cur: cur.execute('SELECT count(*) FROM "LibraryItem"') assert cur.fetchone()[0] == 1 + + +def test_scan_chunks_across_iterations_and_totals_correctly(conn, tmp_path): + for i in range(3): + _album(tmp_path, "Artist", f"Album {i}") + _set(conn, "scan.requested", "true") + _set(conn, "scan.chunkSize", "2") + + # first call starts the scan and processes chunk 1 (2 albums), leaving it in progress + maybe_run_scan(conn, FakeResolver(), FakeProbe(), FakeMbBrowser(), _cfg(conn), dest_root=str(tmp_path)) + assert _get(conn, "scan.inProgress") == "true" + assert _get(conn, "scan.requested") == "false" + assert _get(conn, "scan.progress") == "2/0" + + # second call resumes from the cursor, finishes the last album + maybe_run_scan(conn, FakeResolver(), FakeProbe(), FakeMbBrowser(), _cfg(conn), dest_root=str(tmp_path)) + assert _get(conn, "scan.inProgress") == "false" + assert _get(conn, "scan.result") == "imported 3, skipped 0" + with conn.cursor() as cur: + cur.execute('SELECT count(*) FROM "LibraryItem"') + assert cur.fetchone()[0] == 3 + + +def test_maybe_run_scan_idle_when_not_requested_or_in_progress(conn, tmp_path): + # no flags set -> no-op, no result written + maybe_run_scan(conn, FakeResolver(), FakeProbe(), FakeMbBrowser(), _cfg(conn), dest_root=str(tmp_path)) + assert _get(conn, "scan.result") is None + assert _get(conn, "scan.inProgress") is None