Lyra/worker/lyra_worker/main.py

import os
import sys
import threading
import time
import uuid

import psycopg

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
from lyra_worker.discovery import DiscoveryConfig, DiscoveryResult, run_discovery
from lyra_worker.library import clear_staging_root
from lyra_worker.monitor import MonitorConfig, fulfill_owned_releases, reconcile, sweep
from lyra_worker.pipeline import run_pipeline
from lyra_worker.registry import (
    build_adapters, build_browser, build_probe, build_resolver,
    build_similarity_sources, build_tagger,
)
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)
MONITOR_TICK_SECONDS = 60.0
DISCOVER_TICK_SECONDS = 300.0  # check the discovery interval every 5 min (interval itself is hours)
CACHE_PRUNE_SECONDS = 3600.0  # prune the shared ApiCache at most hourly (drops rows > 90d old)
DEST_ROOT = "/music"  # must match run_pipeline's default library root
# Per-job download staging. Defaults inside the library (/music/.staging) to preserve prior
# behavior; set STAGING_ROOT (mounted as a separate volume) to keep temp download I/O off a
# network-share library.
STAGING_ROOT = os.environ.get("STAGING_ROOT") or f"{DEST_ROOT}/.staging"
DEFAULT_SCAN_CHUNK = 25  # albums resolved per loop iteration (~25-50s at MB ~1 req/s)
_TRUE = {"1", "true", "yes", "on"}


