Files
KiraTV/backend/core/restream.py
T
joaquin f5d6745ed0 Initial commit: KiraTV IPTV proxy server
Full-stack IPTV management system with Xtream Codes API compatibility,
multi-account provider fan-out, Jellyfin local library integration,
real-time dashboard, EPG support, and A/V-sync-safe stream recovery.
2026-05-17 23:22:05 +02:00

351 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
opens a fresh upstream connection. If the connection fails or stalls the loop
sends SENTINEL to every client (forcing the player to reconnect from scratch),
sleeps briefly, then retries the upstream connection.
Sending SENTINEL on reconnect is intentional: keeping the TCP connection alive
while splicing in chunks from a new provider connection causes a PTS/DTS jump
that most players cannot reconcile silently — the result is permanent audio
desync. A clean player reconnect gives the player a fresh timestamp timeline
and guarantees A/V sync from the first frame of the new connection.
Stall detection
---------------
Each chunk read is wrapped in asyncio.wait_for(STALL_TIMEOUT). If the
upstream stops sending bytes for STALL_TIMEOUT seconds we treat it as a dead
stream and reconnect immediately rather than waiting for a TCP timeout that
could take minutes.
"""
import asyncio
import logging
import time
import uuid
from dataclasses import dataclass, field
from typing import AsyncIterator
import aiohttp
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 of silence → reconnect (fast stall detection)
MAX_RETRIES = 999 # effectively infinite — pool stops the group when 0 clients remain
# Back-off: 0.5 s, 1 s, 2 s, 4 s … capped at 5 s
_BACKOFF_CAP = 5.0
_SENTINEL = object() # signals end-of-stream to client queues
@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 = {} # populated by background ffprobe
self._reconnect_count: int = 0
self._status: str = "connecting" # connecting | ok | reconnecting | stopped
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}"
)
# Probe stream info in background — does NOT delay the first chunk
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,
# Health
"status": self._status,
"reconnect_count": self._reconnect_count,
"last_data_ago": round(time.time() - self._last_data_at, 1),
# Stream metadata from ffprobe
"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:
"""
Outer retry loop. Each iteration calls _pump_once() which runs until
the upstream connection dies or stalls. On failure we back off and
retry without touching client connections.
"""
consecutive = 0
base_wait = 1.0
try:
while self._running:
try:
await self._pump_once()
# _pump_once returned cleanly (self._running went False)
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}] Upstream error (attempt {consecutive}), "
f"reconnecting in {wait:.1f}s: {e}"
)
# Send SENTINEL to every client so the player drops its
# own buffers and reconnects from scratch. This is the
# only reliable way to prevent A/V desync: if we keep the
# TCP connection alive but splice in chunks from a new
# provider connection, the player sees a PTS/DTS jump it
# cannot reconcile silently. A clean player reconnect
# gives it a fresh timeline and guarantees sync.
async with self._lock:
for handle in self._clients.values():
_drain_queue(handle.queue)
try:
handle.queue.put_nowait(_SENTINEL)
except asyncio.QueueFull:
pass
await asyncio.sleep(wait)
finally:
self._status = "stopped"
self._running = False
# Signal all waiting clients that the stream is over
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:
"""
Open one upstream connection and pump until EOF, stall, or error.
Resets consecutive-failure counter on first successful chunk.
"""
timeout = aiohttp.ClientTimeout(
connect=settings.stream_connect_timeout,
sock_read=None, # we control read timeouts via wait_for ourselves
)
async with aiohttp.ClientSession(timeout=timeout) as session:
async with session.get(self.stream_url, ssl=False) as resp:
if resp.status not in (200, 206):
raise RuntimeError(f"Provider returned HTTP {resp.status}")
self._status = "ok"
self._last_data_at = time.time()
logger.info(
f"[ch{self.channel_id}] Stream connected "
f"(reconnects so far: {self._reconnect_count})"
)
while self._running:
# --- read one chunk with stall timeout ---
try:
chunk = await asyncio.wait_for(
resp.content.readany(), timeout=STALL_TIMEOUT
)
except asyncio.TimeoutError:
raise RuntimeError(
f"Stream stalled: no data for {STALL_TIMEOUT:.0f}s"
)
if not chunk:
raise RuntimeError("Provider closed the stream (EOF)")
# --- 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
# --- distribute to clients ---
await self._dispatch(chunk, now)
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:
# Drop-newest: discard this chunk for the slow client.
# Preserves the already-buffered sequence → no A/V desync from gaps.
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