feat: reserve downloads and protect slot cleanup

This commit is contained in:
Media
2026-08-06 17:34:11 -05:00
parent 3ff9ba4637
commit 45083207b4
3 changed files with 611 additions and 0 deletions
+395
View File
@@ -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)
+17
View File
@@ -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] = []
+199
View File
@@ -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()