From 1ae43339227d8a3d9a69c9e3f7b20c2962b9450c Mon Sep 17 00:00:00 2001 From: npeter83 Date: Fri, 3 Jul 2026 00:08:18 +0200 Subject: [PATCH] =?UTF-8?q?feat(downloads):=20M2=20=E2=80=94=20worker=20do?= =?UTF-8?q?wnload=20loop=20+=20cache/dedup?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - downloads/formats.py: profile spec -> format_sig (cache key) + yt-dlp opts (format selector, audio extract, sponsorblock, embed subs/chapters/thumbnail, jpg poster) - downloads/storage.py: Plex tree path ({Channel}/Season {Year}/… [id].ext) + .nfo + poster.jpg sidecars; sanitization; asset-file deletion with empty-dir pruning - downloads/service.py: enqueue + get-or-create shared MediaAsset (race-safe) + instant cache-hit path (ready asset -> job done, ref_count++) - worker.py: N claim-loop threads (FOR UPDATE SKIP LOCKED), asset compare-and-set pending->downloading so identical requests never double-download, throttled progress writes, cooperative pause/cancel via the progress hook, crash-safe orphan recovery Verified on localdev: real download -> Plex tree + .nfo + poster on the host bind-mount, correct metadata/codecs/size; 2nd user = instant cache hit (shared asset, ref_count=2); audio-only MP3 extraction path produces a distinct cached asset. --- backend/app/worker.py | 350 +++++++++++++++++++++++++++++++++++++++--- 1 file changed, 331 insertions(+), 19 deletions(-) diff --git a/backend/app/worker.py b/backend/app/worker.py index 1b7564f..5276e0b 100644 --- a/backend/app/worker.py +++ b/backend/app/worker.py @@ -1,28 +1,339 @@ """Download worker process entry point. Runs in a DEDICATED container (`command: python -m app.worker`, `WORKER_ENABLED=1`), separate -from the API so long yt-dlp downloads + ffmpeg postprocessing never block web requests. It -claims `queued` download jobs from the DB with FOR UPDATE SKIP LOCKED and runs them, writing -live progress back to the job row (the frontend polls it — the worker is a separate process and -can't reach the API's in-process WebSocket registry). +from the API so long yt-dlp downloads + ffmpeg postprocessing never block web requests. -M1: thin, idle entry point (keeps the container healthy). The claim loop + yt-dlp integration -land in M2. +Design: + * N loop threads (download_worker_concurrency) each claim a `queued` DownloadJob with + FOR UPDATE SKIP LOCKED and process it. + * The shared MediaAsset is the dedup unit. A job whose asset is already `ready` finishes + instantly (no re-download); a fresh asset is claimed pending→downloading by exactly one + worker (a DB compare-and-set), so two jobs for the same video+format never double-download. + * Live progress is written to the job row (short-lived sessions, throttled ~1/s) — the + frontend polls it; the worker can't reach the API's in-process WebSocket registry. + * Pause/cancel are cooperative: the progress hook re-reads the job status and aborts the + yt-dlp download if it's no longer `running`. + * Crash-safe: on startup, orphaned `running` jobs are requeued and `downloading` assets reset. """ import logging +import shutil import signal +import threading import time +from datetime import date, datetime, timedelta, timezone +from pathlib import Path + +import yt_dlp +from sqlalchemy import text from app.config import settings +from app.db import SessionLocal +from app.downloads import formats, service, storage +from app.models import DownloadJob, MediaAsset log = logging.getLogger("siftlode.worker") -_stop = False +_stop = threading.Event() + + +class _Aborted(Exception): + """Raised inside the progress hook when the job was paused/canceled mid-download.""" def _handle_signal(signum, frame): - global _stop - _stop = True + _stop.set() + + +# --- small DB helpers (short-lived sessions; the download itself holds no transaction) ------- + +def _set_job(job_id: int, **fields) -> None: + with SessionLocal() as db: + job = db.get(DownloadJob, job_id) + if job is None: + return + for k, v in fields.items(): + setattr(job, k, v) + db.commit() + + +def _job_status(job_id: int) -> str | None: + with SessionLocal() as db: + return db.execute( + text("SELECT status FROM download_jobs WHERE id = :id"), {"id": job_id} + ).scalar() + + +def _parse_upload_date(raw) -> date | None: + if not raw: + return None + try: + return datetime.strptime(str(raw), "%Y%m%d").date() + except ValueError: + return None + + +# --- claim + recovery ----------------------------------------------------------------------- + +def _recover_orphans() -> None: + """Reset state left behind by a previous crash so nothing is stuck.""" + with SessionLocal() as db: + db.execute(text("UPDATE download_jobs SET status='queued' WHERE status='running'")) + db.execute(text("UPDATE media_assets SET status='pending' WHERE status='downloading'")) + db.commit() + + +def _claim_job() -> int | None: + with SessionLocal() as db: + row = db.execute( + text( + """ + UPDATE download_jobs SET status='running', phase='starting', updated_at=now() + WHERE id = ( + SELECT id FROM download_jobs + WHERE status='queued' + ORDER BY queue_pos, id + FOR UPDATE SKIP LOCKED + LIMIT 1 + ) + RETURNING id + """ + ) + ).fetchone() + db.commit() + return row[0] if row else None + + +def _claim_asset(asset_id: int) -> bool: + """Compare-and-set pending→downloading. True if THIS worker won the right to download.""" + with SessionLocal() as db: + row = db.execute( + text( + "UPDATE media_assets SET status='downloading' " + "WHERE id = :id AND status='pending' RETURNING id" + ), + {"id": asset_id}, + ).fetchone() + db.commit() + return row is not None + + +# --- processing ----------------------------------------------------------------------------- + +def _process_job(job_id: int) -> None: + with SessionLocal() as db: + job = db.get(DownloadJob, job_id) + if job is None: + return + asset = db.get(MediaAsset, job.asset_id) if job.asset_id else None + spec = dict(job.profile_snapshot or {}) + source_kind, source_ref = job.source_kind, job.source_ref + asset_id = asset.id if asset else None + asset_status = asset.status if asset else None + + if asset is None: + _set_job(job_id, status="error", error="No media asset", phase=None) + return + + if asset_status == "ready": + _finish_from_cache(job_id, asset_id) + return + if asset_status == "error": + _set_job(job_id, status="error", error=asset.error or "Download failed", phase=None) + return + if asset_status == "downloading": + # Another worker owns this asset; wait and let the job be reclaimed later. + _set_job(job_id, status="queued", phase="waiting") + time.sleep(2) + return + + # asset_status == "pending": try to win the download. + if not _claim_asset(asset_id): + _set_job(job_id, status="queued", phase="waiting") + time.sleep(2) + return + + _download(job_id, asset_id, source_kind, source_ref, spec) + + +def _finish_from_cache(job_id: int, asset_id: int) -> None: + with SessionLocal() as db: + job = db.get(DownloadJob, job_id) + asset = db.get(MediaAsset, asset_id) + if job and asset: + job.status = "done" + job.progress = 100 + job.phase = None + if not job.display_name: + job.display_name = asset.title + asset.last_access_at = datetime.now(timezone.utc) + db.commit() + + +def _make_progress_hook(job_id: int): + state = {"last_db": 0.0, "last_check": 0.0} + + def hook(d: dict) -> None: + now = time.monotonic() + # Cooperative pause/cancel: if the job is no longer 'running', abort the download. + if now - state["last_check"] >= 2.0: + state["last_check"] = now + if _job_status(job_id) not in ("running", None): + raise _Aborted + if d.get("status") != "downloading": + return + if now - state["last_db"] < 1.0: + return + state["last_db"] = now + total = d.get("total_bytes") or d.get("total_bytes_estimate") or 0 + done = d.get("downloaded_bytes") or 0 + pct = int(done * 100 / total) if total else 0 + _set_job( + job_id, + progress=min(pct, 99), + speed_bps=int(d["speed"]) if d.get("speed") else None, + eta_s=int(d["eta"]) if d.get("eta") else None, + phase="downloading", + ) + + return hook + + +def _staging_dir(asset_id: int) -> Path: + return Path(settings.download_root) / ".staging" / f"asset-{asset_id}" + + +def _find_thumbnail(staging: Path, media: Path) -> Path | None: + for p in staging.iterdir(): + if p != media and p.suffix.lower() in (".jpg", ".jpeg", ".png", ".webp"): + return p + return None + + +def _download(job_id: int, asset_id: int, source_kind: str, source_ref: str, spec: dict) -> None: + staging = _staging_dir(asset_id) + try: + if staging.exists(): + shutil.rmtree(staging, ignore_errors=True) + staging.mkdir(parents=True, exist_ok=True) + + outtmpl = str(staging / "%(id)s.%(ext)s") + opts = formats.build_ydl_opts(spec, outtmpl, _make_progress_hook(job_id)) + url = service.source_url(source_kind, source_ref) + + _set_job(job_id, phase="downloading", progress=0) + with yt_dlp.YoutubeDL(opts) as ydl: + info = ydl.extract_info(url, download=True) + + req = (info.get("requested_downloads") or [{}])[0] + produced = Path(req.get("filepath") or "") + if not produced.exists(): + # Fallback: the single media file left in staging (ignore the thumbnail). + cands = [p for p in staging.iterdir() if p.suffix.lower() not in (".jpg", ".jpeg", ".png", ".webp")] + if not cands: + raise RuntimeError("yt-dlp produced no output file") + produced = cands[0] + + ext = produced.suffix.lstrip(".").lower() + meta = storage.MediaMeta( + video_id=info.get("id") or source_ref, + title=info.get("title") or source_ref, + uploader=info.get("uploader") or info.get("channel") or "Unknown Channel", + upload_date=_parse_upload_date(info.get("upload_date")), + duration_s=int(info["duration"]) if info.get("duration") else None, + description=info.get("description"), + thumbnail_url=info.get("thumbnail"), + ) + + rel = storage.rel_path(meta, ext, settings.download_layout) + thumb = _find_thumbnail(staging, produced) + storage.place_file(produced, settings.download_root, rel) + nfo_ok = storage.write_sidecars(settings.download_root, rel, meta, thumb) + + final = storage.abs_path(settings.download_root, rel) + size = final.stat().st_size if final.exists() else None + expires = datetime.now(timezone.utc) + timedelta(days=settings.download_retention_days) + + with SessionLocal() as db: + asset = db.get(MediaAsset, asset_id) + if asset: + asset.status = "ready" + asset.rel_path = rel + asset.size_bytes = size + asset.duration_s = meta.duration_s + asset.width = int(info["width"]) if info.get("width") else None + asset.height = int(info["height"]) if info.get("height") else None + asset.container = ext + asset.vcodec = info.get("vcodec") if info.get("vcodec") not in (None, "none") else None + asset.acodec = info.get("acodec") if info.get("acodec") not in (None, "none") else None + asset.title = meta.title + asset.uploader = meta.uploader + asset.upload_date = meta.upload_date + asset.thumbnail_url = meta.thumbnail_url + asset.nfo_written = nfo_ok + asset.error = None + asset.last_access_at = datetime.now(timezone.utc) + asset.expires_at = expires + job = db.get(DownloadJob, job_id) + if job: + job.status = "done" + job.progress = 100 + job.phase = None + job.speed_bps = None + job.eta_s = None + if not job.display_name: + job.display_name = meta.title + db.commit() + log.info("job %s done: %s", job_id, rel) + + except _Aborted: + # Paused/canceled mid-download: free the asset for a later retry; keep the job's status + # (the route set it to paused/canceled). ref_count is adjusted by the route. + _reset_asset_pending(asset_id) + log.info("job %s aborted (pause/cancel)", job_id) + except Exception as exc: # noqa: BLE001 — surface any failure to the user, keep worker alive + msg = str(exc)[:500] + with SessionLocal() as db: + asset = db.get(MediaAsset, asset_id) + if asset: + asset.status = "error" + asset.error = msg + job = db.get(DownloadJob, job_id) + if job: + job.status = "error" + job.error = msg + job.phase = None + db.commit() + log.warning("job %s failed: %s", job_id, msg) + finally: + shutil.rmtree(staging, ignore_errors=True) + + +def _reset_asset_pending(asset_id: int) -> None: + with SessionLocal() as db: + asset = db.get(MediaAsset, asset_id) + if asset and asset.status == "downloading": + asset.status = "pending" + db.commit() + + +# --- loop ----------------------------------------------------------------------------------- + +def _worker_loop(idx: int) -> None: + while not _stop.is_set(): + try: + job_id = _claim_job() + except Exception: # noqa: BLE001 + log.exception("claim failed") + _stop.wait(settings.download_poll_seconds) + continue + if job_id is None: + _stop.wait(settings.download_poll_seconds) + continue + try: + _process_job(job_id) + except Exception: # noqa: BLE001 + log.exception("processing job %s crashed", job_id) + _set_job(job_id, status="error", error="Worker error", phase=None) def main() -> None: @@ -34,16 +345,17 @@ def main() -> None: log.info("WORKER_ENABLED is off; worker exiting (nothing to do).") return - log.info( - "Download worker started (concurrency=%d, root=%s). Claim loop lands in M2.", - settings.download_worker_concurrency, - settings.download_root, - ) - while not _stop: - # M2: claim queued jobs (FOR UPDATE SKIP LOCKED) and download them. - time.sleep(settings.download_poll_seconds) - - log.info("Download worker stopped.") + _recover_orphans() + n = max(1, settings.download_worker_concurrency) + log.info("Download worker started (concurrency=%d, root=%s).", n, settings.download_root) + threads = [threading.Thread(target=_worker_loop, args=(i,), daemon=True) for i in range(n)] + for t in threads: + t.start() + while not _stop.is_set(): + _stop.wait(1.0) + log.info("Download worker stopping…") + for t in threads: + t.join(timeout=5) if __name__ == "__main__":