import asyncio import os import subprocess import sys import tempfile from pathlib import Path from secrets import token_hex from typing import TYPE_CHECKING, ClassVar from aiohttp import web from app.features.streaming.library.ffprobe import ffprobe from app.features.streaming.library.segment_encoders import ( detect_qsv_capabilities, encoder_fallback_chain, get_builder_for_codec, has_dri_devices, select_encoder, ) from app.library.config import SUPPORTED_CODECS, Config from app.library.log import get_logger if TYPE_CHECKING: from asyncio.subprocess import Process from .ffprobe import FFProbeResult from .segment_encoders import EncoderBuilder LOG = get_logger() class Segments: """ Build and stream MPEG-TS segments using ffmpeg with optional on-the-fly transcoding. The class can auto-detect an available hardware encoder on the first video segment and cache that choice for subsequent segments to avoid repeated probing. """ # Cache of selected video encoder across all segment instances _cached_vcodec: ClassVar[str | None] = None _cache_initialized: ClassVar[bool] = False def __init__(self, download_path: str, index: int, duration: float, vconvert: bool, aconvert: bool): config: Config = Config.get_instance() self.download_path: str = download_path "The path where files are downloaded." self.index = int(index) "The index of the segment." self.duration = float(duration) "The duration of the segment." self.vconvert = bool(vconvert) "Whether to convert video." self.aconvert = bool(aconvert) "Whether to convert audio." # Default to configured codec; can be replaced by auto-detection self.vcodec: str = config.streamer_vcodec "The video codec to use." self.acodec: str = config.streamer_acodec "The audio codec to use." # sadly due to unforeseen circumstances, we have to convert the video for now. self.vconvert = True "Whether to convert video." self.aconvert = True "Whether to convert audio." self.attempted: set[str] = set() "The set of attempted codecs." def _make_stream_input(self, file: Path) -> Path: while True: stream_input = Path(tempfile.gettempdir()).joinpath(f"ytptube_stream.{token_hex(8)}") try: stream_input.symlink_to(file, target_is_directory=False) return stream_input except FileExistsError: continue async def build_ffmpeg_args(self, file: Path, s_codec: str, *, stream_input: Path | None = None) -> list[str]: try: ff: FFProbeResult = await ffprobe(file) except UnicodeDecodeError: pass input_path = stream_input or file startTime: str = f"{0:.6f}" if self.index == 0 else f"{self.duration * self.index:.6f}" ctx = { "is_linux": sys.platform.startswith("linux"), "has_dri": has_dri_devices(), "vaapi_device": Config.get_instance().vaapi_device, "qsv": {"full": False, "lp": False}, } caps: dict[str, dict[str, bool]] = detect_qsv_capabilities() base_codec: str = s_codec.split("_", maxsplit=1)[0] codec_caps: dict[str, bool] = caps.get(base_codec, {"full": False, "lp": False}) ctx["qsv"] = codec_caps if self.vconvert and ff and hasattr(ff, "has_video") and ff.has_video(): builder: EncoderBuilder = get_builder_for_codec(s_codec) else: builder = None ctx: dict = {} # Collect encoder-specific input/global flags that must precede the input input_args: list[str] = [] if builder: input_args = builder.input_args(ctx) fargs: list[str] = [ "-xerror", "-hide_banner", "-loglevel", "error", # input trimming before the input file "-ss", str(startTime), "-t", str(f"{self.duration:.6f}"), "-copyts", # hardware/global input options must come before -i *input_args, "-i", f"file:{input_path}", "-map_metadata", "-1", ] v_args: list[str] = [] if builder: v_args = builder.add_video_args(["-g", "52", "-map", "0:v:0", "-strict", "-2"], ctx) else: v_args += ["-codec:v", "copy"] fargs += v_args if ff and ff.has_audio(): fargs += ["-map", "0:a:0", "-codec:a", self.acodec if self.aconvert else "copy"] fargs += ["-sn", "-muxdelay", "0", "-f", "mpegts", "pipe:1"] return fargs async def _run(self, resp: web.StreamResponse, file: Path, args: list[str]) -> tuple[bool, int, bool, str]: proc: Process = await asyncio.create_subprocess_exec( "ffmpeg", *args, stdin=asyncio.subprocess.DEVNULL, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, creationflags=subprocess.CREATE_NO_WINDOW if os.name == "nt" else 0, ) client_disconnected_local = False wrote_any = False stderr_buf = bytearray() async def _drain_stderr() -> None: try: assert proc.stderr is not None while True: chunk: bytes = await proc.stderr.read(4096) if not chunk: break stderr_buf.extend(chunk) except Exception: # best-effort only pass stderr_task: asyncio.Task[None] = asyncio.create_task(_drain_stderr()) LOG.debug( "Streaming segment %s for '%s'.", self.index, file, extra={"file": str(file), "segment_index": self.index, "ffmpeg_args": args}, ) try: assert proc.stdout is not None while True: chunk: bytes = await proc.stdout.read(1024 * 64) if not chunk: break wrote_any = True try: await resp.write(chunk) except (asyncio.CancelledError, ConnectionResetError, BrokenPipeError, ConnectionError): LOG.warning("Client disconnected or connection reset while writing.") client_disconnected_local = True break except asyncio.CancelledError: LOG.warning("Client disconnected. Terminating ffmpeg.") client_disconnected_local = True proc.terminate() try: await asyncio.wait_for(proc.wait(), timeout=5) except TimeoutError: LOG.warning( "Segment stream for '%s' did not stop ffmpeg in time after the client disconnected; killing it.", file, ) proc.kill() raise except ConnectionResetError: LOG.warning("Connection reset by peer. Skipping further writes.") client_disconnected_local = True finally: # Ensure process is stopped when client disconnected, to avoid hangs if client_disconnected_local: proc.terminate() try: rc: int = await asyncio.wait_for(proc.wait(), timeout=5) except TimeoutError: LOG.warning( "Segment stream for '%s' did not stop ffmpeg in time after the client disconnected; killing it.", file, ) proc.kill() try: rc = await asyncio.wait_for(proc.wait(), timeout=5) except Exception: rc = -1 else: # Normal termination path with a safety timeout try: rc = await asyncio.wait_for(proc.wait(), timeout=30) except TimeoutError: LOG.warning("Segment stream for '%s' did not stop ffmpeg in time; killing the process.", file) proc.kill() try: rc = await asyncio.wait_for(proc.wait(), timeout=5) except Exception: rc = -1 try: await asyncio.wait_for(stderr_task, timeout=1) except Exception: pass return (wrote_any, rc, client_disconnected_local, stderr_buf.decode("utf-8", errors="ignore").strip()) async def stream(self, file: Path, resp: web.StreamResponse) -> None: codec: str = self.vcodec if Segments._cache_initialized and Segments._cached_vcodec: codec = Segments._cached_vcodec if not codec or codec not in SUPPORTED_CODECS: codec: str = select_encoder(self.vcodec or "") LOG.debug("Selected video codec '%s' for segment streaming.", codec, extra={"codec": codec}) if Segments._cached_vcodec and Segments._cache_initialized: codecs: list[str] = [Segments._cached_vcodec] else: codecs: list[str] = [codec, *list(encoder_fallback_chain(codec))] stream_input = self._make_stream_input(file) try: for s_codec in codecs: if s_codec in self.attempted: continue ffmpeg_args: list[str] = await self.build_ffmpeg_args(file, s_codec, stream_input=stream_input) _, rc, client_disconnected, stderr_text = await self._run(resp, file, ffmpeg_args) if 0 == rc: Segments._cached_vcodec = s_codec Segments._cache_initialized = True return if client_disconnected: return if 0 != rc: err: str = stderr_text[:500] if stderr_text else "no error output" LOG.warning( "Retrying segment %s for '%s' because hardware encoder '%s' failed: %s.", self.index, file, s_codec, err, extra={"ffmpeg_args": ffmpeg_args, "returncode": rc, "stderr": err, "codec": s_codec}, ) self.attempted.add(s_codec) finally: if stream_input.is_symlink(): stream_input.unlink()