"""Bounded, generation-safe media downloading and canonical slot cleanup.""" from __future__ import annotations import logging import os import signal import shutil import subprocess import tempfile import threading import time import uuid from collections.abc import Iterator from contextlib import contextmanager 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 _working_cookie_copy(self) -> Path | None: source = self.config.downloads.cookie_path if source is None: return None if source.is_file(): descriptor, temporary_name = tempfile.mkstemp( prefix="ytplayer-cookies-", suffix=".txt", ) temporary_path = Path(temporary_name) try: with os.fdopen(descriptor, "wb") as destination: with source.open("rb") as source_file: shutil.copyfileobj(source_file, destination) destination.flush() os.fsync(destination.fileno()) temporary_path.chmod(0o600) return temporary_path except BaseException: temporary_path.unlink(missing_ok=True) raise if self.config.downloads.cookie_required: raise FileNotFoundError( f"required yt-dlp cookie file is missing: {source}" ) LOG.warning( "yt-dlp cookie file is unavailable; continuing without it: %s", source, ) return None @contextmanager def _yt_dlp_arguments(self) -> Iterator[list[str]]: working_cookie = self._working_cookie_copy() arguments = [ self.config.downloads.yt_dlp_binary, "--ignore-config", "--no-js-runtimes", "--js-runtimes", "node", "--cache-dir", "/tmp/yt-dlp-cache", ] if working_cookie is not None: arguments.extend(("--cookies", str(working_cookie))) try: yield arguments finally: if working_cookie is not None: working_cookie.unlink(missing_ok=True) def _get_title(self, url: str) -> tuple[str, str | None]: try: with self._yt_dlp_arguments() as base_arguments: arguments = base_arguments + [ "--flat-playlist", "--print", "%(playlist_title)s|||%(title)s|||%(id)s", "--playlist-end", "1", url, ] 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) with self._yt_dlp_arguments() as base_arguments: arguments = base_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 any( guard.slot == token.slot for guard in self.store.list_live_guards() ): LOG.warning( "deferring interrupted deletion for protected slot %s", token.slot, ) continue 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 schedule_pending_deletion_recovery(self) -> bool: """Run recoverable trash cleanup without delaying service startup.""" with self._lock: if self._stopping: return False try: future = self._executor.submit(self.recover_pending_deletions) except BaseException: LOG.exception("could not schedule interrupted deletion recovery") return False with self._lock: self._futures.add(future) future.add_done_callback(self._forget_future) return True 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)