Files
KiraTV/backend/core/restream.py
T
joaquin ac06691656 feat: show ffmpeg PID and CPU% per stream in dashboard
Backend samples /proc/{pid}/stat every 2s to calculate CPU usage of
each ffmpeg process and exposes ffmpeg_pid + ffmpeg_cpu_pct in stats().
Dashboard shows a cyan ffmpeg badge with live CPU% per stream row,
plus a new StatCard with total ffmpeg CPU across all active streams.
2026-05-17 23:50:10 +02:00

432 lines
16 KiB
Python

"""BroadcastGroup: fan-out a single provider stream to N simultaneous clients.
Recovery model
--------------
The pump loop runs indefinitely while self._running is True. Each iteration
spawns an ffmpeg process in copy (remux) mode. ffmpeg owns the upstream
HTTP connection and uses its built-in -reconnect_streamed flag to handle
short provider interruptions (e.g. expiring session tokens) transparently,
normalising PTS/DTS in the output so clients never see an A/V discontinuity.
If ffmpeg itself exits (unrecoverable error or stall beyond its own timeout)
the outer loop drains all client queues, backs off briefly, and spawns a
fresh ffmpeg process — again without disconnecting clients.
SENTINEL is only pushed when stop() is called or the pump truly gives up.
Stall detection
---------------
ffmpeg's -timeout option (socket-level read timeout) handles provider stalls.
A safety-net asyncio.wait_for wraps each stdout read in case ffmpeg hangs
without producing output or exiting.
A/V sync
--------
ffmpeg re-muxes the incoming MPEG-TS into a fresh output stream. On every
reconnect the output PTS/DTS sequence is continuous regardless of provider-
side timestamp resets, eliminating audio desync entirely. No video or audio
is re-encoded — quality is bit-for-bit identical to the source.
"""
import asyncio
import logging
import os
import shutil
import time
import uuid
from dataclasses import dataclass, field
from typing import AsyncIterator
from ..config import settings
logger = logging.getLogger(__name__)
CHUNK_SIZE = settings.stream_chunk_size
QUEUE_MAX = settings.stream_queue_maxsize
STALL_TIMEOUT = 10.0 # seconds — passed to ffmpeg -timeout (in µs below)
MAX_RETRIES = 999 # effectively infinite; pool stops the group when 0 clients
_BACKOFF_CAP = 5.0 # max seconds between outer retries
_SENTINEL = object() # signals end-of-stream to client queues
_FFMPEG = shutil.which("ffmpeg") or "ffmpeg"
@dataclass
class ClientHandle:
client_id: str
queue: asyncio.Queue = field(default_factory=lambda: asyncio.Queue(maxsize=QUEUE_MAX))
connected_at: float = field(default_factory=time.time)
is_priority: bool = False
username: str = ""
_overflow_count: int = 0
_overflow_since: float | None = None
async def read(self) -> AsyncIterator[bytes]:
while True:
chunk = await self.queue.get()
if chunk is _SENTINEL:
break
yield chunk
class BroadcastGroup:
"""Streams one provider URL and distributes chunks to all registered clients."""
def __init__(
self,
channel_id: int,
stream_url: str,
provider_account_id: int,
channel_name: str = "",
provider_name: str = "",
):
self.channel_id = channel_id
self.channel_name = channel_name
self.stream_url = stream_url
self.provider_account_id = provider_account_id
self.provider_name = provider_name
self.started_at = time.time()
self._clients: dict[str, ClientHandle] = {}
self._lock = asyncio.Lock()
self._task: asyncio.Task | None = None
self._running = False
# Throughput tracking
self._bytes_pumped: int = 0
self._bytes_prev: int = 0
self._bps_time: float = time.time()
self._bps_down: float = 0.0
# Health / metadata
self.stream_info: dict = {}
self._reconnect_count: int = 0
self._status: str = "connecting"
self._last_data_at: float = time.time()
# ffmpeg process tracking
self._ffmpeg_pid: int | None = None
self._ffmpeg_cpu: float = 0.0
self._cpu_task: asyncio.Task | None = None
# ------------------------------------------------------------------
# Public API
# ------------------------------------------------------------------
@property
def client_count(self) -> int:
return len(self._clients)
@property
def client_ids(self) -> list[str]:
return list(self._clients.keys())
@property
def all_non_priority(self) -> bool:
return all(not h.is_priority for h in self._clients.values())
async def start(self) -> None:
self._running = True
self._task = asyncio.create_task(
self._pump_with_retry(), name=f"pump-ch{self.channel_id}"
)
self._cpu_task = asyncio.create_task(
self._monitor_cpu(), name=f"cpu-ch{self.channel_id}"
)
asyncio.create_task(self._probe_info(), name=f"probe-ch{self.channel_id}")
async def add_client(self, user_id: str, username: str = "", is_priority: bool = False) -> ClientHandle:
client_id = f"{user_id}-{uuid.uuid4().hex[:8]}"
handle = ClientHandle(client_id=client_id, is_priority=is_priority, username=username)
async with self._lock:
self._clients[client_id] = handle
logger.info(
f"Client {client_id} ({username}) joined channel {self.channel_id} "
f"(total={len(self._clients)})"
)
return handle
async def force_disconnect_client(self, client_id: str) -> bool:
async with self._lock:
handle = self._clients.get(client_id)
if handle is None:
return False
_drain_queue(handle.queue)
try:
handle.queue.put_nowait(_SENTINEL)
except asyncio.QueueFull:
pass
return True
async def remove_client(self, client_id: str) -> int:
async with self._lock:
self._clients.pop(client_id, None)
remaining = len(self._clients)
logger.info(f"Client {client_id} left channel {self.channel_id} (remaining={remaining})")
return remaining
async def stop(self) -> None:
self._running = False
self._status = "stopped"
if self._cpu_task and not self._cpu_task.done():
self._cpu_task.cancel()
if self._task and not self._task.done():
self._task.cancel()
try:
await self._task
except asyncio.CancelledError:
pass
async with self._lock:
for handle in self._clients.values():
try:
handle.queue.put_nowait(_SENTINEL)
except asyncio.QueueFull:
pass
def stats(self) -> dict:
n = self.client_count or 1
return {
"stream_type": "live",
"channel_id": self.channel_id,
"channel_name": self.channel_name or f"Canal #{self.channel_id}",
"provider_account_id": self.provider_account_id,
"provider_name": self.provider_name,
"client_count": self.client_count,
"clients": [
{
"client_id": h.client_id,
"username": h.username or h.client_id.rsplit("-", 1)[0],
}
for h in self._clients.values()
],
"started_at": self.started_at,
"bytes_pumped": self._bytes_pumped,
"bps_down": self._bps_down,
"bps_up": self._bps_down * n,
"running": self._running,
"status": self._status,
"reconnect_count": self._reconnect_count,
"last_data_ago": round(time.time() - self._last_data_at, 1),
"stream_info": self.stream_info,
# ffmpeg process info
"ffmpeg_pid": self._ffmpeg_pid,
"ffmpeg_cpu_pct": round(self._ffmpeg_cpu, 1),
}
# ------------------------------------------------------------------
# Internal helpers
# ------------------------------------------------------------------
async def _monitor_cpu(self) -> None:
"""Sample /proc/{pid}/stat every 2 s and update self._ffmpeg_cpu (%)."""
clk_tck = float(os.sysconf(os.sysconf_names.get("SC_CLK_TCK", 2)) or 100)
prev_ticks: int | None = None
prev_wall: float | None = None
while self._running:
await asyncio.sleep(2)
pid = self._ffmpeg_pid
if pid is None:
self._ffmpeg_cpu = 0.0
prev_ticks = prev_wall = None
continue
try:
with open(f"/proc/{pid}/stat") as f:
parts = f.read().split()
ticks = int(parts[13]) + int(parts[14]) # utime + stime
wall = time.time()
if prev_ticks is not None and prev_wall is not None:
dt = wall - prev_wall
if dt > 0:
self._ffmpeg_cpu = (ticks - prev_ticks) / clk_tck / dt * 100
prev_ticks = ticks
prev_wall = wall
except (FileNotFoundError, IndexError, ValueError, OSError):
self._ffmpeg_cpu = 0.0
prev_ticks = prev_wall = None
async def _probe_info(self) -> None:
from .probe import probe_stream
try:
info = await probe_stream(self.stream_url)
self.stream_info.update(info)
except Exception as e:
logger.debug(f"probe_info failed for ch{self.channel_id}: {e}")
async def _pump_with_retry(self) -> None:
consecutive = 0
base_wait = 0.5
try:
while self._running:
try:
await self._pump_once()
break
except asyncio.CancelledError:
raise
except Exception as e:
if not self._running:
break
consecutive += 1
self._reconnect_count += 1
wait = min(base_wait * (2 ** (consecutive - 1)), _BACKOFF_CAP)
self._status = "reconnecting"
logger.warning(
f"[ch{self.channel_id}] ffmpeg exited (attempt {consecutive}), "
f"restarting in {wait:.1f}s: {e}"
)
# Drain queues so stale buffered chunks are not delivered
# after the fresh ffmpeg process starts.
async with self._lock:
for handle in self._clients.values():
_drain_queue(handle.queue)
await asyncio.sleep(wait)
finally:
self._status = "stopped"
self._running = False
async with self._lock:
for handle in self._clients.values():
_drain_queue(handle.queue)
try:
handle.queue.put_nowait(_SENTINEL)
except asyncio.QueueFull:
pass
async def _pump_once(self) -> None:
"""
Spawn ffmpeg in copy (remux) mode and pump its stdout to all client
queues until ffmpeg exits or self._running goes False.
ffmpeg handles HTTP reconnects internally (-reconnect_streamed), which
covers provider session-token refreshes without any visible interruption
to clients. PTS/DTS values are normalised by ffmpeg across reconnects,
eliminating A/V desync entirely.
"""
cmd = [
_FFMPEG,
"-hide_banner", "-loglevel", "warning",
# Follow HTTP redirects and reconnect when the connection drops.
# -reconnect_streamed is the key flag: it retries even on non-
# seekable (live) streams, re-issuing the original GET so that
# providers using short-lived 302 tokens get a fresh token.
"-reconnect", "1",
"-reconnect_streamed", "1",
"-reconnect_delay_max", "4",
# Socket-level read timeout in microseconds. If the provider
# stops sending bytes for STALL_TIMEOUT seconds ffmpeg closes
# the connection and tries to reconnect (or exits if it gives up).
"-timeout", str(int(STALL_TIMEOUT * 1_000_000)),
"-i", self.stream_url,
"-c", "copy", # remux only — zero transcoding, bit-for-bit quality
"-f", "mpegts",
"pipe:1",
]
proc = await asyncio.create_subprocess_exec(
*cmd,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE,
)
self._status = "ok"
self._last_data_at = time.time()
self._ffmpeg_pid = proc.pid
logger.info(
f"[ch{self.channel_id}] ffmpeg started pid={proc.pid} "
f"(reconnects so far: {self._reconnect_count})"
)
# Drain stderr in background so the pipe never fills and blocks ffmpeg.
stderr_task = asyncio.create_task(self._log_ffmpeg_stderr(proc))
try:
while self._running:
try:
chunk = await asyncio.wait_for(
proc.stdout.read(CHUNK_SIZE),
timeout=STALL_TIMEOUT + 15, # safety net beyond ffmpeg's own timeout
)
except asyncio.TimeoutError:
raise RuntimeError(
f"ffmpeg produced no output for {STALL_TIMEOUT + 15:.0f}s"
)
if not chunk:
# ffmpeg exited cleanly or with error
rc = proc.returncode
raise RuntimeError(f"ffmpeg exited (rc={rc})")
# Throughput accounting
self._bytes_pumped += len(chunk)
self._last_data_at = time.time()
now = self._last_data_at
if now - self._bps_time >= 2.0:
self._bps_down = (self._bytes_pumped - self._bytes_prev) / (now - self._bps_time)
self._bytes_prev = self._bytes_pumped
self._bps_time = now
await self._dispatch(chunk, now)
finally:
self._ffmpeg_pid = None
self._ffmpeg_cpu = 0.0
stderr_task.cancel()
try:
proc.kill()
except ProcessLookupError:
pass
try:
await asyncio.wait_for(proc.wait(), timeout=5)
except asyncio.TimeoutError:
pass
async def _log_ffmpeg_stderr(self, proc: asyncio.subprocess.Process) -> None:
"""Read ffmpeg stderr lines and forward them to the logger."""
try:
async for line in proc.stderr:
text = line.decode(errors="replace").rstrip()
if text:
logger.warning(f"[ch{self.channel_id}] ffmpeg: {text}")
except asyncio.CancelledError:
pass
except Exception:
pass
async def _dispatch(self, chunk: bytes, now: float) -> None:
"""Push chunk to every client queue; evict clients that have been full for >30 s."""
async with self._lock:
dead: list[str] = []
for cid, handle in self._clients.items():
try:
handle.queue.put_nowait(chunk)
handle._overflow_count = 0
handle._overflow_since = None
except asyncio.QueueFull:
handle._overflow_count += 1
if handle._overflow_since is None:
handle._overflow_since = now
elif now - handle._overflow_since > 30:
dead.append(cid)
for cid in dead:
logger.warning(
f"[ch{self.channel_id}] Evicting unresponsive client {cid} "
f"after 30 s of sustained queue overflow"
)
dropped = self._clients.pop(cid, None)
if dropped is not None:
_drain_queue(dropped.queue)
try:
dropped.queue.put_nowait(_SENTINEL)
except asyncio.QueueFull:
pass
def _drain_queue(q: asyncio.Queue) -> None:
"""Empty a queue non-blockingly."""
while not q.empty():
try:
q.get_nowait()
except asyncio.QueueEmpty:
break