284 lines
10 KiB
Python
284 lines
10 KiB
Python
import asyncio
|
|
import os
|
|
import subprocess # type: ignore
|
|
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:
|
|
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 = Segments._cached_vcodec if Segments._cache_initialized else self.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()
|