Files
siftlode/backend/app/worker.py
T
peter 0ed82fe6b6 feat(downloads): run each download/edit in a supervised child process with a hard-kill watchdog
Defense-in-depth after the ReDoS incident. The heavy work (yt-dlp download / ffmpeg edit)
now runs in a child process (app.downloads.runner) supervised by the worker thread, not in
the thread itself. A thread can't be killed and a C-level spin can't even be cooperatively
cancelled, so one wedged download used to pin a core, hold the GIL and freeze the whole
queue; a child process can be SIGTERM/SIGKILLed.

The supervisor kills the child (and its process group, so ffmpeg grandchildren die too)
when: a route pauses/cancels the job, it makes no progress for download_stall_timeout_seconds
(150s), it runs past download_wall_clock_max_seconds (2h), the worker is shutting down, or
the child crashes without setting state. On a kill it frees the asset and fails the job with
a retryable network error (stalls are not auto-requeued). A clean exit is left alone — the
child already wrote the terminal state, including a cooperative pause via the progress hook.

- config.py: stall / wall-clock / supervisor-poll / kill-grace settings.
- worker.py: _run_supervised + _spawn_download_child / _job_liveness / _stop_child /
  _cleanup_staging; _process_job claims the asset, then hands the heavy work to the child.
- app/downloads/runner.py: the child entrypoint.
- docker-compose.localdev.yml: init:true on the worker so tini (PID 1) reaps orphaned
  ffmpeg grandchildren left by a hard-kill (the worker reaps its own direct children).
- tests: watchdog decision logic + SIGTERM->SIGKILL escalation (fake proc, no real subprocess).
2026-08-12 22:39:17 +02:00

1030 lines
44 KiB
Python

