1ae4333922
- 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.
362 lines
13 KiB
Python
362 lines
13 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 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 = 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:
|
|
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:
|
|
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()
|