live/backend/main.py

import os
from contextlib import asynccontextmanager
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse
from starlette.responses import RedirectResponse
from starlette.background import BackgroundTask
import httpx
import asyncio
import time
import traceback
import psutil
from database import init_db, session_factory
from services.transcoder import transcode_mgr
from services.stream_service import StreamService
from services.media_service import media_service
from services.mediamtx_manager import mediamtx_mgr
from models.stream import StreamStatus
from routes import streams_router, nodes_router, auth_router
from routes.settings_route import router as settings_router
from config import settings
from services.load_balancer import lb
from services.node_store import load_nodes


async def _origin_paths() -> dict[str, dict]:
    """Fetch MediaMTX path list keyed by path name. Empty dict when origin is unreachable."""
    try:
        async with httpx.AsyncClient(timeout=3) as client:
            response = await client.get(
                f"http://127.0.0.1:{settings.media_origin_api_port}/v3/paths/list"
            )
            response.raise_for_status()
            return {item.get("name"): item for item in response.json().get("items", [])}
    except Exception:
        return {}


async def supervise_transcoders() -> list[str]:
    """Restart transcoders that died or drifted from the stream's config.

    A live stream whose upstream is publishing (MediaMTX path ready) but whose
    local ffmpeg processes have exited would otherwise stay black forever: the
    HLS output stops updating and nothing re-triggers a start.

    Returns the stream keys that were (re)started this tick.
    """
    restarted: list[str] = []
    paths = await _origin_paths()
    if not paths:
        return restarted
    async with session_factory() as db:
        service = StreamService(db)
        for stream in await service.get_all_streams():
            if stream.status != StreamStatus.live:
                continue
            path = paths.get(stream.stream_key) or paths.get(f"live/{stream.stream_key}")
            # Only act when the publisher is actually connected: restarting against
            # a missing source just burns ffmpeg processes that exit immediately.
            if not path or not path.get("ready"):
                continue
            # Must match what start_transcode is given below: needs_restart
            # compares against the signature recorded at the last start, so
            # passing a different list here would restart the encoders every tick.
            qualities = service.transcode_renditions(stream)
            if transcode_mgr.needs_restart(stream.stream_key, qualities):
                transcode_mgr.start_transcode(
                    stream.stream_key,
                    media_service.get_rtmp_url(stream.stream_key),
                    qualities,
                )
                restarted.append(stream.stream_key)
    return restarted


async def _metrics_tick(previous, previous_time):
    current = psutil.net_io_counters()
    current_time = time.monotonic()
    elapsed = current_time - previous_time
    network_in = ((current.bytes_recv - previous.bytes_recv) * 8 / elapsed / 1e6) if elapsed > 0 else None
    network_out = ((current.bytes_sent - previous.bytes_sent) * 8 / elapsed / 1e6) if elapsed > 0 else None
    viewer_count = None
    try:
        async with httpx.AsyncClient(timeout=3) as client:
            response = await client.get(f"http://127.0.0.1:{settings.media_origin_api_port}/v3/paths/list")
            response.raise_for_status()
            payload = response.json()
            # MediaMTX reports readers as an array of reader objects; there is no
            # readersCount field in the version this project ships. Summing the
            # array directly raised TypeError as soon as anyone was watching,
            # which the bare except below turned into a permanent "no viewers".
            viewer_count = sum(len(item.get("readers") or []) for item in payload.get("items", []))
    except Exception:
        pass
    lb.report_health("default", {
        "cpu_percent": psutil.cpu_percent(),
        "memory_percent": psutil.virtual_memory().percent,
        "bandwidth_mbps": network_in + network_out if network_in is not None and network_out is not None else None,
        "network_in_mbps": network_in,
        "network_out_mbps": network_out,
        "active_connections": len(psutil.net_connections(kind="tcp")),
        "hls_readers": viewer_count,
        "viewer_count": viewer_count,
        "frame_drop_indicators": None,
        "packet_loss_indicators": None,
        "media_metrics_available": False,
    })
    return current, current_time


async def monitor_origin():
    """Background supervisor loop.

    Every step is individually guarded: this task is created with create_task, so
    an exception escaping it would silently kill supervision forever with no
    traceback reaching the logs. Each phase catches its own errors instead.
    """
    previous = psutil.net_io_counters()
    previous_time = time.monotonic()
    while True:
        await asyncio.sleep(10)
        try:
            previous, previous_time = await _metrics_tick(previous, previous_time)
        except Exception:
            traceback.print_exc()
        try:
            # MediaMTX first: the transcoder pass below only acts on a stream whose
            # MediaMTX path is ready, so a dead origin would otherwise be reported
            # as "nothing to do" forever.
            if await mediamtx_mgr.ensure_running():
                await mediamtx_mgr.wait_until_ready()
        except asyncio.CancelledError:
            raise
        except Exception:
            traceback.print_exc()
        try:
            restarted = await supervise_transcoders()
            for stream_key in restarted:
                print(f"[supervisor] restarted transcoder for {stream_key}", flush=True)
        except asyncio.CancelledError:
            raise
        except Exception:
            traceback.print_exc()


