415 lines
16 KiB
Python
415 lines
16 KiB
Python
import asyncio
|
|
import hashlib
|
|
import logging
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
from collections.abc import Iterable
|
|
from pathlib import Path
|
|
from typing import TYPE_CHECKING, ClassVar
|
|
|
|
from aiohttp import web
|
|
|
|
from .config import Config
|
|
from .ffprobe import ffprobe
|
|
|
|
if TYPE_CHECKING:
|
|
from asyncio.subprocess import Process
|
|
|
|
from app.library.ffprobe import FFProbeResult
|
|
|
|
LOG: logging.Logger = logging.getLogger("player.segments")
|
|
|
|
|
|
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.
|
|
|
|
Supported encoder families (preference order):
|
|
- amd -> h264_amf
|
|
- nvidia -> h264_nvenc
|
|
- intel -> h264_qsv
|
|
- apple -> h264_videotoolbox
|
|
- software -> libx264
|
|
"""
|
|
|
|
# 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."
|
|
|
|
@staticmethod
|
|
def _encoder_map() -> dict[str, str]:
|
|
"""Return mapping from logical encoder family to ffmpeg codec name."""
|
|
return {
|
|
"intel": "h264_qsv",
|
|
"nvidia": "h264_nvenc",
|
|
"amd": "h264_amf",
|
|
"apple": "h264_videotoolbox",
|
|
"software": "libx264",
|
|
}
|
|
|
|
@staticmethod
|
|
def _preferred_families(configured: str) -> list[str]:
|
|
"""
|
|
Build a preference list based on configuration value.
|
|
|
|
Accepts either a logical family name (amd|nvidia|intel|apple|software),
|
|
a concrete ffmpeg encoder (e.g., h264_nvenc/libx264), or any other
|
|
string. If the value is not recognized, we try GPU families first and
|
|
fall back to software.
|
|
"""
|
|
families: list[str] = ["intel", "nvidia", "amd", "apple", "software"]
|
|
enc_map: dict[str, str] = Segments._encoder_map()
|
|
|
|
# If user configured a concrete ffmpeg encoder, prioritize that and
|
|
# include software as fallback.
|
|
if configured in enc_map.values():
|
|
# Map concrete to a family for consistent probing order
|
|
family: str | None = next((k for k, v in enc_map.items() if v == configured), None)
|
|
return [family, "software"] if family else ["software"]
|
|
|
|
cfg_lower: str = configured.strip().lower()
|
|
if cfg_lower in enc_map:
|
|
# Specific family requested; then software fallback
|
|
return [cfg_lower, "software"]
|
|
|
|
# Unknown/legacy value (e.g., libx264 default): try GPUs first, then software
|
|
return families
|
|
|
|
@staticmethod
|
|
def _ffmpeg_encoders() -> set[str]:
|
|
"""
|
|
Return a set of available ffmpeg encoder names by parsing
|
|
`ffmpeg -encoders` output. On any error, return an empty set to signal
|
|
that probing failed (callers should fall back to software).
|
|
"""
|
|
try:
|
|
result: subprocess.CompletedProcess[str] = subprocess.run(
|
|
["ffmpeg", "-hide_banner", "-loglevel", "error", "-encoders"], # noqa: S607
|
|
capture_output=True,
|
|
text=True,
|
|
check=False,
|
|
creationflags=subprocess.CREATE_NO_WINDOW if os.name == "nt" else 0,
|
|
)
|
|
out: str = (result.stdout or "") + (result.stderr or "")
|
|
except Exception:
|
|
return set()
|
|
|
|
encoders: set[str] = set()
|
|
if not out:
|
|
return encoders
|
|
|
|
# Each encoder appears as a token in the list; match known names
|
|
for name in Segments._encoder_map().values():
|
|
if name in out:
|
|
encoders.add(name)
|
|
|
|
return encoders
|
|
|
|
@staticmethod
|
|
def _select_encoder(preferred: Iterable[str]) -> str:
|
|
"""
|
|
Select the first available encoder according to the preferred family
|
|
order. Falls back to software if probing fails or none found.
|
|
"""
|
|
enc_map: dict[str, str] = Segments._encoder_map()
|
|
|
|
# Only on Linux, if no /dev/dri device, fall back to software.
|
|
# On macOS/Windows, this check is not applicable.
|
|
if sys.platform.startswith("linux") and not Segments._has_dri_devices():
|
|
return enc_map["software"]
|
|
|
|
available: set[str] = Segments._ffmpeg_encoders()
|
|
|
|
# If probing failed, prefer software to minimize failure risk
|
|
if not available:
|
|
return enc_map["software"]
|
|
|
|
for family in preferred:
|
|
name: str | None = enc_map.get(family)
|
|
if name and name in available:
|
|
return name
|
|
|
|
# Final fallback
|
|
return enc_map["software"]
|
|
|
|
@staticmethod
|
|
def _has_dri_devices() -> bool:
|
|
"""Check if /dev/dri exists and has entries (card0, renderD128, etc)."""
|
|
try:
|
|
dri = Path("/dev/dri")
|
|
if not dri.exists() or not dri.is_dir():
|
|
return False
|
|
for _ in dri.iterdir():
|
|
return True
|
|
return False
|
|
except Exception:
|
|
return False
|
|
|
|
async def build_ffmpeg_args(self, file: Path) -> list[str]:
|
|
try:
|
|
ff: FFProbeResult = await ffprobe(file)
|
|
except UnicodeDecodeError:
|
|
pass
|
|
|
|
tmpFile: Path = Path(tempfile.gettempdir()).joinpath(
|
|
f"ytptube_stream.{hashlib.sha256(str(file).encode()).hexdigest()}"
|
|
)
|
|
|
|
if not tmpFile.exists():
|
|
tmpFile.symlink_to(file, target_is_directory=False)
|
|
|
|
startTime: str = f"{0:.6f}" if self.index == 0 else f"{self.duration * self.index:.6f}"
|
|
|
|
fargs: list[str] = [
|
|
"-xerror",
|
|
"-hide_banner",
|
|
"-loglevel",
|
|
"error",
|
|
"-ss",
|
|
str(startTime),
|
|
"-t",
|
|
str(f"{self.duration:.6f}"),
|
|
"-copyts",
|
|
"-i",
|
|
f"file:{tmpFile}",
|
|
"-map_metadata",
|
|
"-1",
|
|
]
|
|
|
|
# Decide on encoder for the first segment that actually has video
|
|
if ff and hasattr(ff, "has_video") and ff.has_video():
|
|
# Initialize cached encoder only once, on first video segment
|
|
if not Segments._cache_initialized and self.index == 0 and self.vconvert:
|
|
preferred: list[str] = Segments._preferred_families(self.vcodec)
|
|
Segments._cached_vcodec = Segments._select_encoder(preferred)
|
|
Segments._cache_initialized = True
|
|
self.vcodec = Segments._cached_vcodec
|
|
LOG.info(f"Selected streaming video encoder: {self.vcodec}")
|
|
elif Segments._cache_initialized and Segments._cached_vcodec:
|
|
self.vcodec = Segments._cached_vcodec
|
|
|
|
# Base video args
|
|
v_args: list[str] = [
|
|
"-g",
|
|
"52",
|
|
"-map",
|
|
"0:v:0",
|
|
"-strict",
|
|
"-2",
|
|
]
|
|
|
|
# For hardware encoders, avoid forcing pix_fmt yuv420p.
|
|
if self.vconvert:
|
|
if self.vcodec in ("h264_nvenc", "h264_amf", "h264_videotoolbox"):
|
|
# Let ffmpeg pick appropriate conversions; do not force pix_fmt
|
|
pass
|
|
elif self.vcodec == "h264_qsv":
|
|
# On Linux with /dev/dri present, initialize QSV device and ensure
|
|
# frames are properly aligned and uploaded to GPU.
|
|
if sys.platform.startswith("linux") and Segments._has_dri_devices():
|
|
v_args += [
|
|
"-init_hw_device",
|
|
"qsv=hw:/dev/dri/renderD128",
|
|
"-filter_hw_device",
|
|
"hw",
|
|
"-vf",
|
|
"scale=trunc(iw/2)*2:trunc(ih/2)*2,format=nv12,hwupload=extra_hw_frames=64",
|
|
]
|
|
else:
|
|
# Software path: ensure broad compatibility
|
|
v_args = ["-pix_fmt", "yuv420p", *v_args]
|
|
|
|
v_args += ["-codec:v", self.vcodec if self.vconvert else "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 stream(self, file: Path, resp: web.StreamResponse):
|
|
ffmpeg_args: list[str] = await self.build_ffmpeg_args(file)
|
|
|
|
async def _run_and_stream(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 = await proc.stderr.read(4096)
|
|
if not chunk:
|
|
break
|
|
stderr_buf.extend(chunk)
|
|
except Exception:
|
|
# best-effort only
|
|
pass
|
|
|
|
stderr_task = asyncio.create_task(_drain_stderr())
|
|
|
|
LOG.debug(f"Streaming '{file}' segment '{self.index}'. ffmpeg: {' '.join(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.error("ffmpeg process did not terminate in time. Killing it.")
|
|
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.error("ffmpeg process did not terminate in time after disconnect. Killing it.")
|
|
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.error("ffmpeg process wait timed out. Killing it.")
|
|
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())
|
|
|
|
wrote, rc, client_disconnected, stderr_text = await _run_and_stream(ffmpeg_args)
|
|
|
|
is_software: bool = self.vcodec == Segments._encoder_map()["software"]
|
|
if not is_software and not wrote and rc != 0 and not client_disconnected:
|
|
msg: str = stderr_text[:500] if stderr_text else "no error output"
|
|
LOG.warning(f"Hardware encoder failed (rc={rc}): {msg}. Retrying with software (libx264).")
|
|
Segments._cached_vcodec = Segments._encoder_map()["software"]
|
|
Segments._cache_initialized = True
|
|
|
|
sw_args: list[str] = list(ffmpeg_args)
|
|
# Replace codec with software
|
|
try:
|
|
idx: int = sw_args.index("-codec:v")
|
|
sw_args[idx + 1] = Segments._encoder_map()["software"]
|
|
except ValueError:
|
|
sw_args += ["-codec:v", Segments._encoder_map()["software"]]
|
|
|
|
# Remove QSV-specific flags if present prior to software fallback.
|
|
# -init_hw_device qsv=...
|
|
try:
|
|
while True:
|
|
i_idx = sw_args.index("-init_hw_device")
|
|
del sw_args[i_idx : min(i_idx + 2, len(sw_args))]
|
|
except ValueError:
|
|
pass
|
|
# -filter_hw_device <name>
|
|
try:
|
|
while True:
|
|
f_idx = sw_args.index("-filter_hw_device")
|
|
del sw_args[f_idx : min(f_idx + 2, len(sw_args))]
|
|
except ValueError:
|
|
pass
|
|
# -vf QSV chain
|
|
try:
|
|
vf_idx = sw_args.index("-vf")
|
|
if (
|
|
vf_idx + 1 < len(sw_args)
|
|
and sw_args[vf_idx + 1]
|
|
== "scale=trunc(iw/2)*2:trunc(ih/2)*2,format=nv12,hwupload=extra_hw_frames=64"
|
|
):
|
|
del sw_args[vf_idx : vf_idx + 2]
|
|
except ValueError:
|
|
pass
|
|
|
|
# Ensure software path uses yuv420p
|
|
if "-pix_fmt" in sw_args:
|
|
try:
|
|
pf_idx = sw_args.index("-pix_fmt")
|
|
if pf_idx + 1 < len(sw_args):
|
|
sw_args[pf_idx + 1] = "yuv420p"
|
|
except Exception:
|
|
pass
|
|
else:
|
|
# Insert near the start, before mapping/gop where possible
|
|
insert_pos = 0
|
|
if "-g" in sw_args:
|
|
insert_pos = max(0, sw_args.index("-g") - 2)
|
|
sw_args[insert_pos:insert_pos] = ["-pix_fmt", "yuv420p"]
|
|
|
|
wrote, rc, client_disconnected, _ = await _run_and_stream(sw_args)
|
|
|
|
if not client_disconnected:
|
|
try:
|
|
await resp.write_eof()
|
|
except (ConnectionResetError, RuntimeError):
|
|
LOG.warning("Failed to write EOF; client already disconnected.")
|