"""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. 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, quota, service, storage from app.models import DownloadJob, MediaAsset from app.titles import normalize_title log = logging.getLogger("siftlode.worker") _stop = threading.Event() class _Aborted(Exception): """Raised inside the progress hook when the job was paused/canceled mid-download.""" def _handle_signal(signum, frame): _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: """Claim the oldest queued job whose user is under their per-user concurrency limit. Fetches a window of locked candidates (FOR UPDATE SKIP LOCKED) and picks the first eligible one, so one user flooding the queue can't starve others beyond their max_concurrent.""" with SessionLocal() as db: rows = db.execute( text( """ SELECT id, user_id FROM download_jobs WHERE status='queued' ORDER BY queue_pos, id FOR UPDATE SKIP LOCKED LIMIT 50 """ ) ).fetchall() chosen = next( (jid for jid, uid in rows if not quota.at_concurrency_limit(db, uid)), None ) if chosen is None: db.commit() # release the row locks; nothing eligible right now return None db.execute( text( "UPDATE download_jobs SET status='running', phase='starting', updated_at=now() " "WHERE id = :id" ), {"id": chosen}, ) db.commit() return chosen 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 _fill_asset_meta_early(asset_id: int, info: dict) -> None: """Populate an asset's display metadata from yt-dlp's info_dict the first time it's seen (only if not already filled at enqueue from our catalog) — so an ad-hoc URL's row shows a real title + thumbnail while it downloads, not just the URL. Zero extra network cost.""" with SessionLocal() as db: asset = db.get(MediaAsset, asset_id) if asset is None or asset.title: return asset.title = normalize_title(info.get("title")) asset.uploader = info.get("uploader") or info.get("channel") asset.thumbnail_url = info.get("thumbnail") asset.duration_s = int(info["duration"]) if info.get("duration") else None asset.upload_date = _parse_upload_date(info.get("upload_date")) db.commit() def _make_progress_hook(job_id: int, asset_id: int): state = {"last_db": 0.0, "last_check": 0.0, "meta_done": False} 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 not state["meta_done"] and d.get("info_dict"): state["meta_done"] = True try: _fill_asset_meta_early(asset_id, d["info_dict"]) except Exception: # noqa: BLE001 — metadata is best-effort, never fail a download pass if d.get("status") != "downloading": return if now - state["last_db"] < 1.0: return state["last_db"] = now # yt-dlp downloads the video and audio streams SEPARATELY (each 0→100%), then merges. # Label the phase so the user understands the bar resetting between streams instead of # seeing a mysterious 95%→5% jump. info = d.get("info_dict") or {} vcodec = info.get("vcodec") or "none" acodec = info.get("acodec") or "none" if vcodec != "none" and acodec == "none": phase = "video" elif acodec != "none" and vcodec == "none": phase = "audio" else: phase = "media" 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=phase, ) return hook def _pp_phase(pp: str) -> str: """Map a yt-dlp postprocessor class name to a user-facing phase label. ffmpeg post-steps aren't byte-progress operations (no %), so the UI shows the step name + an indeterminate pulse — this makes the long merge/finalize legible instead of a silent 'Processing'.""" if "Merg" in pp: return "merging" if "ExtractAudio" in pp: return "audio_extract" if "Thumbnail" in pp: return "thumbnail" if "SponsorBlock" in pp or "ModifyChapters" in pp: return "sponsorblock" if "Metadata" in pp: return "metadata" return "processing" def _make_pp_hook(job_id: int): """Postprocessor hook: surface the current ffmpeg step so the bar shows a concrete label (Merging / Embedding thumbnail / …) instead of freezing at 100% after the streams finish.""" def hook(d: dict) -> None: if d.get("status") in ("started", "processing"): _set_job(job_id, phase=_pp_phase(d.get("postprocessor") or ""), speed_bps=None, eta_s=None) 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, asset_id)) opts["postprocessor_hooks"] = [_make_pp_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=normalize_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: logging.basicConfig(level=logging.INFO) signal.signal(signal.SIGTERM, _handle_signal) signal.signal(signal.SIGINT, _handle_signal) if not settings.worker_enabled: log.info("WORKER_ENABLED is off; worker exiting (nothing to do).") return _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__": main()