diff --git a/worker/lyra_worker/main.py b/worker/lyra_worker/main.py index af3fd62..69ed3f1 100644 --- a/worker/lyra_worker/main.py +++ b/worker/lyra_worker/main.py @@ -1,5 +1,6 @@ import os import time +import uuid import psycopg @@ -14,7 +15,7 @@ from lyra_worker.registry import ( build_adapters, build_browser, build_probe, build_resolver, build_similarity_sources, build_tagger, ) -from lyra_worker.scan import scan_chunk +from lyra_worker.scan import build_worklist, clear_worklist, scan_chunk IDLE_SLEEP = 2.0 HEARTBEAT_SECONDS = 15.0 # how often to stamp worker.heartbeat (its updatedAt = liveness) @@ -58,29 +59,39 @@ def _parse_progress(s: str) -> tuple[int, int]: 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).""" + fresh scan (walk once -> ScanWorkItem worklist); `scan.inProgress`/`scan.id`/`scan.progress` + persist state until the worklist drains, at which point `scan.result` is written.""" 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 + scan_id = uuid.uuid4().hex _set_config(conn, "scan.inProgress", "true") _set_config(conn, "scan.requested", "false") - _set_config(conn, "scan.cursor", "") + _set_config(conn, "scan.id", scan_id) _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 "" + try: + build_worklist(conn, scan_id, dest_root) + except Exception as e: # a scan error must never kill the worker + print(f"worker: library scan worklist build failed: {e}", flush=True) + conn.rollback() + _set_config(conn, "scan.inProgress", "false") + _set_config(conn, "scan.requested", "false") + return + imported_so_far, skipped_so_far = 0, 0 + else: # resume the in-progress scan + scan_id = config.get("scan.id", "") or "" imported_so_far, skipped_so_far = _parse_progress(config.get("scan.progress", "")) + if not scan_id: # corrupt/absent state -> abandon; a new request restarts it + _set_config(conn, "scan.inProgress", "false") + return try: - new_cursor, imp, skp, done = scan_chunk( - conn, resolver, probe, browser, dest_root, cursor, _scan_chunk_size(config) - ) + imp, skp, done = scan_chunk(conn, resolver, probe, browser, scan_id, _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.inProgress", "false") _set_config(conn, "scan.requested", "false") return imported_total = imported_so_far + imp @@ -88,11 +99,11 @@ def maybe_run_scan(conn, resolver, probe, browser, config, dest_root: str = DEST 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") + clear_worklist(conn, scan_id) + _set_config(conn, "scan.id", "") 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}") diff --git a/worker/lyra_worker/scan.py b/worker/lyra_worker/scan.py index ceeba75..88b8c70 100644 --- a/worker/lyra_worker/scan.py +++ b/worker/lyra_worker/scan.py @@ -1,5 +1,6 @@ import os import re +import uuid from dataclasses import dataclass import psycopg @@ -107,17 +108,6 @@ def _iter_album_entries(dest_root: str): yield artist_name, album_folder, album_path -def _cursor_key(artist_name: str, album_folder: str) -> str: - return f"{artist_name}/{album_folder}" - - -def _cursor_tuple(cursor: str) -> tuple[str, str]: - # A cursor is "/"; artist and album are single path segments (no "/"), - # so partitioning on the first "/" recovers them, and tuple order matches the walk. - artist, _, album = (cursor or "").partition("/") - return artist, album - - def _process_album(conn, resolver, probe, browser, artist_name, album_folder, album_path, browsed) -> str: """Identify and record one album folder. Returns 'imported' (newly recorded), 'recorded' (matched but already present), 'skipped' (no MB match) or 'none' (no audio).""" @@ -139,38 +129,76 @@ def _process_album(conn, resolver, probe, browser, artist_name, album_folder, al return "imported" if _record_owned(conn, target, pr.quality_class, pr.fmt, album_path, watched_id) else "recorded" -def scan_chunk(conn: psycopg.Connection, resolver, probe: AudioProbe, browser, dest_root: str = "/music", - cursor: str = "", limit: int | None = None) -> tuple[str, int, int, bool]: - """Resolve and record album folders whose (artist, album) sorts after `cursor`, up to - `limit` of them (unbounded when limit is None). Returns (new_cursor, imported, skipped, - done); `done` is True when the walk was exhausted within this chunk. All writes are - ON CONFLICT, so an overlapping or repeated cursor is idempotent.""" - after = _cursor_tuple(cursor) - imported = skipped = processed = 0 +def build_worklist(conn: psycopg.Connection, scan_id: str, dest_root: str = "/music") -> int: + """Walk the tree once and persist this scan's album worklist. Idempotent + (ON CONFLICT DO NOTHING), so a re-run after a crash adds only missing rows. + Returns the number of rows inserted.""" + inserted = 0 + with conn.cursor() as cur: + for artist_name, album_folder, album_path in _iter_album_entries(dest_root): + cur.execute( + 'INSERT INTO "ScanWorkItem" (id, "scanId", artist, album, path, done) ' + "VALUES (gen_random_uuid()::text, %s, %s, %s, %s, false) " + 'ON CONFLICT ("scanId", path) DO NOTHING', + (scan_id, artist_name, album_folder, album_path), + ) + inserted += cur.rowcount + conn.commit() + return inserted + + +def scan_chunk(conn: psycopg.Connection, resolver, probe: AudioProbe, browser, + scan_id: str, limit: int) -> tuple[int, int, bool]: + """Process up to `limit` not-yet-done worklist items for `scan_id`, in stable + (artist, album) order, marking each done. Returns (imported, skipped, done) where + `done` is True when no undone rows remain. Idempotent: all record writes are + ON CONFLICT, and an item is marked done only after it is processed.""" + with conn.cursor() as cur: + cur.execute( + 'SELECT id, artist, album, path FROM "ScanWorkItem" ' + 'WHERE "scanId" = %s AND NOT done ORDER BY artist, album LIMIT %s', + (scan_id, limit), + ) + items = cur.fetchall() + + imported = skipped = 0 browsed: set[str] = set() # artists whose discography we've populated this chunk - new_cursor = cursor - done = True - for artist_name, album_folder, album_path in _iter_album_entries(dest_root): - if (artist_name, album_folder) <= after: - continue # already processed by an earlier chunk - if limit is not None and processed >= limit: - done = False # stopped early; more albums remain after new_cursor - break + for item_id, artist_name, album_folder, album_path in items: outcome = _process_album(conn, resolver, probe, browser, artist_name, album_folder, album_path, browsed) if outcome == "imported": imported += 1 elif outcome == "skipped": skipped += 1 - new_cursor = _cursor_key(artist_name, album_folder) - processed += 1 - return new_cursor, imported, skipped, done + with conn.cursor() as cur: + cur.execute('UPDATE "ScanWorkItem" SET done = true WHERE id = %s', (item_id,)) + conn.commit() + + with conn.cursor() as cur: + cur.execute('SELECT count(*) FROM "ScanWorkItem" WHERE "scanId" = %s AND NOT done', (scan_id,)) + remaining = cur.fetchone()[0] + return imported, skipped, remaining == 0 -def scan_library(conn: psycopg.Connection, resolver, probe: AudioProbe, browser, dest_root: str = "/music") -> ScanResult: - """Walk `{dest_root}/{Artist}/{Album}/` in one unbounded pass, resolving each album - against MusicBrainz and recording matched albums as have + monitored + followed (plus - each followed artist's full discography as unmonitored releases). Unmatched or audio-less - folders are skipped. This is `scan_chunk` with no limit; the worker loop uses the chunked - form so a large library doesn't block job-claim.""" - _cursor, imported, skipped, _done = scan_chunk(conn, resolver, probe, browser, dest_root, cursor="", limit=None) +def clear_worklist(conn: psycopg.Connection, scan_id: str) -> None: + with conn.cursor() as cur: + cur.execute('DELETE FROM "ScanWorkItem" WHERE "scanId" = %s', (scan_id,)) + conn.commit() + + +def scan_library(conn: psycopg.Connection, resolver, probe: AudioProbe, browser, + dest_root: str = "/music") -> ScanResult: + """Walk `{dest_root}/{Artist}/{Album}/` in one unbounded pass (build the worklist, + drain it, clean up), recording matched albums as have + monitored + followed plus + each followed artist's full discography. Test/one-shot convenience; the worker loop + uses the chunked build_worklist + scan_chunk so a large library doesn't block + job-claim.""" + scan_id = "oneshot-" + uuid.uuid4().hex + build_worklist(conn, scan_id, dest_root) + imported = skipped = 0 + done = False + while not done: + imp, skp, done = scan_chunk(conn, resolver, probe, browser, scan_id, limit=1000) + imported += imp + skipped += skp + clear_worklist(conn, scan_id) return ScanResult(imported, skipped) diff --git a/worker/tests/test_scan_chunk.py b/worker/tests/test_scan_chunk.py index bfdfc3a..fdc6127 100644 --- a/worker/tests/test_scan_chunk.py +++ b/worker/tests/test_scan_chunk.py @@ -1,13 +1,12 @@ from lyra_worker.adapters.fakes import FakeMbBrowser from lyra_worker.browser import ReleaseGroupInfo -from lyra_worker.scan import scan_chunk +from lyra_worker.scan import build_worklist, scan_chunk, clear_worklist from lyra_worker.types import MBTarget from tests.test_scan import FakeProbe, FakeResolver, _album, _counts def _tree(tmp_path): - """Three albums across two artists, in nested-sorted order: - Artist A/Alpha, Artist A/Beta, Artist B/Gamma.""" + """Three albums across two artists: Artist A/Alpha, Artist A/Beta, Artist B/Gamma.""" _album(tmp_path, "Artist A", "Alpha (2001)") _album(tmp_path, "Artist A", "Beta (2002)") _album(tmp_path, "Artist B", "Gamma (2003)") @@ -24,58 +23,66 @@ def _tree(tmp_path): return resolver, browser -def test_scan_chunk_respects_limit_and_returns_cursor(conn, tmp_path): - resolver, browser = _tree(tmp_path) - cursor, imported, skipped, done = scan_chunk( - conn, resolver, FakeProbe(quality_class=2), browser, str(tmp_path), cursor="", limit=2) - assert (imported, skipped, done) == (2, 0, False) # processed 2 of 3, more remain - assert cursor == "Artist A/Beta (2002)" # last processed key +def test_build_worklist_inserts_one_row_per_album(conn, tmp_path): + resolver, _ = _tree(tmp_path) + n = build_worklist(conn, "scan1", str(tmp_path)) + assert n == 3 with conn.cursor() as cur: + cur.execute('SELECT count(*) FROM "ScanWorkItem" WHERE "scanId" = %s AND NOT done', ("scan1",)) + assert cur.fetchone()[0] == 3 + + +def test_build_worklist_is_idempotent(conn, tmp_path): + resolver, _ = _tree(tmp_path) + assert build_worklist(conn, "scan1", str(tmp_path)) == 3 + assert build_worklist(conn, "scan1", str(tmp_path)) == 0 # ON CONFLICT DO NOTHING + with conn.cursor() as cur: + cur.execute('SELECT count(*) FROM "ScanWorkItem" WHERE "scanId" = %s', ("scan1",)) + assert cur.fetchone()[0] == 3 + + +def test_scan_chunk_respects_limit_and_marks_done(conn, tmp_path): + resolver, browser = _tree(tmp_path) + build_worklist(conn, "scan1", str(tmp_path)) + imported, skipped, done = scan_chunk(conn, resolver, FakeProbe(quality_class=2), browser, "scan1", limit=2) + assert (imported, skipped, done) == (2, 0, False) # 2 of 3 processed, more remain + with conn.cursor() as cur: + cur.execute('SELECT count(*) FROM "ScanWorkItem" WHERE "scanId" = %s AND done', ("scan1",)) + assert cur.fetchone()[0] == 2 cur.execute('SELECT count(*) FROM "LibraryItem"') assert cur.fetchone()[0] == 2 -def test_scan_chunk_resumes_from_cursor(conn, tmp_path): +def test_scan_chunk_resumes_and_finishes(conn, tmp_path): resolver, browser = _tree(tmp_path) - cursor, imported, skipped, done = scan_chunk( - conn, resolver, FakeProbe(quality_class=2), browser, str(tmp_path), - cursor="Artist A/Beta (2002)", limit=25) - assert (imported, skipped, done) == (1, 0, True) # only Gamma remains - assert cursor == "Artist B/Gamma (2003)" + build_worklist(conn, "scan1", str(tmp_path)) + scan_chunk(conn, resolver, FakeProbe(quality_class=2), browser, "scan1", limit=2) + imported, skipped, done = scan_chunk(conn, resolver, FakeProbe(quality_class=2), browser, "scan1", limit=25) + assert (imported, skipped, done) == (1, 0, True) # last album, then drained with conn.cursor() as cur: - cur.execute('SELECT album FROM "LibraryItem"') - assert cur.fetchall() == [("Gamma",)] # only Gamma imported + cur.execute('SELECT album FROM "LibraryItem" ORDER BY album') + assert cur.fetchall() == [("Alpha",), ("Beta",), ("Gamma",)] -def test_scan_chunk_cursor_past_end_is_noop(conn, tmp_path): +def test_scan_chunk_one_per_chunk_totals_correctly(conn, tmp_path): resolver, browser = _tree(tmp_path) - cursor, imported, skipped, done = scan_chunk( - conn, resolver, FakeProbe(), browser, str(tmp_path), cursor="zzzz/zzzz", limit=10) - assert (imported, skipped, done) == (0, 0, True) - assert cursor == "zzzz/zzzz" # unchanged - assert _counts(conn)["LibraryItem"] == 0 - - -def test_scan_chunk_splits_artist_across_chunks_idempotently(conn, tmp_path): - # limit=1 forces Artist A's two albums into separate chunks, re-browsing its - # discography — which must stay idempotent (ON CONFLICT) and total the same as one scan. - resolver, browser = _tree(tmp_path) - cursor, done, total_i, total_s, guard = "", False, 0, 0, 0 + build_worklist(conn, "scan1", str(tmp_path)) + total_i = total_s = guard = 0 + done = False while not done and guard < 10: - cursor, i, s, done = scan_chunk( - conn, resolver, FakeProbe(quality_class=2), browser, str(tmp_path), cursor=cursor, limit=1) + i, s, done = scan_chunk(conn, resolver, FakeProbe(quality_class=2), browser, "scan1", limit=1) total_i += i total_s += s guard += 1 assert (total_i, total_s) == (3, 0) - assert guard == 3 # three albums, one per chunk, then done + assert guard == 3 assert _counts(conn) == {"LibraryItem": 3, "MonitoredRelease": 3, "WatchedArtist": 2} -def test_scan_chunk_unbounded_matches_full_walk(conn, tmp_path): - resolver, browser = _tree(tmp_path) - cursor, imported, skipped, done = scan_chunk( - conn, resolver, FakeProbe(quality_class=2), browser, str(tmp_path), cursor="", limit=None) - assert (imported, skipped, done) == (3, 0, True) - assert cursor == "Artist B/Gamma (2003)" - assert _counts(conn) == {"LibraryItem": 3, "MonitoredRelease": 3, "WatchedArtist": 2} +def test_clear_worklist_removes_this_scans_rows(conn, tmp_path): + resolver, _ = _tree(tmp_path) + build_worklist(conn, "scan1", str(tmp_path)) + clear_worklist(conn, "scan1") + with conn.cursor() as cur: + cur.execute('SELECT count(*) FROM "ScanWorkItem" WHERE "scanId" = %s', ("scan1",)) + assert cur.fetchone()[0] == 0