feat: add isolated utility export modes
This commit is contained in:
Media
+767
-3
@@ -2,14 +2,18 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import signal
|
||||
import shutil
|
||||
import stat
|
||||
import subprocess
|
||||
import tempfile
|
||||
import threading
|
||||
import time
|
||||
import unicodedata
|
||||
import uuid
|
||||
from collections.abc import Iterator
|
||||
from contextlib import contextmanager
|
||||
@@ -20,12 +24,31 @@ from typing import Any, Callable
|
||||
|
||||
from .config import ControllerConfig
|
||||
from .contracts import SLOT_IDS
|
||||
from .models import DeleteToken, DownloadReservation, SlotProtectedError, StaleOperationError
|
||||
from .models import (
|
||||
DeleteToken,
|
||||
DownloadReservation,
|
||||
ExportReservation,
|
||||
SlotProtectedError,
|
||||
StaleOperationError,
|
||||
UTILITY_EXPORT_MODES,
|
||||
)
|
||||
from .store import ControllerStore
|
||||
|
||||
|
||||
LOG = logging.getLogger(__name__)
|
||||
MEDIA_EXTENSIONS = frozenset({".mp4", ".mkv", ".webm", ".mov", ".m4v"})
|
||||
VIDEO_EXPORT_INPUT_EXTENSIONS = frozenset(
|
||||
{*MEDIA_EXTENSIONS, ".avi", ".flv", ".ts"}
|
||||
)
|
||||
EXPORT_AUDIO_EXTENSIONS = {"m4a": ".m4a", "mp3": ".mp3"}
|
||||
MAX_EXPORT_STEM_BYTES = 180
|
||||
EXPORT_MANIFEST_NAME = "manifest.json"
|
||||
_INVALID_FILENAME_CHARACTERS = frozenset('<>:"/\\|?*')
|
||||
_WINDOWS_RESERVED_NAMES = frozenset(
|
||||
{"CON", "PRN", "AUX", "NUL"}
|
||||
| {f"COM{index}" for index in range(1, 10)}
|
||||
| {f"LPT{index}" for index in range(1, 10)}
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
@@ -43,6 +66,16 @@ class DownloadEvent:
|
||||
reused_slot: bool = False
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MediaInspection:
|
||||
format_names: frozenset[str]
|
||||
video_codec: str | None
|
||||
audio_codec: str | None
|
||||
pixel_format: str | None
|
||||
width: int | None
|
||||
height: int | None
|
||||
|
||||
|
||||
EventCallback = Callable[[DownloadEvent], None]
|
||||
Runner = Callable[..., subprocess.CompletedProcess[str]]
|
||||
|
||||
@@ -63,6 +96,54 @@ def player_format_selector(max_height: int) -> str:
|
||||
return f"bv*[height<={max_height}]+ba/b[height<={max_height}]"
|
||||
|
||||
|
||||
def mp4_export_format_selector(max_height: int) -> str:
|
||||
"""Select quality first while making native AVC/AAC the equal-quality preference."""
|
||||
|
||||
return f"bv*[height<={max_height}]+ba/b[height<={max_height}]"
|
||||
|
||||
|
||||
def safe_export_filename(
|
||||
stem: str,
|
||||
extension: str,
|
||||
*,
|
||||
used_casefolded: set[str] | None = None,
|
||||
) -> str:
|
||||
"""Return a bounded Unicode filename that cannot introduce a path component."""
|
||||
|
||||
if extension not in {".mp4", ".m4a", ".mp3"}:
|
||||
raise ValueError("unsupported export filename extension")
|
||||
normalized = unicodedata.normalize("NFC", str(stem))
|
||||
cleaned = "".join(
|
||||
"_"
|
||||
if character in _INVALID_FILENAME_CHARACTERS
|
||||
or ord(character) < 32
|
||||
or ord(character) == 127
|
||||
else character
|
||||
for character in normalized
|
||||
)
|
||||
cleaned = re.sub(r"\s+", " ", cleaned).strip(" .") or "untitled"
|
||||
if cleaned.split(".", 1)[0].upper() in _WINDOWS_RESERVED_NAMES:
|
||||
cleaned = f"_{cleaned}"
|
||||
|
||||
def bounded(value: str, suffix: str) -> str:
|
||||
available = MAX_EXPORT_STEM_BYTES - len(suffix.encode("utf-8"))
|
||||
candidate = value
|
||||
while candidate and len(candidate.encode("utf-8")) > available:
|
||||
candidate = candidate[:-1]
|
||||
return (candidate.rstrip(" .") or "untitled") + suffix
|
||||
|
||||
used = used_casefolded if used_casefolded is not None else set()
|
||||
counter = 1
|
||||
while True:
|
||||
collision_suffix = "" if counter == 1 else f" ({counter})"
|
||||
filename = bounded(cleaned, collision_suffix) + extension
|
||||
folded = filename.casefold()
|
||||
if folded not in used:
|
||||
used.add(folded)
|
||||
return filename
|
||||
counter += 1
|
||||
|
||||
|
||||
def title_from_filename(path: Path) -> str:
|
||||
stem = path.stem
|
||||
if " - " in stem:
|
||||
@@ -83,6 +164,7 @@ class DownloadManager:
|
||||
on_event: EventCallback | None = None,
|
||||
runner: Runner | None = None,
|
||||
executor: ThreadPoolExecutor | None = None,
|
||||
export_executor: ThreadPoolExecutor | None = None,
|
||||
):
|
||||
self.config = config
|
||||
self.store = store
|
||||
@@ -93,6 +175,19 @@ class DownloadManager:
|
||||
thread_name_prefix="yt-download",
|
||||
)
|
||||
self._owns_executor = executor is None
|
||||
if export_executor is not None:
|
||||
self._export_executor = export_executor
|
||||
self._owns_export_executor = False
|
||||
elif executor is not None:
|
||||
# Deterministic callers/tests may inject one executor explicitly.
|
||||
self._export_executor = executor
|
||||
self._owns_export_executor = False
|
||||
else:
|
||||
self._export_executor = ThreadPoolExecutor(
|
||||
max_workers=config.downloads.export_max_workers,
|
||||
thread_name_prefix="yt-export",
|
||||
)
|
||||
self._owns_export_executor = True
|
||||
self._lock = threading.Lock()
|
||||
self._futures: set[Future[Any]] = set()
|
||||
self._processes: set[subprocess.Popen[str]] = set()
|
||||
@@ -167,6 +262,49 @@ class DownloadManager:
|
||||
future.add_done_callback(self._forget_future)
|
||||
return reservation
|
||||
|
||||
def enqueue_export(
|
||||
self,
|
||||
url: str,
|
||||
mode: str,
|
||||
submitter: str | None,
|
||||
note: str | None,
|
||||
) -> ExportReservation:
|
||||
"""Reserve and queue a utility export without touching Player slots."""
|
||||
|
||||
if mode not in UTILITY_EXPORT_MODES:
|
||||
raise ValueError("unsupported utility export mode")
|
||||
with self._lock:
|
||||
if self._stopping:
|
||||
return ExportReservation(status="shutting_down")
|
||||
reservation = self.store.reserve_export_job(
|
||||
url,
|
||||
mode,
|
||||
submitter,
|
||||
note,
|
||||
max_active=self.config.downloads.max_queued_exports,
|
||||
)
|
||||
if reservation.status != "queued":
|
||||
return reservation
|
||||
assert reservation.job_id is not None
|
||||
try:
|
||||
future = self._export_executor.submit(
|
||||
self._export,
|
||||
url,
|
||||
mode,
|
||||
reservation.job_id,
|
||||
)
|
||||
except BaseException as exc:
|
||||
message = f"could not start export worker: {exc}"
|
||||
self.store.fail_export_job(reservation.job_id, message)
|
||||
return ExportReservation(
|
||||
status="worker_error",
|
||||
job_id=reservation.job_id,
|
||||
)
|
||||
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)
|
||||
@@ -374,12 +512,636 @@ class DownloadManager:
|
||||
)
|
||||
)
|
||||
|
||||
def _export(
|
||||
self,
|
||||
url: str,
|
||||
mode: str,
|
||||
job_id: str,
|
||||
) -> None:
|
||||
if not self.store.mark_export_running(job_id):
|
||||
return
|
||||
work_directory: Path | None = None
|
||||
published = False
|
||||
completed = False
|
||||
try:
|
||||
work_directory, source_directory, ready_directory, final_parent = (
|
||||
self._prepare_export_directories(job_id, mode)
|
||||
)
|
||||
with self._yt_dlp_arguments() as base_arguments:
|
||||
arguments = base_arguments + self._export_download_arguments(
|
||||
mode,
|
||||
source_directory,
|
||||
url,
|
||||
)
|
||||
self._run_command(
|
||||
arguments,
|
||||
check=True,
|
||||
text=True,
|
||||
timeout=self.config.downloads.process_timeout_seconds,
|
||||
)
|
||||
|
||||
sources = self._export_source_files(source_directory, mode)
|
||||
if not sources:
|
||||
raise RuntimeError("yt-dlp completed without producing export media")
|
||||
used_names: set[str] = set()
|
||||
ready_outputs: list[Path] = []
|
||||
for source in sources:
|
||||
extension = ".mp4" if mode == "mp4" else EXPORT_AUDIO_EXTENSIONS[mode]
|
||||
output = ready_directory / safe_export_filename(
|
||||
source.stem,
|
||||
extension,
|
||||
used_casefolded=used_names,
|
||||
)
|
||||
if mode == "mp4":
|
||||
self._produce_compatible_mp4(source, output)
|
||||
else:
|
||||
inspection = self._inspect_media(source)
|
||||
self._validate_audio_export(inspection, mode)
|
||||
os.replace(source, output)
|
||||
self._fsync_file(output)
|
||||
ready_outputs.append(output)
|
||||
|
||||
self._fsync_directory(ready_directory)
|
||||
self._write_export_manifest(
|
||||
work_directory,
|
||||
job_id=job_id,
|
||||
mode=mode,
|
||||
filenames=[output.name for output in ready_outputs],
|
||||
)
|
||||
published_directory = final_parent / job_id
|
||||
if os.path.lexists(published_directory):
|
||||
raise FileExistsError(
|
||||
f"refusing to overwrite existing export directory: {published_directory}"
|
||||
)
|
||||
os.rename(ready_directory, published_directory)
|
||||
self._fsync_directory(final_parent)
|
||||
published = True
|
||||
media_base = self.config.storage.media_base.resolve()
|
||||
output_paths = tuple(
|
||||
str((published_directory / output.name).resolve().relative_to(media_base))
|
||||
for output in ready_outputs
|
||||
)
|
||||
if not self.store.complete_export_job(job_id, output_paths):
|
||||
raise StaleOperationError(
|
||||
"published export could not transition its durable job"
|
||||
)
|
||||
completed = True
|
||||
except BaseException as exc:
|
||||
message = self._failure_message(exc)
|
||||
try:
|
||||
self.store.fail_export_job(job_id, message)
|
||||
except BaseException:
|
||||
LOG.exception("could not persist utility export failure for %s", job_id)
|
||||
if published:
|
||||
LOG.error(
|
||||
"utility export %s was published but needs restart reconciliation: %s",
|
||||
job_id,
|
||||
message,
|
||||
)
|
||||
else:
|
||||
LOG.warning("utility export %s failed: %s", job_id, message)
|
||||
finally:
|
||||
if work_directory is not None and (not published or completed):
|
||||
self._remove_export_work_directory(work_directory)
|
||||
|
||||
def _export_download_arguments(
|
||||
self,
|
||||
mode: str,
|
||||
source_directory: Path,
|
||||
url: str,
|
||||
) -> list[str]:
|
||||
arguments = [
|
||||
"--yes-playlist",
|
||||
"--no-overwrites",
|
||||
"--no-progress",
|
||||
"--windows-filenames",
|
||||
"--trim-filenames",
|
||||
str(MAX_EXPORT_STEM_BYTES),
|
||||
"-o",
|
||||
str(
|
||||
source_directory
|
||||
/ "%(playlist_index)03d - %(title)s [%(id)s].%(ext)s"
|
||||
),
|
||||
]
|
||||
if mode == "mp4":
|
||||
arguments.extend(
|
||||
(
|
||||
"-f",
|
||||
mp4_export_format_selector(
|
||||
self.config.downloads.export_max_height
|
||||
),
|
||||
"-S",
|
||||
(
|
||||
f"res:{self.config.downloads.export_max_height},"
|
||||
"vcodec:h264,acodec:aac"
|
||||
),
|
||||
)
|
||||
)
|
||||
elif mode == "m4a":
|
||||
arguments.extend(
|
||||
(
|
||||
"-f",
|
||||
"ba/b",
|
||||
"-S",
|
||||
"abr,acodec:aac",
|
||||
"-x",
|
||||
"--audio-format",
|
||||
"m4a",
|
||||
"--audio-quality",
|
||||
"0",
|
||||
)
|
||||
)
|
||||
elif mode == "mp3":
|
||||
arguments.extend(
|
||||
(
|
||||
"-f",
|
||||
"ba/b",
|
||||
"-x",
|
||||
"--audio-format",
|
||||
"mp3",
|
||||
"--audio-quality",
|
||||
"0",
|
||||
)
|
||||
)
|
||||
else:
|
||||
raise ValueError("unsupported utility export mode")
|
||||
arguments.append(url)
|
||||
return arguments
|
||||
|
||||
def _prepare_export_directories(
|
||||
self,
|
||||
job_id: str,
|
||||
mode: str,
|
||||
) -> tuple[Path, Path, Path, Path]:
|
||||
if not re.fullmatch(r"[0-9a-f]{32}", job_id):
|
||||
raise ValueError("invalid export job identity")
|
||||
media_base = self.config.storage.media_base
|
||||
media_base.mkdir(parents=True, exist_ok=True)
|
||||
export_root = media_base
|
||||
for part in self.config.downloads.export_directory.parts:
|
||||
export_root = export_root / part
|
||||
self._ensure_real_directory(export_root)
|
||||
work_root = export_root / ".work"
|
||||
final_parent = export_root / ("video" if mode == "mp4" else "audio")
|
||||
self._ensure_real_directory(work_root)
|
||||
self._ensure_real_directory(final_parent)
|
||||
work_directory = work_root / job_id
|
||||
work_directory.mkdir(mode=0o750, exist_ok=False)
|
||||
source_directory = work_directory / "source"
|
||||
ready_directory = work_directory / "ready"
|
||||
source_directory.mkdir(mode=0o750)
|
||||
ready_directory.mkdir(mode=0o750)
|
||||
self._fsync_directory(work_root)
|
||||
return work_directory, source_directory, ready_directory, final_parent
|
||||
|
||||
@staticmethod
|
||||
def _ensure_real_directory(path: Path) -> None:
|
||||
try:
|
||||
path.mkdir(mode=0o750)
|
||||
except FileExistsError:
|
||||
pass
|
||||
metadata = path.lstat()
|
||||
if stat.S_ISLNK(metadata.st_mode) or not stat.S_ISDIR(metadata.st_mode):
|
||||
raise RuntimeError(f"refusing unsafe export directory: {path}")
|
||||
|
||||
def _write_export_manifest(
|
||||
self,
|
||||
work_directory: Path,
|
||||
*,
|
||||
job_id: str,
|
||||
mode: str,
|
||||
filenames: list[str],
|
||||
) -> None:
|
||||
if not filenames:
|
||||
raise ValueError("export manifest must contain output filenames")
|
||||
payload = json.dumps(
|
||||
{
|
||||
"version": 1,
|
||||
"job_id": job_id,
|
||||
"mode": mode,
|
||||
"filenames": filenames,
|
||||
},
|
||||
ensure_ascii=False,
|
||||
sort_keys=True,
|
||||
)
|
||||
temporary = work_directory / f".{EXPORT_MANIFEST_NAME}.tmp"
|
||||
manifest = work_directory / EXPORT_MANIFEST_NAME
|
||||
try:
|
||||
with temporary.open("x", encoding="utf-8", newline="\n") as stream:
|
||||
stream.write(payload + "\n")
|
||||
stream.flush()
|
||||
os.fsync(stream.fileno())
|
||||
os.replace(temporary, manifest)
|
||||
self._fsync_directory(work_directory)
|
||||
finally:
|
||||
temporary.unlink(missing_ok=True)
|
||||
|
||||
@staticmethod
|
||||
def _export_source_files(source_directory: Path, mode: str) -> list[Path]:
|
||||
allowed = (
|
||||
VIDEO_EXPORT_INPUT_EXTENSIONS
|
||||
if mode == "mp4"
|
||||
else frozenset({EXPORT_AUDIO_EXTENSIONS[mode]})
|
||||
)
|
||||
media: list[Path] = []
|
||||
for candidate in source_directory.iterdir():
|
||||
if candidate.is_symlink() or not candidate.is_file():
|
||||
raise RuntimeError(
|
||||
f"yt-dlp produced an unsafe export entry: {candidate.name}"
|
||||
)
|
||||
if candidate.suffix.lower() not in allowed:
|
||||
raise RuntimeError(
|
||||
f"yt-dlp produced an unexpected export file: {candidate.name}"
|
||||
)
|
||||
media.append(candidate)
|
||||
return sorted(media, key=lambda path: path.name.casefold())
|
||||
|
||||
def _inspect_media(self, path: Path) -> MediaInspection:
|
||||
result = self._run_command(
|
||||
[
|
||||
"ffprobe",
|
||||
"-v",
|
||||
"error",
|
||||
"-show_streams",
|
||||
"-show_format",
|
||||
"-of",
|
||||
"json",
|
||||
str(path),
|
||||
],
|
||||
check=True,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=min(120.0, self.config.downloads.process_timeout_seconds),
|
||||
)
|
||||
try:
|
||||
payload = json.loads(result.stdout or "")
|
||||
streams = payload["streams"]
|
||||
format_name = payload["format"]["format_name"]
|
||||
except (KeyError, TypeError, ValueError, json.JSONDecodeError) as exc:
|
||||
raise RuntimeError("ffprobe returned invalid media metadata") from exc
|
||||
if not isinstance(streams, list) or not isinstance(format_name, str):
|
||||
raise RuntimeError("ffprobe returned invalid media metadata")
|
||||
video = next(
|
||||
(
|
||||
stream
|
||||
for stream in streams
|
||||
if isinstance(stream, dict) and stream.get("codec_type") == "video"
|
||||
),
|
||||
None,
|
||||
)
|
||||
audio = next(
|
||||
(
|
||||
stream
|
||||
for stream in streams
|
||||
if isinstance(stream, dict) and stream.get("codec_type") == "audio"
|
||||
),
|
||||
None,
|
||||
)
|
||||
|
||||
def positive_integer(value: Any) -> int | None:
|
||||
return (
|
||||
value
|
||||
if isinstance(value, int)
|
||||
and not isinstance(value, bool)
|
||||
and value > 0
|
||||
else None
|
||||
)
|
||||
|
||||
return MediaInspection(
|
||||
format_names=frozenset(
|
||||
value.strip().lower()
|
||||
for value in format_name.split(",")
|
||||
if value.strip()
|
||||
),
|
||||
video_codec=(
|
||||
video.get("codec_name") if isinstance(video, dict) else None
|
||||
),
|
||||
audio_codec=(
|
||||
audio.get("codec_name") if isinstance(audio, dict) else None
|
||||
),
|
||||
pixel_format=(
|
||||
video.get("pix_fmt") if isinstance(video, dict) else None
|
||||
),
|
||||
width=positive_integer(video.get("width") if isinstance(video, dict) else None),
|
||||
height=positive_integer(video.get("height") if isinstance(video, dict) else None),
|
||||
)
|
||||
|
||||
def _produce_compatible_mp4(self, source: Path, output: Path) -> None:
|
||||
source_info = self._inspect_media(source)
|
||||
if source_info.video_codec is None or source_info.audio_codec is None:
|
||||
raise RuntimeError("MP4 export requires both video and audio")
|
||||
if source_info.height is None:
|
||||
raise RuntimeError("MP4 export source has no verifiable video height")
|
||||
if source_info.height > self.config.downloads.export_max_height:
|
||||
raise RuntimeError("MP4 export source exceeds configured maximum height")
|
||||
video_compatible = (
|
||||
source_info.video_codec == "h264"
|
||||
and source_info.pixel_format in {"yuv420p", "yuvj420p"}
|
||||
)
|
||||
audio_compatible = source_info.audio_codec == "aac"
|
||||
container_compatible = "mp4" in source_info.format_names
|
||||
if video_compatible and audio_compatible and container_compatible:
|
||||
os.replace(source, output)
|
||||
else:
|
||||
temporary = output.with_name(
|
||||
f".{output.stem}.{uuid.uuid4().hex}.tmp.mp4"
|
||||
)
|
||||
arguments = [
|
||||
"ffmpeg",
|
||||
"-nostdin",
|
||||
"-hide_banner",
|
||||
"-loglevel",
|
||||
"error",
|
||||
"-n",
|
||||
"-i",
|
||||
str(source),
|
||||
"-map",
|
||||
"0:v:0",
|
||||
"-map",
|
||||
"0:a:0",
|
||||
"-map_metadata",
|
||||
"0",
|
||||
]
|
||||
if video_compatible:
|
||||
arguments.extend(("-c:v", "copy"))
|
||||
else:
|
||||
arguments.extend(
|
||||
(
|
||||
"-c:v",
|
||||
"libx264",
|
||||
"-preset",
|
||||
"medium",
|
||||
"-crf",
|
||||
"18",
|
||||
"-pix_fmt",
|
||||
"yuv420p",
|
||||
)
|
||||
)
|
||||
if audio_compatible:
|
||||
arguments.extend(("-c:a", "copy"))
|
||||
else:
|
||||
arguments.extend(("-c:a", "aac", "-b:a", "192k"))
|
||||
arguments.extend(("-movflags", "+faststart", str(temporary)))
|
||||
try:
|
||||
self._run_command(
|
||||
arguments,
|
||||
check=True,
|
||||
text=True,
|
||||
timeout=self.config.downloads.process_timeout_seconds,
|
||||
)
|
||||
converted = self._inspect_media(temporary)
|
||||
self._validate_compatible_mp4(converted, source_info.height)
|
||||
os.replace(temporary, output)
|
||||
finally:
|
||||
temporary.unlink(missing_ok=True)
|
||||
self._validate_compatible_mp4(
|
||||
self._inspect_media(output),
|
||||
source_info.height,
|
||||
)
|
||||
|
||||
def _validate_compatible_mp4(
|
||||
self,
|
||||
inspection: MediaInspection,
|
||||
source_height: int,
|
||||
) -> None:
|
||||
if (
|
||||
"mp4" not in inspection.format_names
|
||||
or inspection.video_codec != "h264"
|
||||
or inspection.audio_codec != "aac"
|
||||
or inspection.pixel_format not in {"yuv420p", "yuvj420p"}
|
||||
):
|
||||
raise RuntimeError("MP4 compatibility verification failed")
|
||||
if (
|
||||
inspection.height != source_height
|
||||
or inspection.height > self.config.downloads.export_max_height
|
||||
):
|
||||
raise RuntimeError("MP4 export resolution verification failed")
|
||||
|
||||
@staticmethod
|
||||
def _validate_audio_export(
|
||||
inspection: MediaInspection,
|
||||
mode: str,
|
||||
) -> None:
|
||||
expected_codec = "aac" if mode == "m4a" else "mp3"
|
||||
expected_container = "mp4" if mode == "m4a" else "mp3"
|
||||
if inspection.video_codec is not None:
|
||||
raise RuntimeError("audio export unexpectedly contains video")
|
||||
if inspection.audio_codec != expected_codec:
|
||||
raise RuntimeError(f"{mode.upper()} export codec verification failed")
|
||||
if expected_container not in inspection.format_names:
|
||||
raise RuntimeError(f"{mode.upper()} export container verification failed")
|
||||
|
||||
@staticmethod
|
||||
def _fsync_file(path: Path) -> None:
|
||||
descriptor = os.open(path, os.O_RDONLY)
|
||||
try:
|
||||
os.fsync(descriptor)
|
||||
finally:
|
||||
os.close(descriptor)
|
||||
|
||||
def _remove_export_work_directory(self, path: Path) -> None:
|
||||
try:
|
||||
parent = path.parent
|
||||
if os.path.lexists(path):
|
||||
if path.is_symlink() or path.is_file():
|
||||
path.unlink()
|
||||
else:
|
||||
shutil.rmtree(path)
|
||||
self._fsync_directory(parent)
|
||||
except OSError:
|
||||
LOG.exception("could not remove export work directory %s", path)
|
||||
|
||||
def reconcile_export_artifacts(self) -> None:
|
||||
"""Recover verified crash-published exports and discard incomplete staging."""
|
||||
|
||||
export_root = self.config.storage.media_base / self.config.downloads.export_directory
|
||||
if not os.path.lexists(export_root):
|
||||
return
|
||||
try:
|
||||
self._verify_existing_export_tree(export_root)
|
||||
except OSError as exc:
|
||||
LOG.error("could not inspect utility export tree: %s", exc)
|
||||
return
|
||||
except RuntimeError as exc:
|
||||
LOG.error("utility export reconciliation refused unsafe storage: %s", exc)
|
||||
return
|
||||
|
||||
jobs = {
|
||||
job.job_id: job for job in self.store.list_reconcilable_export_jobs()
|
||||
}
|
||||
for job in jobs.values():
|
||||
try:
|
||||
self._reconcile_export_job(export_root, job.job_id, job.mode)
|
||||
except BaseException:
|
||||
LOG.exception("could not reconcile utility export %s", job.job_id)
|
||||
|
||||
work_root = export_root / ".work"
|
||||
if not work_root.is_dir() or work_root.is_symlink():
|
||||
return
|
||||
for candidate in work_root.iterdir():
|
||||
if not re.fullmatch(r"[0-9a-f]{32}", candidate.name):
|
||||
LOG.warning("leaving unknown utility export staging entry %s", candidate)
|
||||
continue
|
||||
job = self.store.get_export_job(candidate.name)
|
||||
if job is None:
|
||||
LOG.warning("leaving untracked utility export staging entry %s", candidate)
|
||||
continue
|
||||
if job.status == "completed":
|
||||
self._remove_export_work_directory(candidate)
|
||||
elif job.job_id not in jobs:
|
||||
LOG.warning(
|
||||
"leaving utility export staging for active job %s",
|
||||
candidate.name,
|
||||
)
|
||||
|
||||
def _verify_existing_export_tree(self, export_root: Path) -> None:
|
||||
current = self.config.storage.media_base
|
||||
for part in self.config.downloads.export_directory.parts:
|
||||
current = current / part
|
||||
metadata = current.lstat()
|
||||
if stat.S_ISLNK(metadata.st_mode) or not stat.S_ISDIR(metadata.st_mode):
|
||||
raise RuntimeError(f"unsafe export directory: {current}")
|
||||
for child_name in (".work", "video", "audio"):
|
||||
child = export_root / child_name
|
||||
if not os.path.lexists(child):
|
||||
continue
|
||||
metadata = child.lstat()
|
||||
if stat.S_ISLNK(metadata.st_mode) or not stat.S_ISDIR(metadata.st_mode):
|
||||
raise RuntimeError(f"unsafe export directory: {child}")
|
||||
|
||||
def _reconcile_export_job(
|
||||
self,
|
||||
export_root: Path,
|
||||
job_id: str,
|
||||
mode: str,
|
||||
) -> None:
|
||||
work_directory = export_root / ".work" / job_id
|
||||
final_parent = export_root / ("video" if mode == "mp4" else "audio")
|
||||
final_directory = final_parent / job_id
|
||||
if os.path.lexists(work_directory):
|
||||
work_metadata = work_directory.lstat()
|
||||
if stat.S_ISLNK(work_metadata.st_mode) or not stat.S_ISDIR(
|
||||
work_metadata.st_mode
|
||||
):
|
||||
raise RuntimeError("utility export staging directory is unsafe")
|
||||
manifest = self._read_export_manifest(work_directory, job_id, mode)
|
||||
|
||||
if os.path.lexists(final_directory):
|
||||
filenames = None if manifest is None else manifest
|
||||
outputs = self._validate_export_output_directory(
|
||||
final_directory,
|
||||
mode,
|
||||
expected_filenames=filenames,
|
||||
)
|
||||
elif manifest is not None:
|
||||
ready_directory = work_directory / "ready"
|
||||
outputs = self._validate_export_output_directory(
|
||||
ready_directory,
|
||||
mode,
|
||||
expected_filenames=manifest,
|
||||
)
|
||||
self._ensure_real_directory(final_parent)
|
||||
os.rename(ready_directory, final_directory)
|
||||
self._fsync_directory(final_parent)
|
||||
else:
|
||||
if os.path.lexists(work_directory):
|
||||
self._remove_export_work_directory(work_directory)
|
||||
return
|
||||
|
||||
media_base = self.config.storage.media_base.resolve()
|
||||
output_paths = tuple(
|
||||
str((final_directory / output.name).resolve().relative_to(media_base))
|
||||
for output in outputs
|
||||
)
|
||||
if self.store.recover_export_job(job_id, output_paths):
|
||||
LOG.warning("recovered crash-published utility export %s", job_id)
|
||||
if os.path.lexists(work_directory):
|
||||
self._remove_export_work_directory(work_directory)
|
||||
|
||||
@staticmethod
|
||||
def _read_export_manifest(
|
||||
work_directory: Path,
|
||||
job_id: str,
|
||||
mode: str,
|
||||
) -> list[str] | None:
|
||||
manifest = work_directory / EXPORT_MANIFEST_NAME
|
||||
if not manifest.is_file() or manifest.is_symlink():
|
||||
return None
|
||||
if manifest.stat().st_size > 8 * 1024 * 1024:
|
||||
raise RuntimeError("utility export manifest is too large")
|
||||
try:
|
||||
payload = json.loads(manifest.read_text(encoding="utf-8"))
|
||||
except (OSError, UnicodeError, json.JSONDecodeError) as exc:
|
||||
raise RuntimeError("utility export manifest is invalid") from exc
|
||||
if (
|
||||
not isinstance(payload, dict)
|
||||
or payload.get("version") != 1
|
||||
or payload.get("job_id") != job_id
|
||||
or payload.get("mode") != mode
|
||||
):
|
||||
raise RuntimeError("utility export manifest identity is invalid")
|
||||
filenames = payload.get("filenames")
|
||||
if not isinstance(filenames, list) or not filenames or len(filenames) > 10000:
|
||||
raise RuntimeError("utility export manifest filenames are invalid")
|
||||
expected_extension = ".mp4" if mode == "mp4" else EXPORT_AUDIO_EXTENSIONS[mode]
|
||||
normalized: list[str] = []
|
||||
for filename in filenames:
|
||||
if (
|
||||
not isinstance(filename, str)
|
||||
or not filename
|
||||
or Path(filename).name != filename
|
||||
or Path(filename).suffix.lower() != expected_extension
|
||||
):
|
||||
raise RuntimeError("utility export manifest contains an unsafe filename")
|
||||
normalized.append(filename)
|
||||
if len({value.casefold() for value in normalized}) != len(normalized):
|
||||
raise RuntimeError("utility export manifest contains duplicate filenames")
|
||||
return normalized
|
||||
|
||||
def _validate_export_output_directory(
|
||||
self,
|
||||
directory: Path,
|
||||
mode: str,
|
||||
*,
|
||||
expected_filenames: list[str] | None,
|
||||
) -> list[Path]:
|
||||
metadata = directory.lstat()
|
||||
if stat.S_ISLNK(metadata.st_mode) or not stat.S_ISDIR(metadata.st_mode):
|
||||
raise RuntimeError("utility export output directory is unsafe")
|
||||
outputs = sorted(directory.iterdir(), key=lambda path: path.name.casefold())
|
||||
if not outputs:
|
||||
raise RuntimeError("utility export output directory is empty")
|
||||
if expected_filenames is not None and {
|
||||
path.name for path in outputs
|
||||
} != set(expected_filenames):
|
||||
raise RuntimeError("utility export output does not match its manifest")
|
||||
expected_extension = ".mp4" if mode == "mp4" else EXPORT_AUDIO_EXTENSIONS[mode]
|
||||
for output in outputs:
|
||||
if (
|
||||
output.is_symlink()
|
||||
or not output.is_file()
|
||||
or output.suffix.lower() != expected_extension
|
||||
):
|
||||
raise RuntimeError("utility export output contains an unsafe entry")
|
||||
inspection = self._inspect_media(output)
|
||||
if mode == "mp4":
|
||||
if inspection.height is None:
|
||||
raise RuntimeError("MP4 export output has no verifiable height")
|
||||
self._validate_compatible_mp4(inspection, inspection.height)
|
||||
else:
|
||||
self._validate_audio_export(inspection, mode)
|
||||
return outputs
|
||||
|
||||
@staticmethod
|
||||
def _failure_message(exc: BaseException) -> str:
|
||||
command = getattr(exc, "cmd", None)
|
||||
program = "media command"
|
||||
if isinstance(command, (list, tuple)) and command:
|
||||
program = Path(str(command[0])).name
|
||||
if isinstance(exc, subprocess.TimeoutExpired):
|
||||
return f"yt-dlp timed out after {exc.timeout} seconds"
|
||||
return f"{program} timed out after {exc.timeout} seconds"
|
||||
if isinstance(exc, subprocess.CalledProcessError):
|
||||
return f"yt-dlp failed with exit code {exc.returncode}"
|
||||
return f"{program} failed with exit code {exc.returncode}"
|
||||
return f"download failed: {exc}"
|
||||
|
||||
def _run_command(
|
||||
@@ -708,6 +1470,8 @@ class DownloadManager:
|
||||
self._terminate_process(process, force=True)
|
||||
if self._owns_executor:
|
||||
self._executor.shutdown(wait=wait, cancel_futures=True)
|
||||
if self._owns_export_executor:
|
||||
self._export_executor.shutdown(wait=wait, cancel_futures=True)
|
||||
|
||||
def active_future_count(self) -> int:
|
||||
with self._lock:
|
||||
|
||||
Reference in New Issue
Block a user