diff --git a/controller/downloads.py b/controller/downloads.py new file mode 100644 index 0000000..2d20ddb --- /dev/null +++ b/controller/downloads.py @@ -0,0 +1,395 @@ +"""Bounded, generation-safe media downloading and canonical slot cleanup.""" + +from __future__ import annotations + +import logging +import os +import shutil +import subprocess +import threading +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 + + +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 = subprocess.run, + 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._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, + ) + ) + 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.runner( + 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.runner( + 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}" + + @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 + 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) diff --git a/controller/store.py b/controller/store.py index 473c9e7..7647c1d 100644 --- a/controller/store.py +++ b/controller/store.py @@ -1195,6 +1195,23 @@ class ControllerStore: ) return True + def pending_deletions(self) -> list[DeleteToken]: + with self._lock: + rows = self.connection.execute( + """ + SELECT slot_number, generation FROM slots + WHERE media_state = 'deleting' ORDER BY slot_number + """ + ).fetchall() + return [ + DeleteToken( + slot=row["slot_number"], + generation=row["generation"], + reason="recovered_interrupted_deletion", + ) + for row in rows + ] + def begin_expired_deletions(self, now: float | None = None) -> list[DeleteToken]: now = self._now() if now is None else now tokens: list[DeleteToken] = [] diff --git a/tests/test_downloads.py b/tests/test_downloads.py new file mode 100644 index 0000000..fc73b6e --- /dev/null +++ b/tests/test_downloads.py @@ -0,0 +1,199 @@ +import os +import subprocess +import tempfile +import unittest +from concurrent.futures import Future +from pathlib import Path + +from controller.config import ControllerConfig, DownloadConfig, StorageConfig +from controller.downloads import DownloadManager, canonical_playlist_path +from controller.models import SlotProtectedError +from controller.store import ControllerStore + + +class ImmediateExecutor: + def submit(self, function, *args, **kwargs): + future = Future() + try: + future.set_result(function(*args, **kwargs)) + except BaseException as exc: + future.set_exception(exc) + return future + + +class FakeYtDlp: + def __init__(self, media_base: Path, *, fail_download=False, no_media=False): + self.media_base = media_base + self.fail_download = fail_download + self.no_media = no_media + self.calls = [] + + def __call__(self, arguments, **kwargs): + self.calls.append((list(arguments), dict(kwargs))) + if "--flat-playlist" in arguments: + return subprocess.CompletedProcess( + arguments, + 0, + stdout="Test Playlist|||First Video|||playlist-id\n", + stderr="", + ) + if self.fail_download: + raise subprocess.CalledProcessError(7, arguments) + if not self.no_media: + output_index = arguments.index("-o") + 1 + output_template = Path(arguments[output_index]) + directory = output_template.parent + directory.mkdir(parents=True, exist_ok=True) + (directory / "001 - First [one].mp4").write_bytes(b"video-one") + (directory / "002 - Second [two].webm").write_bytes(b"video-two") + return subprocess.CompletedProcess(arguments, 0, stdout="", stderr="") + + +class DownloadManagerTests(unittest.TestCase): + def setUp(self): + self.tempdir = tempfile.TemporaryDirectory() + base = Path(self.tempdir.name) + self.media = base / "media" + self.data = base / "data" + self.cookie = base / "cookies.txt" + self.cookie.write_text("cookie data", encoding="utf-8") + self.config = ControllerConfig( + storage=StorageConfig( + media_base=self.media, + data_base=self.data, + retention_seconds=100, + ), + downloads=DownloadConfig( + yt_dlp_binary="yt-dlp-test", + cookie_path=self.cookie, + cookie_required=True, + max_workers=1, + process_timeout_seconds=30, + max_height=720, + ), + ) + self.store = ControllerStore(self.config.storage).open() + + def tearDown(self): + self.store.close() + self.tempdir.cleanup() + + def manager(self, runner, events=None): + event_list = events if events is not None else [] + return DownloadManager( + self.config, + self.store, + runner=runner, + executor=ImmediateExecutor(), + on_event=event_list.append, + ) + + def test_success_uses_configured_cookie_and_atomically_builds_playlist(self): + events = [] + runner = FakeYtDlp(self.media) + manager = self.manager(runner, events) + result = manager.enqueue("https://example.test/playlist", "Alice", "Service") + self.assertEqual(result.status, "reserved") + self.assertEqual([event.state for event in events], ["queued", "downloading", "ready"]) + slot = self.store.get_slot(1) + self.assertEqual(slot.media_state, "ready") + self.assertEqual(slot.title, "Test Playlist") + self.assertEqual(slot.item_count, 2) + playlist = canonical_playlist_path(self.media, 1) + self.assertEqual( + playlist.read_text(encoding="utf-8"), + "001 - First [one].mp4\n002 - Second [two].webm\n", + ) + self.assertEqual(list(playlist.parent.glob(".playlist.*.tmp")), []) + for arguments, _ in runner.calls: + self.assertIn("--ignore-config", arguments) + self.assertIn("--cookies", arguments) + self.assertIn(str(self.cookie), arguments) + download_arguments = runner.calls[-1][0] + self.assertIn("bv*[height<=720]+ba/b[height<=720]/b", download_arguments) + + def test_worker_failure_is_contained_and_persisted(self): + events = [] + manager = self.manager(FakeYtDlp(self.media, fail_download=True), events) + result = manager.enqueue("https://example.test/failure", None, None) + self.assertEqual(result.status, "reserved") + slot = self.store.get_slot(1) + self.assertEqual(slot.media_state, "error") + self.assertIn("exit code 7", slot.last_error) + self.assertEqual(events[-1].state, "error") + + def test_missing_required_cookie_fails_job_without_escaping(self): + self.cookie.unlink() + events = [] + runner = FakeYtDlp(self.media) + manager = self.manager(runner, events) + manager.enqueue("https://example.test/no-cookie", None, None) + self.assertEqual(self.store.get_slot(1).media_state, "error") + self.assertIn("required yt-dlp cookie", self.store.get_slot(1).last_error) + self.assertEqual(runner.calls, []) + + def test_delete_moves_only_canonical_slot_and_keeps_other_identity(self): + first = self.store.reserve_download("https://example.test/one", None, None, job_id="one") + second = self.store.reserve_download("https://example.test/two", None, None, job_id="two") + self.store.fail_download("one", 1, first.generation, "failed") + self.store.fail_download("two", 2, second.generation, "failed") + slot_one = self.media / "slot1" + slot_two = self.media / "slot2" + slot_one.mkdir(parents=True) + slot_two.mkdir(parents=True) + (slot_one / "one.mp4").write_bytes(b"one") + (slot_two / "two.mp4").write_bytes(b"two") + manager = self.manager(FakeYtDlp(self.media)) + self.assertTrue(manager.delete_slot(1, reason="manual")) + self.assertFalse(slot_one.exists()) + self.assertTrue((slot_two / "two.mp4").exists()) + self.assertEqual(self.store.get_slot(1).media_state, "empty") + self.assertEqual(self.store.get_slot(2).url, "https://example.test/two") + + def test_delete_refuses_live_download_and_playback_guard(self): + reservation = self.store.reserve_download( + "https://example.test/live", None, None, job_id="live" + ) + manager = self.manager(FakeYtDlp(self.media)) + with self.assertRaises(SlotProtectedError): + manager.delete_slot(1, reason="manual") + self.store.fail_download("live", 1, reservation.generation, "cancelled") + self.store.protect_slot(1, "player_active", guard_id="active") + with self.assertRaises(SlotProtectedError): + manager.delete_slot(1, reason="manual") + + def test_symlink_slot_is_unlinked_without_touching_target(self): + reservation = self.store.reserve_download( + "https://example.test/link", None, None, job_id="link" + ) + self.store.fail_download("link", 1, reservation.generation, "failed") + target = Path(self.tempdir.name) / "outside" + target.mkdir() + protected = target / "keep.txt" + protected.write_text("keep", encoding="utf-8") + self.media.mkdir(parents=True) + os.symlink(target, self.media / "slot1") + manager = self.manager(FakeYtDlp(self.media)) + self.assertTrue(manager.delete_slot(1, reason="manual")) + self.assertEqual(protected.read_text(encoding="utf-8"), "keep") + + def test_interrupted_deletion_is_recoverable(self): + reservation = self.store.reserve_download( + "https://example.test/recover", None, None, job_id="recover" + ) + self.store.fail_download("recover", 1, reservation.generation, "failed") + token = self.store.begin_delete(1, reason="manual") + slot = self.media / "slot1" + slot.mkdir(parents=True) + (slot / "partial.mp4").write_bytes(b"partial") + self.assertEqual(self.store.pending_deletions(), [ + type(token)(slot=1, generation=token.generation, reason="recovered_interrupted_deletion") + ]) + manager = self.manager(FakeYtDlp(self.media)) + self.assertEqual(manager.recover_pending_deletions(), [1]) + self.assertFalse(slot.exists()) + self.assertEqual(self.store.get_slot(1).media_state, "empty") + + +if __name__ == "__main__": + unittest.main()