feat: reserve downloads and protect slot cleanup
This commit is contained in:
Media
@@ -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)
|
||||
@@ -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] = []
|
||||
|
||||
Reference in New Issue
Block a user