"""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 os
import re
import shutil
import signal
import subprocess
import sys
import threading
import time
from datetime import date, datetime, timedelta, timezone
from pathlib import Path
import yt_dlp
from sqlalchemy import text
from app import sysconfig
from app.config import settings
from app.db import SessionLocal
from app.downloads import edit as editmod
from app.downloads import errors as dlerrors
from app.downloads import formats, quota, service, storage
from app.models import DownloadJob, MediaAsset
from app.notifications import create_notification
from app.titles import normalize_title
log = logging.getLogger("siftlode.worker")
def _notify_download_done(
db, user_id: int | None, title: str | None, job_id, *, kind: str
) -> None:
"""Post a best-effort 'ready' notice to the owner's inbox — downloads/edits used to finish
silently while scheduler jobs and playlist-push already notified (R11 feedback symmetry). A
notification failure must never mask the completed job, so it is isolated + rolled back (mirrors
scheduler._notify_done). `kind` ("download"/"edit") lets the inbox pick accurate copy; the
English title here is only a fallback for an unknown type. `user_id`/`title` are captured by the
caller BEFORE its commit, so reading them here needs no post-commit refresh of an expired row."""
if user_id is None:
return
try:
create_notification(
db,
user_id=user_id,
type="download_done",
title=f"Ready to watch: {title}" if title else "Ready to watch",
data={"title": title, "kind": kind},
)
except Exception:
db.rollback()
log.exception("failed to write download-done notification for job %s", job_id)
_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 _requeue_if_running(job_id: int, phase: str) -> None:
"""Put the job back to 'queued' for a later retry, but ONLY while it is still 'running'. The
route flips status to paused/cancelled to stop a job; a blind `status='queued'` here would
clobber that and silently revive a job the user just cancelled. The conditional WHERE makes the
requeue a no-op unless we still own a running job."""
with SessionLocal() as db:
db.execute(
text(
"UPDATE download_jobs SET status='queued', phase=:phase, updated_at=now() "
"WHERE id=:id AND status='running'"
),
{"id": job_id, "phase": phase},
)
db.commit()
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) -> bool:
"""Block until the download tables exist before doing any DB work. Returns False if they never
appeared (or we were asked to stop while waiting).
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 — and if the schema
is still absent when the timeout runs out, say so and let the caller exit, because "continuing"
only moved the same crash one call further down (into `_recover_orphans`), where it read as an
unrelated failure."""
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 True
except Exception as exc: # noqa: BLE001 — table not migrated yet / DB not up yet
if time.monotonic() > deadline:
log.error(
"download schema still absent after %ss — the API's migrations have not landed; "
"exiting so the container restarts: %s",
timeout, str(exc)[:120],
)
return False
log.info("waiting for the database schema (API migrations)…")
_stop.wait(3.0)
return False
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
# --- supervised child execution + watchdog --------------------------------------------------
# The heavy work runs in a child process so a wedged yt-dlp/ffmpeg is HARD-killable. A worker thread
# can't be killed, a C-level spin (e.g. a catastrophic regex) can't even be cooperatively cancelled,
# and one such thread starves the rest via the GIL — the incident that motivated this. The parent
# thread supervises: it kills the child (and its process group, so ffmpeg grandchildren die too) on
# a route pause/cancel, on a no-progress stall, or past an absolute wall-clock cap.
_CHILD_OK = 0
def _spawn_download_child(job_id: int) -> subprocess.Popen:
"""Start the child that does the actual download/edit for `job_id`, in its OWN session/process
group so stopping it also stops any ffmpeg / helper processes yt-dlp spawned."""
return subprocess.Popen(
[sys.executable, "-m", "app.downloads.runner", str(job_id)],
start_new_session=True,
)
def _job_liveness(job_id: int) -> tuple[str | None, float | None]:
"""(status, seconds-since-last-row-update) for the job, computed with the DB clock so there is
no parent↔DB skew. A live download rewrites its row ~1/s; a stalled one never does. Returns
(None, None) if the row is gone."""
with SessionLocal() as db:
row = db.execute(
text(
"SELECT status, EXTRACT(EPOCH FROM (now() - updated_at)) "
"FROM download_jobs WHERE id = :id"
),
{"id": job_id},
).one_or_none()
if row is None:
return None, None
return row[0], (float(row[1]) if row[1] is not None else None)
def _signal_group(proc: subprocess.Popen, pgid: int | None, sig: int) -> None:
"""Signal the child's whole process group (falling back to the child alone), swallowing the race
where it has already exited."""
try:
if pgid is not None:
os.killpg(pgid, sig)
else:
proc.send_signal(sig)
except (ProcessLookupError, OSError):
pass
def _stop_child(proc: subprocess.Popen, grace_s: float) -> None:
"""SIGTERM the child's process group, then SIGKILL if it outlives the grace period — the same
terminate→wait→kill escalation as run_ffmpeg, so a child that ignores SIGTERM is still reaped."""
if proc.poll() is not None:
return
try:
pgid: int | None = os.getpgid(proc.pid)
except OSError:
pgid = None
_signal_group(proc, pgid, signal.SIGTERM)
try:
proc.wait(timeout=grace_s)
return
except subprocess.TimeoutExpired:
pass
_signal_group(proc, pgid, signal.SIGKILL)
try:
proc.wait(timeout=grace_s)
except subprocess.TimeoutExpired:
log.error("download child pid %s survived SIGKILL", proc.pid)
def _run_supervised(job_id: int, asset_id: int) -> None:
"""Run the heavy download/edit for an already-claimed asset in a child process and watchdog it.
A clean child exit means the child already wrote the terminal job/asset state (done/error, or a
cooperative pause via the in-child progress hook) — we leave it. We step in only when the child
can't finish or self-abort: a route paused/cancelled it (kill, free the asset, keep the route's
status), it made no progress for `stall_timeout` or ran past `wall_clock_max` (kill, fail the
job), the worker is shutting down (kill, requeue), or it crashed without setting state."""
poll_s = max(1, settings.download_supervisor_poll_seconds)
stall_s = settings.download_stall_timeout_seconds
wall_s = settings.download_wall_clock_max_seconds
grace_s = settings.download_kill_grace_seconds
try:
proc = _spawn_download_child(job_id)
except Exception: # noqa: BLE001 — spawn failed; don't leave the asset stuck 'downloading'
log.exception("failed to spawn download child for job %s", job_id)
_reset_asset_pending(asset_id)
_requeue_if_running(job_id, None)
return
started = time.monotonic()
reason: str | None = None
while True:
try:
proc.wait(timeout=poll_s)
break # child finished on its own
except subprocess.TimeoutExpired:
pass
status, stale = _job_liveness(job_id)
if status is None or status != "running":
reason = "cancelled" # a route paused/cancelled it (or the row vanished)
break
if stale is not None and stale > stall_s:
reason = "stalled"
break
if time.monotonic() - started > wall_s:
reason = "timeout"
break
if _stop.is_set():
reason = "shutdown"
break
if reason is None:
# Child exited. If it died WITHOUT setting a terminal state (non-zero exit / crash / OOM),
# the job is still 'running' — fail it here so nothing is left stuck.
if proc.returncode != _CHILD_OK:
status, _ = _job_liveness(job_id)
if status == "running":
_reset_asset_pending(asset_id)
_cleanup_staging(asset_id)
_set_job(
job_id, status="error", error_code="unknown", phase=None,
error=f"The download process exited unexpectedly (code {proc.returncode}).",
)
log.warning("job %s: child exited %s without finishing", job_id, proc.returncode)
return
_stop_child(proc, grace_s)
_reset_asset_pending(asset_id)
_cleanup_staging(asset_id)
if reason == "cancelled":
# The route already set paused/cancelled (or the job is gone); just free the asset.
log.info("job %s: child stopped (route pause/cancel)", job_id)
return
if reason == "shutdown":
_requeue_if_running(job_id, None) # resume on the next boot
return
msg = (
f"The download stalled (no progress for {int(stall_s)}s) and was stopped."
if reason == "stalled"
else f"The download exceeded the {int(wall_s)}s time limit and was stopped."
)
_set_job(job_id, status="error", error=msg, error_code="network", phase=None)
log.warning("job %s: stopped by watchdog (%s)", job_id, reason)
def _cleanup_staging(asset_id: int) -> None:
"""Best-effort removal of a killed child's leftover staging dirs (its own `finally` never ran).
The next attempt on this asset also wipes staging on start, so this is just tidiness/disk."""
for path in (_staging_dir(asset_id), _edit_staging_dir(asset_id)):
shutil.rmtree(path, ignore_errors=True)
# --- processing -----------------------------------------------------------------------------
def _process_job(job_id: int) -> None:
"""Dispatch a claimed job. The trivial outcomes (cache hit, dead/contended asset) are handled
inline and cheaply; the actual heavy work (yt-dlp download / ffmpeg edit) is claimed here and
then run in a supervised CHILD PROCESS (`_run_supervised`) so a wedged yt-dlp/ffmpeg can be
hard-killed — a worker THREAD cannot be killed, nor (when it spins in a C call) even
cooperatively cancelled, and one such thread starves the others via the GIL."""
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
asset_id = asset.id if asset else None
asset_status = asset.status if asset else None
asset_error = asset.error 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.
_requeue_if_running(job_id, "waiting")
time.sleep(2)
return
# asset_status == "pending": try to win the download / edit.
if not _claim_asset(asset_id):
_requeue_if_running(job_id, "waiting")
time.sleep(2)
return
# The asset is now ours ('downloading'): run the heavy work in a supervised child process.
_run_supervised(job_id, asset_id)
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 build_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/<id>. 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
# `or info.get("id")` avoids a blank card while downloading when the cleaned title is empty
# (e.g. a title that was only a social engagement prefix); the completion path later sets
# the final title with its own source_ref fallback.
asset.title = normalize_title(info.get("title")) or info.get("id")
asset.uploader = info.get("uploader") or info.get("channel")
asset.uploader_url = build_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 = {"meta_done": False}
cancelled = _cancel_poll(job_id) # same policy as the ffmpeg path
allow_write = _rate_limiter(_DB_WRITE_INTERVAL_S)
def hook(d: dict) -> None:
# Cooperative pause/cancel: if the job is no longer 'running', abort the download.
if cancelled():
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 not allow_write():
return
# 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
# Both progress sources call back FAR faster than the database should be touched: yt-dlp fires its
# hook per chunk, and ffmpeg's `-progress pipe:1` emits a block of ~12 lines every ~0.5s while
# run_ffmpeg calls its callbacks once per LINE. Every callback here is a round-trip (`_job_status`
# opens a session and SELECTs; `_set_job` opens one and UPDATE+COMMITs), so unthrottled they cost
# ~10^5 queries on a long job. ONE interval, ONE implementation, used by both paths — three
# hand-rolled copies is how the ffmpeg side ended up unthrottled while the download side was fine.
_DB_WRITE_INTERVAL_S = 1.0
_CANCEL_POLL_INTERVAL_S = 1.0
def _rate_limiter(interval_s: float, *, now=time.monotonic):
"""`should_run()` → True at most once per `interval_s` (the first call always passes).
`now` is injectable so tests drive a clock instead of patching the global one."""
# `None`, not 0.0: "never ran" has to be a VALUE, not a falsy timestamp. A truthiness test here
# reads a legitimate `now() == 0.0` as "never ran" and lets a second call straight through —
# unreachable with real `time.monotonic()`, but `now` exists to be injected, and a test clock
# starting at 0 is the obvious choice. That would let a throttle test encode the wrong
# behaviour and still pass.
state: dict[str, float | None] = {"at": None}
def should_run() -> bool:
t = now()
last = state["at"]
if last is not None and t - last < interval_s:
return False
state["at"] = t
return True
return should_run
def _cancel_poll(job_id: int, interval_s: float = _CANCEL_POLL_INTERVAL_S, *, now=time.monotonic):
"""Cooperative "has this job been paused/cancelled?" poll, rate-limited and sticky: once it
says yes it never queries again. Shared by the yt-dlp hook and run_ffmpeg so cancel latency is
one policy, not one per call site."""
allow = _rate_limiter(interval_s, now=now)
state = {"cancelled": False}
def cancelled() -> bool:
if not state["cancelled"] and allow():
state["cancelled"] = _job_status(job_id) not in ("running", None)
return state["cancelled"]
return cancelled
def _ffmpeg_progress_writer(job_id: int, phase: str, *, now=time.monotonic):
"""`on_progress` for run_ffmpeg: persist at most one row write per interval, and only when the
whole-percent value actually changed (the bar shows integers)."""
allow = _rate_limiter(_DB_WRITE_INTERVAL_S, now=now)
state = {"pct": -1}
def write(pct: float) -> None:
whole = int(pct)
if whole == state["pct"] or not allow():
return
state["pct"] = whole
_set_job(job_id, progress=whole, phase=phase)
return write
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}"
_THUMB_SUFFIXES = (".jpg", ".jpeg", ".png", ".webp")
# yt-dlp's working files in the staging dir. None of them is ever the finished media, but several
# CAN be the biggest file there (a video-only per-format stream, an abandoned fragment run), so the
# staging fallback must exclude them by name — picking one would store e.g. a silent video-only
# stream as a ready asset.
_YTDLP_TEMP_RE = re.compile(
r"\.part(-Frag\d+)?$" # in-flight download / fragment
r"|\.ytdl$" # resume state
r"|\.temp\.[^.]+$" # merge/postprocess temp
r"|\.f\d+\.[^.]+$", # per-format stream kept before the merge
re.IGNORECASE,
)
def _find_thumbnail(staging: Path, media: Path) -> Path | None:
for p in staging.iterdir():
if p != media and p.suffix.lower() in _THUMB_SUFFIXES:
return p
return None
def _produced_file(info: dict, staging: Path) -> Path:
"""The media file yt-dlp actually wrote.
`requested_downloads[0].filepath` is the authoritative answer on YouTube, but an arbitrary-URL
extractor may not fill it in at all — and `Path("")` is `Path(".")`, which *passes* `.exists()`,
so a missing value used to sail past the fallback and blow up two steps later (`with_name` on an
empty name: the BBC-news failure). Anything that isn't a real file counts as absent; then fall
back to the largest file left in staging that is neither a thumbnail nor one of yt-dlp's own
working files (see `_YTDLP_TEMP_RE`)."""
req = (info.get("requested_downloads") or [{}])[0]
fp = req.get("filepath") or req.get("_filename") or info.get("filepath") or info.get("_filename")
if fp:
produced = Path(str(fp))
if produced.is_file():
return produced
cands = [
p
for p in staging.iterdir()
if p.is_file()
and p.suffix.lower() not in _THUMB_SUFFIXES
and not _YTDLP_TEMP_RE.search(p.name)
]
if not cands:
raise RuntimeError("yt-dlp produced no output file")
return max(cands, key=lambda p: p.stat().st_size)
def _media_info(info: dict) -> dict:
"""Unwrap a playlist/multi_video extraction to the entry that carries the download.
A single non-YouTube page (a news article's video, a multi-rendition embed) can extract as a
wrapper whose own dict has no `requested_downloads` — the media lives on an entry. Pick the
FIRST entry that actually downloaded something: on a page with a skipped teaser in front of the
real video, entries[0] carries no file, and taking it blindly would label the produced file
with the wrong id/title/duration."""
entries = [e for e in (info.get("entries") or []) if e]
if not entries or info.get("requested_downloads"):
return info
chosen = next((e for e in entries if e.get("requested_downloads")), entries[0])
return {**info, **chosen}
# 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_settings() -> tuple[str, int]:
"""Layout + retention days from the DB config (admin-overridable on the Configuration page),
falling back to the env defaults. Read here, not from `settings`, so an admin edit takes effect."""
with SessionLocal() as db:
return (
sysconfig.get_str(db, "download_layout"),
sysconfig.get_int(db, "download_retention_days"),
)
def _ensure_browser_playable(job_id: int, produced: Path, spec: dict, info: dict) -> Path:
"""Make an mp4-compat-profile download play in every browser: H.264 + AAC/MP3 in mp4. If the
download (whatever the source shipped) isn't already that, re-encode the offending stream(s) —
or just remux when the codecs are fine but the container isn't. **Scope:** only the mp4 compat
profile is normalized; an explicit `vcodec` or a non-mp4 custom profile is the user opting out
of the guarantee and is left as-is (so a `webm` profile can still ship VP9/Opus). Returns the
file to store (converted, or the original) and corrects `info`'s codec fields to match disk."""
if spec.get("container") != "mp4" or spec.get("vcodec"):
return produced
vcodec, acodec = editmod.probe_codecs(produced)
is_mp4 = produced.suffix.lower() == ".mp4"
needs, reenc_v, reenc_a = editmod.browser_transcode_flags(vcodec, acodec, is_mp4)
if not needs:
return produced
dest = produced.with_name(produced.stem + ".browser.mp4")
action = "re-encode" if (reenc_v or reenc_a) else "remux"
log.info("job %s: %s (v=%s a=%s mp4=%s) -> browser-safe mp4", job_id, action, vcodec, acodec, is_mp4)
_set_job(job_id, phase="processing", progress=0, speed_bps=None, eta_s=None)
# Prefer the info duration; probe the file when the extractor didn't provide one, so a long
# re-encode still shows real progress instead of a frozen 0%.
total_s = int(info["duration"]) if info.get("duration") else editmod.probe_duration(produced)
editmod.run_ffmpeg(
editmod.build_browser_safe_cmd(produced, dest, reenc_v, reenc_a),
total_s=total_s,
on_progress=_ffmpeg_progress_writer(job_id, "processing"),
should_cancel=_cancel_poll(job_id),
)
produced.unlink(missing_ok=True)
# Keep the stored asset metadata truthful about what's now on disk.
if reenc_v:
info["vcodec"] = "h264"
if reenc_a:
info["acodec"] = "aac"
return dest
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 = _media_info(_extract_with_fallback(job_id, asset_id, url, spec, staging))
produced = _produced_file(info, staging)
produced = _ensure_browser_playable(job_id, produced, spec, info)
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"),
)
layout, retention_days = _download_settings()
rel = storage.rel_path(meta, ext, layout, asset_id=asset_id)
thumb = _find_thumbnail(staging, produced)
# Convert the raw thumbnail (webp/png/jpeg) to a real JPEG ourselves — robustly — for the
# poster + og:image sidecar, now that yt-dlp's fatal thumbnail convertor is gone. On failure
# this is None and write_sidecars writes no poster, so ensure_poster cuts a frame instead.
poster_jpg = editmod.thumbnail_to_jpg(thumb, staging / "poster.jpg")
storage.place_file(produced, settings.download_root, rel)
nfo_ok = storage.write_sidecars(settings.download_root, rel, meta, poster_jpg)
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=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 = build_uploader_url(info)
asset.upload_date = meta.upload_date
asset.thumbnail_url = meta.thumbnail_url
# Record a self-hosted poster for every download: it's the reliable link-preview
# (og:image) source — a remote thumbnail, esp. Facebook's signed CDN URL, expires
# and can't be relied on cross-origin. write_sidecars already wrote <base>.jpg from
# the thumbnail; ensure_poster returns that, or cuts a frame if there was none.
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
# Capture the owner before commit expires the row (avoids a refresh SELECT just to notify).
owner_id = job.user_id if job else None
db.commit()
_notify_download_done(db, owner_id, meta.title, job_id, kind="download")
log.info("job %s done: %s", job_id, rel)
except (_Aborted, editmod.EditAborted):
# Paused/canceled mid-download or mid-transcode: 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.error_code = dlerrors.classify(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.error_code = dlerrors.classify(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=_ffmpeg_progress_writer(job_id, "editing"),
should_cancel=_cancel_poll(job_id),
)
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
_, retention_days = _download_settings()
expires = datetime.now(timezone.utc) + timedelta(days=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
# Capture owner + title before commit expires the rows (avoids a refresh SELECT).
owner_id = job.user_id if job else None
dl_title = asset.title if asset else None
db.commit()
_notify_download_done(db, owner_id, dl_title, job_id, kind="edit")
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
# The API may still be applying migrations when we boot. If they never arrive, exit instead of
# pressing on into a guaranteed crash — with a FAILURE status, so the container comes back under
# any restart policy (`on-failure` too, not just the `unless-stopped` our composes ship) and the
# exit code doesn't claim success for a worker that never started working.
if not _wait_for_schema():
if _stop.is_set():
return # asked to shut down while waiting — that IS a clean exit
raise SystemExit(1)
_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()