def _set_config(conn, key: str, value: str) -> None:
    with conn.cursor() as cur:
        cur.execute(
            'INSERT INTO "Config"(key, value, secret, "updatedAt") VALUES (%s, %s, false, now()) '
            'ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value, "updatedAt" = now()',
            (key, value),
        )
    conn.commit()


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 (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.id", scan_id)
        _set_config(conn, "scan.progress", "0/0")
        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.id", "")
            _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:
        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()
        clear_worklist(conn, scan_id)
        _set_config(conn, "scan.id", "")
        _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.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.progress", f"{imported_total}/{skipped_total}")


def _run_discovery(conn, sources, browser, config=None) -> DiscoveryResult:
    cfg = DiscoveryConfig.from_config(config if config is not None else get_config(conn))
    return run_discovery(conn, sources, browser, cfg)


def finish_job(conn, job_id, mcfg) -> None:
    """Feed a finished job's outcome back into its MonitoredRelease. Runs for EVERY job,
    monitor enabled or not: a manually-wanted release must still transition to fulfilled so
    it drops off the Wanted list. A no-op for plain (unlinked) requests."""
    try:
        reconcile(conn, job_id, mcfg)
    except Exception as e:  # reconcile must never take down the worker loop
        print(f"worker: reconcile failed for job {job_id}: {e}", flush=True)
        conn.rollback()


def maybe_run_discovery(conn, sources, browser, config, dcfg, now, last_discover_tick,
                        tick_seconds: float = DISCOVER_TICK_SECONDS) -> float:
    """Advance the discovery sweep by at most one chunk (``discover.chunkSize`` seeds) per call,
    so the worker loop keeps claiming jobs between chunks. A sweep starts on ``discover.requested``
    or a due schedule tick, then continues via ``discover.inProgress`` until a chunk drains 0 seeds
    (every eligible seed refreshed), at which point ``discover.result`` is written. Returns the
    (possibly advanced) last_discover_tick."""
    requested = str(config.get("discover.requested", "")).strip().lower() in _TRUE
    in_progress = str(config.get("discover.inProgress", "")).strip().lower() in _TRUE
    due = dcfg.enabled and now - last_discover_tick >= tick_seconds
    if not (requested or due or in_progress):
        return last_discover_tick

    starting = not in_progress
    if starting:  # begin a fresh sweep
        _set_config(conn, "discover.inProgress", "true")
        _set_config(conn, "discover.requested", "false")

    try:
        result = _run_discovery(conn, sources, browser, config)  # one chunk
    except Exception as e:  # a discovery error must never kill the worker
        print(f"worker: discovery run failed: {e}", flush=True)
        conn.rollback()
        _set_config(conn, "discover.inProgress", "false")  # abandon; a new request restarts it
        _set_config(conn, "discover.requested", "false")
        return now if due else last_discover_tick

    a_prev, b_prev = (0, 0) if starting else _parse_progress(config.get("discover.progress", ""))
    a_total, b_total = a_prev + result.artists, b_prev + result.albums
    if result.seeds == 0:  # every eligible seed refreshed -> sweep drained
        _set_config(conn, "discover.result", f"artists {a_total}, albums {b_total}")
        _set_config(conn, "discover.progress", "")
        _set_config(conn, "discover.inProgress", "false")
        print(f"worker: discovery done — {a_total} artists, {b_total} albums", flush=True)
    else:
        _set_config(conn, "discover.progress", f"{a_total}/{b_total}")
    return now if due else last_discover_tick


def prune_api_cache(conn) -> None:
    """Drop ApiCache rows older than 90 days so the shared cache stays bounded. Best-effort:
    a failure here (e.g. a transient error) must never take down the worker loop."""
    try:
        with conn.cursor() as cur:
            cur.execute('DELETE FROM "ApiCache" WHERE "fetchedAt" < now() - interval \'90 days\'')
        conn.commit()
    except psycopg.OperationalError:
        raise  # a real disconnect — let the loop's reconnect handler deal with it
    except Exception as e:
        print(f"worker: ApiCache prune failed: {e}", flush=True)
        try:
            conn.rollback()
        except Exception:
            pass


def _require_secret_key() -> None:
    """Fail fast (or warn loudly) on a missing/placeholder/public LYRA_SECRET_KEY before
    the worker touches any encrypted credential."""
    problem = secret_key_problem()
    if problem is None:
        return
    msg, fatal = problem
    banner = "=" * 72
    if fatal:
        print(
            f"\n{banner}\nFATAL: {msg}.\nGenerate one:  openssl rand -base64 32\n"
            f"Then set LYRA_SECRET_KEY in your .env and restart.\n{banner}\n",
            flush=True,
        )
        sys.exit(1)
    print(f"\n{banner}\nWARNING: {msg}.\n{banner}\n", flush=True)


HEARTBEAT_FILE = "/tmp/worker.alive"


def _liveness_heartbeat() -> None:
    """Touch HEARTBEAT_FILE every HEARTBEAT_SECONDS from a background daemon thread, so the
    compose healthcheck sees the process as alive even while the main loop is blocked in a
    long download (a loop-driven heartbeat would go stale and falsely report unhealthy)."""
    while True:
        try:
            with open(HEARTBEAT_FILE, "w") as fh:
                fh.write("1")
        except OSError:
            pass
        time.sleep(HEARTBEAT_SECONDS)


def run_forever() -> None:
    _require_secret_key()
    threading.Thread(target=_liveness_heartbeat, daemon=True).start()
    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()
    _dsn = os.environ.get("DATABASE_URL")
    browser = build_browser(_dsn)  # shared ApiCache read-through when a DSN is present
    probe = build_probe()
    similarity_sources = build_similarity_sources(get_config(conn), _dsn)
    # Backfill: releases pressed while the monitor was off never got reconciled, so they linger
    # on the wanted list. Fix any already-owned ones once at startup.
    fixed = fulfill_owned_releases(conn, MonitorConfig.from_config(get_config(conn)))
    if fixed:
        print(f"worker: reconciled {fixed} already-owned release(s) to fulfilled", flush=True)
    print(f"worker: {len(adapters)} adapter(s) enabled: {[a.name for a in adapters]}", flush=True)
    print("worker: waiting for jobs", flush=True)
    last_tick = 0.0
    last_discover_tick = 0.0
    last_heartbeat = 0.0
    last_prune = 0.0
    try:
        while 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:
                config = get_config(conn)
                mcfg = MonitorConfig.from_config(config)
                now = time.monotonic()
                if now - last_heartbeat >= HEARTBEAT_SECONDS:
                    _set_config(conn, "worker.heartbeat", "1")  # value irrelevant; updatedAt is the signal
                    last_heartbeat = now
                if now - last_prune >= CACHE_PRUNE_SECONDS:
                    prune_api_cache(conn)  # keep the shared ApiCache bounded
                    last_prune = now
                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:
                    run_pipeline(conn, job_id, adapters, resolver=resolver, tagger=tagger,
                                 staging_root=STAGING_ROOT, 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()
                finish_job(conn, job_id, mcfg)
                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()


if __name__ == "__main__":
    run_forever()