104 lines
3.6 KiB
Python
104 lines
3.6 KiB
Python
import asyncio
|
|
from collections.abc import AsyncIterator
|
|
|
|
from loguru import logger
|
|
|
|
from restreamer.constants import FFMPEG_READ_CHUNK_SIZE, FFMPEG_TERMINATE_TIMEOUT_SECONDS
|
|
from restreamer.resolver import ResolvedStream
|
|
|
|
_RECONNECT_INPUT_ARGS = ["-reconnect", "1", "-reconnect_streamed", "1", "-reconnect_delay_max", "5"]
|
|
|
|
|
|
def build_ffmpeg_args(stream: ResolvedStream) -> list[str]:
|
|
args = ["ffmpeg", "-loglevel", "warning", "-nostdin"]
|
|
|
|
header_lines = "".join(f"{key}: {value}\r\n" for key, value in stream.http_headers.items())
|
|
|
|
if header_lines:
|
|
args += ["-headers", header_lines]
|
|
args += [*_RECONNECT_INPUT_ARGS, "-i", stream.video_url]
|
|
|
|
if stream.audio_url:
|
|
if header_lines:
|
|
args += ["-headers", header_lines]
|
|
args += [*_RECONNECT_INPUT_ARGS, "-i", stream.audio_url]
|
|
args += ["-map", "0:v:0", "-map", "1:a:0"]
|
|
|
|
args += [
|
|
"-c",
|
|
"copy",
|
|
"-f",
|
|
"mpegts",
|
|
"-mpegts_flags",
|
|
"+resend_headers",
|
|
"pipe:1",
|
|
]
|
|
return args
|
|
|
|
|
|
class FfmpegStreamProcess:
|
|
"""Wraps a running ffmpeg process that copies a resolved stream to stdout as MPEG-TS."""
|
|
|
|
def __init__(self, stream: ResolvedStream, channel_id: str) -> None:
|
|
self._stream = stream
|
|
self._channel_id = channel_id
|
|
self._process: asyncio.subprocess.Process | None = None
|
|
self._stderr_task: asyncio.Task[None] | None = None
|
|
|
|
async def start(self) -> None:
|
|
args = build_ffmpeg_args(self._stream)
|
|
logger.debug("Starting ffmpeg for channel '{}'", self._channel_id)
|
|
self._process = await asyncio.create_subprocess_exec(
|
|
*args,
|
|
stdout=asyncio.subprocess.PIPE,
|
|
stderr=asyncio.subprocess.PIPE,
|
|
)
|
|
self._stderr_task = asyncio.create_task(self._drain_stderr())
|
|
|
|
async def iter_chunks(self) -> AsyncIterator[bytes]:
|
|
if self._process is None or self._process.stdout is None:
|
|
raise RuntimeError("ffmpeg process has not been started")
|
|
|
|
while True:
|
|
chunk = await self._process.stdout.read(FFMPEG_READ_CHUNK_SIZE)
|
|
if not chunk:
|
|
return
|
|
yield chunk
|
|
|
|
async def wait(self) -> int:
|
|
if self._process is None:
|
|
raise RuntimeError("ffmpeg process has not been started")
|
|
return await self._process.wait()
|
|
|
|
async def stop(self) -> None:
|
|
if self._process is None or self._process.returncode is not None:
|
|
return
|
|
|
|
# wait() only returns once every pipe hits EOF, so unread stdout left over
|
|
# from the consumer breaking out of iter_chunks() must be drained, or it hangs forever.
|
|
stdout_drain_task = asyncio.create_task(self._drain_stdout())
|
|
|
|
self._process.terminate()
|
|
try:
|
|
await asyncio.wait_for(self._process.wait(), timeout=FFMPEG_TERMINATE_TIMEOUT_SECONDS)
|
|
except TimeoutError:
|
|
logger.warning("ffmpeg for channel '{}' did not exit in time, killing it", self._channel_id)
|
|
self._process.kill()
|
|
await self._process.wait()
|
|
|
|
await stdout_drain_task
|
|
if self._stderr_task is not None:
|
|
await self._stderr_task
|
|
|
|
async def _drain_stdout(self) -> None:
|
|
if self._process is None or self._process.stdout is None:
|
|
return
|
|
while await self._process.stdout.read(FFMPEG_READ_CHUNK_SIZE):
|
|
pass
|
|
|
|
async def _drain_stderr(self) -> None:
|
|
if self._process is None or self._process.stderr is None:
|
|
return
|
|
async for line in self._process.stderr:
|
|
logger.debug("[ffmpeg:{}] {}", self._channel_id, line.decode(errors="replace").rstrip())
|