@asynccontextmanager
async def lifespan(app: FastAPI):
    load_nodes()
    await init_db()
    os.makedirs(settings.hls_runtime_dir, exist_ok=True)
    # Bring the media origin up before the transcoders: they immediately open RTMP
    # connections to it, and against a dead origin every ffmpeg exits at once.
    if settings.media_origin_autostart and not await mediamtx_mgr.is_healthy():
        mediamtx_mgr.start()
        if not await mediamtx_mgr.wait_until_ready():
            print("[supervisor] MediaMTX not ready yet; it will be retried", flush=True)
    async with session_factory() as db:
        service = StreamService(db)
        for stream in await service.get_all_streams():
            if stream.status == StreamStatus.live:
                transcode_mgr.start_transcode(
                    stream.stream_key,
                    media_service.get_rtmp_url(stream.stream_key),
                    service.transcode_renditions(stream),
                )
    origin_monitor = asyncio.create_task(monitor_origin())
    yield
    origin_monitor.cancel()
    transcode_mgr.stop_all()
    mediamtx_mgr.stop()


app = FastAPI(title="ProgramCore", lifespan=lifespan)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

app.include_router(streams_router)
app.include_router(nodes_router)
app.include_router(auth_router)
app.include_router(settings_router)


@app.get("/api/health")
async def health():
    from services.multicast_detector import multicast_detector
    return {
        "status": "ok",
        "multicast_available": multicast_detector.is_available(),
    }


@app.get("/hls/{path:path}")
async def serve_hls(path: str):
    hls_root = settings.hls_runtime_dir
    fs_path = os.path.join(hls_root, path)
    if os.path.isfile(fs_path):
        return FileResponse(fs_path)

    target = f"http://127.0.0.1:{settings.media_origin_hls_port}/{path}"
    # MediaMTX gates HLS behind a "cookieCheck" round trip: it answers with a 302 to
    # "<path>?cookieCheck=1" and sets a *Secure* cookie. Over plain HTTP (any
    # non-TLS deployment, and every dev machine) browsers refuse to store a Secure
    # cookie, so the redirect never terminates and the player loops until it gives
    # up. The browser should never see that handshake at all, so complete it here
    # and hand back only the final playlist.
    #
    # The client also has to outlive this function: the body is streamed after we
    # return, so an `async with` would close the pool before the first chunk.
    client = httpx.AsyncClient(timeout=httpx.Timeout(10), follow_redirects=True)
    try:
        request = client.build_request("GET", target)
        response = await client.send(request, stream=True)
    except Exception:
        await client.aclose()
        raise HTTPException(status_code=404)

    from starlette.responses import StreamingResponse

    headers = {
        key: value
        for key, value in response.headers.items()
        # Length/transfer headers describe the origin's framing, not ours, and
        # hop-by-hop headers must not be relayed.
        if key.lower() not in {"content-length", "transfer-encoding", "connection"}
    }

    async def relay():
        try:
            async for chunk in response.aiter_bytes():
                yield chunk
        finally:
            await response.aclose()
            await client.aclose()

    return StreamingResponse(
        relay(),
        status_code=response.status_code,
        headers=headers,
        media_type=response.headers.get("content-type"),
    )

    from starlette.responses import StreamingResponse

    headers = {
        key: value
        for key, value in response.headers.items()
        # Length/transfer headers describe the origin's framing, not ours, and
        # hop-by-hop headers must not be relayed.
        if key.lower() not in {"content-length", "transfer-encoding", "connection"}
    }
    # MediaMTX answers HLS requests with a 302 to "<path>?cookieCheck=1" so the
    # player picks up its session cookie. That Location is relative to the origin
    # root, so relaying it unchanged sends the browser to /<path> on the backend
    # host — not an /hls route, hence a 404 instead of a playlist. Re-anchor it.
    location = response.headers.get("location")
    if response.is_redirect and location:
        parts = urlsplit(location)
        if parts.path:
            headers["location"] = f"/hls{parts.path}" + (f"?{parts.query}" if parts.query else "")

    async def relay():
        try:
            async for chunk in response.aiter_bytes():
                yield chunk
        finally:
            await response.aclose()
            await client.aclose()

    return StreamingResponse(
        relay(),
        status_code=response.status_code,
        headers=headers,
        media_type=response.headers.get("content-type"),
    )