Files
KiraTV/backend/core/restream.py
T
joaquin 5179dbfa6d feat: replace aiohttp pump with ffmpeg copy-mode remux
Uses ffmpeg -c copy (no transcoding) with -reconnect_streamed to handle
provider session-token refreshes transparently. ffmpeg normalises PTS/DTS
across reconnects, eliminating A/V desync permanently. Quality is bit-for-
bit identical to the source stream.
2026-05-17 23:45:11 +02:00

387 lines
14 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 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()
# ------------------------------------------------------------------
# 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}"
)
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._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,
}
# ------------------------------------------------------------------
# Internal helpers
# ------------------------------------------------------------------
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()
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:
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