"""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 edit as editmod 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 _wait_for_schema(timeout: float = 300.0) -> None: """Block until the download tables exist before doing any DB work. The API applies migrations on its OWN startup, and the worker is a separate container that may boot first (they start in parallel). Without this, the worker would hit a missing `download_jobs` table and crash-loop until the API's migrations land. Poll patiently instead.""" deadline = time.monotonic() + timeout while not _stop.is_set(): try: with SessionLocal() as db: db.execute(text("SELECT 1 FROM download_jobs LIMIT 1")) return except Exception as exc: # noqa: BLE001 — table not migrated yet / DB not up yet if time.monotonic() > deadline: log.warning("download schema still absent after %ss; continuing: %s", timeout, str(exc)[:120]) return log.info("waiting for the database schema (API migrations)…") _stop.wait(3.0) 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 job_kind = job.job_kind 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 / edit. if not _claim_asset(asset_id): _set_job(job_id, status="queued", phase="waiting") time.sleep(2) return if job_kind == "edit": _process_edit(job_id, asset_id) else: _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: # display_name is VARCHAR(255); a title can be longer (e.g. a recipe in the title). job.display_name = asset.title[:255] if asset.title else None asset.last_access_at = datetime.now(timezone.utc) db.commit() def _uploader_url(info: dict) -> str | None: """The channel/uploader page URL for the auto-clickable channel line. YouTube gives a `channel_url`/`uploader_url` directly; Facebook gives neither but exposes a numeric `uploader_id` (the page id) that resolves at facebook.com/. Returns None for sources with no derivable channel page (e.g. reddit / a raw media URL).""" url = info.get("channel_url") or info.get("uploader_url") if url: return url uid = info.get("uploader_id") tag = f"{info.get('extractor_key') or info.get('extractor') or ''} {info.get('webpage_url') or ''}".lower() if uid and "facebook" in tag: return f"https://www.facebook.com/{uid}" return None 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.uploader_url = _uploader_url(info) 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")) asset.source_webpage_url = info.get("webpage_url") 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 # Errors worth retrying with a different player-client set (YouTube per-client flakiness). _RETRIABLE_MARKERS = ( "not a bot", "DRM protected", "Requested format is not available", "Only images are available", "Sign in to confirm", ) def _extract_with_fallback(job_id: int, asset_id: int, url: str, spec: dict, staging: Path): """Download via yt-dlp, trying the primary player clients then the POT-backed fallback set on a bot/DRM/format failure. Returns the info dict of the successful attempt.""" client_sets = [ c for c in (settings.player_client_list, settings.player_client_fallback_list) if c ] or [None] outtmpl = str(staging / "%(id)s.%(ext)s") last_exc: Exception | None = None for i, clients in enumerate(client_sets): opts = formats.build_ydl_opts( spec, outtmpl, _make_progress_hook(job_id, asset_id), settings.download_pot_base_url, clients, ) opts["postprocessor_hooks"] = [_make_pp_hook(job_id)] try: with yt_dlp.YoutubeDL(opts) as ydl: return ydl.extract_info(url, download=True) except yt_dlp.utils.DownloadError as exc: last_exc = exc retriable = any(m in str(exc) for m in _RETRIABLE_MARKERS) if i < len(client_sets) - 1 and retriable: log.info("job %s: clients %s failed (%s); retrying with fallback", job_id, clients, str(exc)[:70]) for p in list(staging.iterdir()): # drop partials before the retry try: p.unlink() except OSError: pass _set_job(job_id, progress=0, phase="downloading") continue raise raise last_exc # unreachable (loop either returns or raises) 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) _set_job(job_id, phase="downloading", progress=0) url = service.source_url(source_kind, source_ref) info = _extract_with_fallback(job_id, asset_id, url, spec, staging) 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.uploader_url = _uploader_url(info) asset.upload_date = meta.upload_date asset.thumbnail_url = meta.thumbnail_url # No thumbnail from the source (e.g. a direct media URL) → cut a poster frame from # the file so the card / watch page / link preview still shows an image. if not meta.thumbnail_url: asset.poster_path = editmod.ensure_poster(asset, settings.download_root) asset.source_webpage_url = info.get("webpage_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: # display_name is VARCHAR(255) — clip a long title (see _fill_asset_meta_early). job.display_name = meta.title[:255] if meta.title else None 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() # --- editor derivatives (phase 2: trim/crop via ffmpeg) ------------------------------------- def _edit_staging_dir(asset_id: int) -> Path: return Path(settings.download_root) / ".staging" / f"edit-{asset_id}" def _fail_edit(job_id: int, asset_id: int, msg: str) -> None: 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("edit job %s failed: %s", job_id, msg) def _process_edit(job_id: int, asset_id: int) -> None: """Produce a per-user trim/crop clip from the source asset's downloaded file via ffmpeg.""" with SessionLocal() as db: job = db.get(DownloadJob, job_id) src = db.get(MediaAsset, job.source_asset_id) if job and job.source_asset_id else None edit_spec = dict(job.edit_spec or {}) if job else {} user_id = job.user_id if job else None src_ready = src is not None and src.status == "ready" and bool(src.rel_path) src_rel = src.rel_path if src else None src_w, src_h, src_dur = (src.width, src.height, src.duration_s) if src else (None, None, None) if not src_ready or not src_rel: _fail_edit(job_id, asset_id, "The source download is no longer available.") return src_path = storage.abs_path(settings.download_root, src_rel) if not src_path.exists(): _fail_edit(job_id, asset_id, "The source file is missing.") return src_ext = src_path.suffix.lstrip(".").lower() or "mp4" crop = editmod.has_crop(edit_spec) out_ext = "mp4" if editmod.needs_reencode(edit_spec) else src_ext total_s = editmod.clip_duration(edit_spec, src_dur) or src_dur or 0 staging = _edit_staging_dir(asset_id) try: if staging.exists(): shutil.rmtree(staging, ignore_errors=True) staging.mkdir(parents=True, exist_ok=True) dest_staging = staging / f"out.{out_ext}" _set_job(job_id, phase="editing", progress=0, speed_bps=None, eta_s=None) if edit_spec.get("segments"): cmd, prep = editmod.build_concat_plan(src_path, dest_staging, edit_spec, out_ext, staging) for p, content in prep.items(): p.write_text(content, encoding="utf-8") else: cmd = editmod.build_edit_cmd(src_path, dest_staging, edit_spec, out_ext) editmod.run_ffmpeg( cmd, total_s, on_progress=lambda pct: _set_job(job_id, progress=int(pct), phase="editing"), should_cancel=lambda: _job_status(job_id) not in ("running", None), ) if not dest_staging.exists(): raise RuntimeError("ffmpeg produced no output file") rel = storage.edit_rel_path(user_id, asset_id, out_ext) storage.place_file(dest_staging, settings.download_root, rel) 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.container = out_ext asset.duration_s = total_s or asset.duration_s if crop: asset.width = edit_spec["crop"]["w"] asset.height = edit_spec["crop"]["h"] else: asset.width, asset.height = src_w, src_h 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 db.commit() log.info("edit job %s done: %s", job_id, rel) except editmod.EditAborted: # Paused/canceled mid-encode: free the asset for a later retry (ref_count is the route's job). _reset_asset_pending(asset_id) log.info("edit job %s aborted (pause/cancel)", job_id) except Exception as exc: # noqa: BLE001 — surface the failure, keep the worker alive _fail_edit(job_id, asset_id, str(exc)[:500]) finally: shutil.rmtree(staging, ignore_errors=True) # --- 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 _wait_for_schema() # the API may still be applying migrations when we boot if _stop.is_set(): 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()