502 lines
17 KiB
Python
502 lines
17 KiB
Python
"""Bounded, generation-safe media downloading and canonical slot cleanup."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import os
|
|
import signal
|
|
import shutil
|
|
import subprocess
|
|
import threading
|
|
import time
|
|
import uuid
|
|
from concurrent.futures import Future, ThreadPoolExecutor
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any, Callable
|
|
|
|
from .config import ControllerConfig
|
|
from .contracts import SLOT_IDS
|
|
from .models import DeleteToken, DownloadReservation, SlotProtectedError, StaleOperationError
|
|
from .store import ControllerStore
|
|
|
|
|
|
LOG = logging.getLogger(__name__)
|
|
MEDIA_EXTENSIONS = frozenset({".mp4", ".mkv", ".webm", ".mov", ".m4v"})
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class DownloadEvent:
|
|
slot: int
|
|
state: str
|
|
title: str | None = None
|
|
count: int | None = None
|
|
message: str | None = None
|
|
job_id: str | None = None
|
|
generation: int | None = None
|
|
url: str | None = None
|
|
submitter: str | None = None
|
|
note: str | None = None
|
|
reused_slot: bool = False
|
|
|
|
|
|
EventCallback = Callable[[DownloadEvent], None]
|
|
Runner = Callable[..., subprocess.CompletedProcess[str]]
|
|
|
|
|
|
def canonical_slot_directory(media_base: Path, slot: int) -> Path:
|
|
if slot not in SLOT_IDS:
|
|
raise ValueError(f"slot must be in 1..10, got {slot!r}")
|
|
return media_base / f"slot{slot}"
|
|
|
|
|
|
def canonical_playlist_path(media_base: Path, slot: int) -> Path:
|
|
return canonical_slot_directory(media_base, slot) / "playlist.m3u"
|
|
|
|
|
|
def title_from_filename(path: Path) -> str:
|
|
stem = path.stem
|
|
if " - " in stem:
|
|
stem = stem.split(" - ", 1)[1]
|
|
if " [" in stem:
|
|
stem = stem.rsplit(" [", 1)[0]
|
|
return stem
|
|
|
|
|
|
class DownloadManager:
|
|
"""Reserve synchronously, then perform work on a bounded executor."""
|
|
|
|
def __init__(
|
|
self,
|
|
config: ControllerConfig,
|
|
store: ControllerStore,
|
|
*,
|
|
on_event: EventCallback | None = None,
|
|
runner: Runner | None = None,
|
|
executor: ThreadPoolExecutor | None = None,
|
|
):
|
|
self.config = config
|
|
self.store = store
|
|
self.on_event = on_event or (lambda event: None)
|
|
self.runner = runner
|
|
self._executor = executor or ThreadPoolExecutor(
|
|
max_workers=config.downloads.max_workers,
|
|
thread_name_prefix="yt-download",
|
|
)
|
|
self._owns_executor = executor is None
|
|
self._lock = threading.Lock()
|
|
self._futures: set[Future[Any]] = set()
|
|
self._processes: set[subprocess.Popen[str]] = set()
|
|
self._stopping = False
|
|
|
|
def enqueue(
|
|
self,
|
|
url: str,
|
|
submitter: str | None,
|
|
note: str | None,
|
|
) -> DownloadReservation:
|
|
with self._lock:
|
|
if self._stopping:
|
|
return DownloadReservation(status="shutting_down")
|
|
reservation = self.store.reserve_download(url, submitter, note)
|
|
if reservation.status != "reserved":
|
|
return reservation
|
|
assert reservation.slot is not None
|
|
assert reservation.job_id is not None
|
|
assert reservation.generation is not None
|
|
|
|
# Publish/record queued after the transaction commits and before a
|
|
# worker can publish downloading or ready.
|
|
self._emit(
|
|
DownloadEvent(
|
|
slot=reservation.slot,
|
|
state="queued",
|
|
job_id=reservation.job_id,
|
|
generation=reservation.generation,
|
|
url=url,
|
|
submitter=submitter,
|
|
note=note,
|
|
reused_slot=reservation.reused_slot,
|
|
)
|
|
)
|
|
try:
|
|
future = self._executor.submit(
|
|
self._download,
|
|
url,
|
|
reservation.slot,
|
|
reservation.generation,
|
|
reservation.job_id,
|
|
)
|
|
except BaseException as exc:
|
|
message = f"could not start download worker: {exc}"
|
|
self.store.fail_download(
|
|
reservation.job_id,
|
|
reservation.slot,
|
|
reservation.generation,
|
|
message,
|
|
)
|
|
self._emit(
|
|
DownloadEvent(
|
|
slot=reservation.slot,
|
|
state="error",
|
|
message=message,
|
|
job_id=reservation.job_id,
|
|
generation=reservation.generation,
|
|
)
|
|
)
|
|
return DownloadReservation(
|
|
status="worker_error",
|
|
slot=reservation.slot,
|
|
generation=reservation.generation,
|
|
job_id=reservation.job_id,
|
|
reused_slot=reservation.reused_slot,
|
|
)
|
|
|
|
with self._lock:
|
|
self._futures.add(future)
|
|
future.add_done_callback(self._forget_future)
|
|
return reservation
|
|
|
|
def _forget_future(self, future: Future[Any]) -> None:
|
|
with self._lock:
|
|
self._futures.discard(future)
|
|
try:
|
|
future.result()
|
|
except BaseException:
|
|
LOG.exception("download worker escaped its failure boundary")
|
|
|
|
def _emit(self, event: DownloadEvent) -> None:
|
|
try:
|
|
self.on_event(event)
|
|
except BaseException:
|
|
LOG.exception("download event callback failed")
|
|
|
|
def _cookie_arguments(self) -> list[str]:
|
|
path = self.config.downloads.cookie_path
|
|
if path is None:
|
|
return []
|
|
if path.is_file():
|
|
return ["--cookies", str(path)]
|
|
if self.config.downloads.cookie_required:
|
|
raise FileNotFoundError(f"required yt-dlp cookie file is missing: {path}")
|
|
LOG.warning("yt-dlp cookie file is unavailable; continuing without it: %s", path)
|
|
return []
|
|
|
|
def _base_yt_dlp_arguments(self) -> list[str]:
|
|
return [self.config.downloads.yt_dlp_binary, "--ignore-config"] + self._cookie_arguments()
|
|
|
|
def _get_title(self, url: str) -> tuple[str, str | None]:
|
|
arguments = self._base_yt_dlp_arguments() + [
|
|
"--flat-playlist",
|
|
"--print",
|
|
"%(playlist_title)s|||%(title)s|||%(id)s",
|
|
"--playlist-end",
|
|
"1",
|
|
url,
|
|
]
|
|
try:
|
|
result = self._run_command(
|
|
arguments,
|
|
check=True,
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=min(120.0, self.config.downloads.process_timeout_seconds),
|
|
)
|
|
line = (result.stdout or "").strip().splitlines()[0]
|
|
playlist_title, _, remainder = line.partition("|||")
|
|
video_title, _, yt_id = remainder.partition("|||")
|
|
playlist_title = playlist_title.strip()
|
|
video_title = video_title.strip()
|
|
title = (
|
|
playlist_title
|
|
if playlist_title and playlist_title.upper() not in {"NA", "NONE"}
|
|
else video_title
|
|
)
|
|
return title or url, yt_id.strip() or None
|
|
except (IndexError, OSError, subprocess.SubprocessError, ValueError) as exc:
|
|
LOG.warning("could not determine media title for %s: %s", url, exc)
|
|
return url, None
|
|
|
|
def _download(
|
|
self,
|
|
url: str,
|
|
slot: int,
|
|
generation: int,
|
|
job_id: str,
|
|
) -> None:
|
|
if not self.store.mark_download_running(job_id, slot, generation):
|
|
return
|
|
self._emit(
|
|
DownloadEvent(
|
|
slot=slot,
|
|
state="downloading",
|
|
job_id=job_id,
|
|
generation=generation,
|
|
)
|
|
)
|
|
slot_directory = canonical_slot_directory(self.config.storage.media_base, slot)
|
|
temporary_playlist = slot_directory / f".playlist.{job_id}.tmp"
|
|
try:
|
|
slot_directory.mkdir(parents=True, exist_ok=True)
|
|
if slot_directory.is_symlink():
|
|
raise RuntimeError(f"refusing to download through symlink: {slot_directory}")
|
|
title, yt_id = self._get_title(url)
|
|
arguments = self._base_yt_dlp_arguments() + [
|
|
"--yes-playlist",
|
|
"--no-overwrites",
|
|
"--no-progress",
|
|
"-o",
|
|
str(
|
|
slot_directory
|
|
/ "%(playlist_index)03d - %(title)s [%(id)s].%(ext)s"
|
|
),
|
|
"-f",
|
|
(
|
|
f"bv*[height<={self.config.downloads.max_height}]+ba/"
|
|
f"b[height<={self.config.downloads.max_height}]/b"
|
|
),
|
|
url,
|
|
]
|
|
self._run_command(
|
|
arguments,
|
|
check=True,
|
|
text=True,
|
|
timeout=self.config.downloads.process_timeout_seconds,
|
|
)
|
|
media_files = sorted(
|
|
path
|
|
for path in slot_directory.iterdir()
|
|
if path.is_file() and path.suffix.lower() in MEDIA_EXTENSIONS
|
|
)
|
|
if not media_files:
|
|
raise RuntimeError("yt-dlp completed without producing media files")
|
|
|
|
items = [
|
|
(title_from_filename(path), os.path.relpath(path, slot_directory))
|
|
for path in media_files
|
|
]
|
|
with temporary_playlist.open("w", encoding="utf-8", newline="\n") as stream:
|
|
for _, relative_path in items:
|
|
stream.write(relative_path + "\n")
|
|
stream.flush()
|
|
os.fsync(stream.fileno())
|
|
os.replace(
|
|
temporary_playlist,
|
|
canonical_playlist_path(self.config.storage.media_base, slot),
|
|
)
|
|
self._fsync_directory(slot_directory)
|
|
|
|
if not self.store.complete_download(
|
|
job_id,
|
|
slot,
|
|
generation,
|
|
title=title,
|
|
items=items,
|
|
yt_id=yt_id,
|
|
):
|
|
LOG.warning(
|
|
"ignored stale download completion job=%s slot=%s generation=%s",
|
|
job_id,
|
|
slot,
|
|
generation,
|
|
)
|
|
return
|
|
self._emit(
|
|
DownloadEvent(
|
|
slot=slot,
|
|
state="ready",
|
|
title=title,
|
|
count=len(items),
|
|
job_id=job_id,
|
|
generation=generation,
|
|
)
|
|
)
|
|
except BaseException as exc:
|
|
temporary_playlist.unlink(missing_ok=True)
|
|
message = self._failure_message(exc)
|
|
self.store.fail_download(job_id, slot, generation, message)
|
|
self._emit(
|
|
DownloadEvent(
|
|
slot=slot,
|
|
state="error",
|
|
message=message,
|
|
job_id=job_id,
|
|
generation=generation,
|
|
)
|
|
)
|
|
|
|
@staticmethod
|
|
def _failure_message(exc: BaseException) -> str:
|
|
if isinstance(exc, subprocess.TimeoutExpired):
|
|
return f"yt-dlp timed out after {exc.timeout} seconds"
|
|
if isinstance(exc, subprocess.CalledProcessError):
|
|
return f"yt-dlp failed with exit code {exc.returncode}"
|
|
return f"download failed: {exc}"
|
|
|
|
def _run_command(
|
|
self,
|
|
arguments: list[str],
|
|
*,
|
|
check: bool,
|
|
text: bool,
|
|
timeout: float,
|
|
capture_output: bool = False,
|
|
) -> subprocess.CompletedProcess[str]:
|
|
if self.runner is not None:
|
|
return self.runner(
|
|
arguments,
|
|
check=check,
|
|
text=text,
|
|
timeout=timeout,
|
|
capture_output=capture_output,
|
|
)
|
|
|
|
with self._lock:
|
|
if self._stopping:
|
|
raise RuntimeError("Controller is shutting down")
|
|
process = subprocess.Popen(
|
|
arguments,
|
|
stdout=subprocess.PIPE if capture_output else None,
|
|
stderr=subprocess.PIPE if capture_output else None,
|
|
text=text,
|
|
start_new_session=True,
|
|
)
|
|
self._processes.add(process)
|
|
try:
|
|
try:
|
|
stdout, stderr = process.communicate(timeout=timeout)
|
|
except subprocess.TimeoutExpired:
|
|
self._terminate_process(process, force=False)
|
|
try:
|
|
stdout, stderr = process.communicate(timeout=2.0)
|
|
except subprocess.TimeoutExpired:
|
|
self._terminate_process(process, force=True)
|
|
stdout, stderr = process.communicate()
|
|
raise subprocess.TimeoutExpired(
|
|
arguments,
|
|
timeout,
|
|
output=stdout,
|
|
stderr=stderr,
|
|
)
|
|
completed = subprocess.CompletedProcess(
|
|
arguments,
|
|
process.returncode,
|
|
stdout=stdout,
|
|
stderr=stderr,
|
|
)
|
|
if check and process.returncode:
|
|
raise subprocess.CalledProcessError(
|
|
process.returncode,
|
|
arguments,
|
|
output=stdout,
|
|
stderr=stderr,
|
|
)
|
|
return completed
|
|
finally:
|
|
with self._lock:
|
|
self._processes.discard(process)
|
|
|
|
@staticmethod
|
|
def _terminate_process(
|
|
process: subprocess.Popen[str],
|
|
*,
|
|
force: bool,
|
|
) -> None:
|
|
if process.poll() is not None:
|
|
return
|
|
try:
|
|
os.killpg(
|
|
process.pid,
|
|
signal.SIGKILL if force else signal.SIGTERM,
|
|
)
|
|
except ProcessLookupError:
|
|
return
|
|
except OSError:
|
|
try:
|
|
process.kill() if force else process.terminate()
|
|
except ProcessLookupError:
|
|
pass
|
|
|
|
@staticmethod
|
|
def _fsync_directory(path: Path) -> None:
|
|
descriptor = os.open(path, os.O_RDONLY | getattr(os, "O_DIRECTORY", 0))
|
|
try:
|
|
os.fsync(descriptor)
|
|
finally:
|
|
os.close(descriptor)
|
|
|
|
def delete_slot(
|
|
self,
|
|
slot: int,
|
|
*,
|
|
reason: str,
|
|
expected_url: str | None = None,
|
|
) -> bool:
|
|
token = self.store.begin_delete(slot, reason=reason, expected_url=expected_url)
|
|
return self.finish_deletion(token)
|
|
|
|
def finish_deletion(self, token: DeleteToken) -> bool:
|
|
slot_directory = canonical_slot_directory(self.config.storage.media_base, token.slot)
|
|
trash_root = self.config.storage.media_base / ".controller-trash"
|
|
trash_path: Path | None = None
|
|
if os.path.lexists(slot_directory):
|
|
if slot_directory.is_symlink():
|
|
slot_directory.unlink()
|
|
else:
|
|
trash_root.mkdir(parents=True, exist_ok=True)
|
|
trash_path = trash_root / (
|
|
f"slot{token.slot}.generation{token.generation}.{uuid.uuid4().hex}"
|
|
)
|
|
os.replace(slot_directory, trash_path)
|
|
self._fsync_directory(self.config.storage.media_base)
|
|
if trash_path is not None:
|
|
shutil.rmtree(trash_path)
|
|
return self.store.finish_delete(token)
|
|
|
|
def expire_due_slots(self) -> list[int]:
|
|
deleted: list[int] = []
|
|
for token in self.store.begin_expired_deletions():
|
|
try:
|
|
if self.finish_deletion(token):
|
|
deleted.append(token.slot)
|
|
except OSError:
|
|
LOG.exception("could not finish expired deletion for slot %s", token.slot)
|
|
return deleted
|
|
|
|
def recover_pending_deletions(self) -> list[int]:
|
|
recovered: list[int] = []
|
|
trash_root = self.config.storage.media_base / ".controller-trash"
|
|
for token in self.store.pending_deletions():
|
|
if trash_root.is_dir():
|
|
prefix = f"slot{token.slot}.generation{token.generation}."
|
|
for candidate in trash_root.iterdir():
|
|
if candidate.name.startswith(prefix):
|
|
if candidate.is_symlink() or candidate.is_file():
|
|
candidate.unlink()
|
|
elif candidate.is_dir():
|
|
shutil.rmtree(candidate)
|
|
if self.finish_deletion(token):
|
|
recovered.append(token.slot)
|
|
return recovered
|
|
|
|
def shutdown(self, *, wait: bool = True) -> None:
|
|
with self._lock:
|
|
self._stopping = True
|
|
processes = tuple(self._processes)
|
|
for process in processes:
|
|
self._terminate_process(process, force=False)
|
|
deadline = time.monotonic() + 2.0
|
|
for process in processes:
|
|
if process.poll() is not None:
|
|
continue
|
|
try:
|
|
process.wait(timeout=max(0.0, deadline - time.monotonic()))
|
|
except subprocess.TimeoutExpired:
|
|
self._terminate_process(process, force=True)
|
|
if self._owns_executor:
|
|
self._executor.shutdown(wait=wait, cancel_futures=True)
|
|
|
|
def active_future_count(self) -> int:
|
|
with self._lock:
|
|
return len(self._futures)
|