live/backend/services/transcoder.py

import subprocess
import os
from config import settings


class TranscodeManager:
    _processes: dict[str, list[subprocess.Popen]] = {}
    # Signature of the last successful start, so a supervisor can tell whether a
    # running set of processes still matches the stream's current configuration.
    _signatures: dict[str, str] = {}
    _failed_starts: dict[str, int] = {}

    @staticmethod
    def _signature(qualities: list[dict]) -> str:
        return "|".join(
            f"{q.get('label')}:{q.get('width')}x{q.get('height')}:{q.get('bitrate')}:"
            f"{'copy' if q.get('copy') else q.get('codec', 'libx264')}"
            for q in qualities
        )

    def is_running(self, stream_key: str) -> bool:
        """True when at least one transcode process exists and none have exited."""
        procs = self._processes.get(stream_key)
        if not procs:
            return False
        return any(proc.poll() is None for proc in procs)

    def needs_restart(self, stream_key: str, qualities: list[dict]) -> bool:
        if not qualities:
            return False
        return self._signatures.get(stream_key) != self._signature(qualities) or not self.is_running(stream_key)

    def _clear_stale_output(self, hls_dir: str) -> None:
        """Drop the previous run's playlist and segments.

        ffmpeg writes a live playlist with no #EXT-X-ENDLIST, so a playlist left
        behind by an encoder that died looks to hls.js exactly like a stream that
        simply stopped: it keeps reloading a file that will never change and
        buffers forever. Removing it makes the player see a missing manifest
        instead, which it retries until the supervisor has the encoder back up.
        """
        if not os.path.isdir(hls_dir):
            return
        for name in os.listdir(hls_dir):
            if name.endswith(".ts") or name.endswith(".m3u8") or name.endswith(".m3u8.tmp"):
                try:
                    os.remove(os.path.join(hls_dir, name))
                except OSError:
                    # A player holding the file open (Windows) is fine: it will be
                    # overwritten by the restarting encoder.
                    pass

    def start_transcode(self, stream_key: str, input_url: str, qualities: list[dict]):
        if stream_key in self._processes and self.is_running(stream_key):
            return
        self.stop_transcode(stream_key)

        procs = []
        segment_seconds = max(1, settings.hls_segment_seconds)
        for quality in qualities:
            label = quality["label"]
            # The stream-copy ("原画") rendition carries the top rendition's numbers
            # only so the signature stays complete; ffmpeg copies the input rather
            # than encoding to them. See StreamService.transcode_renditions.
            is_copy = bool(quality.get("copy"))
            # Files land under the rendition's path slug, never its label: MediaMTX
            # refuses non-alphanumeric HLS paths and the label may be Chinese.
            path_slug = quality.get("path") or label
            width = quality["width"]
            height = quality["height"]
            fps = quality.get("fps", 30)
            codec = quality.get("codec", "libx264")
            bitrate = quality["bitrate"]
            crf = quality.get("crf", 23)
            preset = quality.get("preset", "ultrafast")
            audio_bitrate = quality.get("audio_bitrate", "128k")
            keyframe_seconds = quality.get("keyframe_seconds", segment_seconds)
            keyframe_frames = keyframe_seconds * fps
            hls_dir = os.path.join(settings.hls_runtime_dir, stream_key, path_slug)
            os.makedirs(hls_dir, exist_ok=True)
            self._clear_stale_output(hls_dir)
            hls_output = [
                "-f", "hls",
                "-hls_time", str(segment_seconds),
                "-hls_list_size", "6",
                "-hls_flags", "delete_segments+independent_segments+program_date_time",
                "-hls_segment_filename", os.path.join(hls_dir, "segment_%03d.ts"),
                os.path.join(hls_dir, "index.m3u8"),
            ]
            # NOTE: ffmpeg's -reconnect* flags are HTTP/HTTPS-only protocol
            # options and make RTMP inputs fail with "Option not found".
            # RTMP drop recovery is handled by the supervisor in main.py,
            # which restarts exited encoders on its next tick.
            hls_input = ["-rw_timeout", "10000000", "-i", input_url]
            if is_copy:
                # Source ("原画") rendition: keep the publisher's video as-is and only
                # re-wrap it, so the output keeps the exact incoming resolution,
                # frame rate and bitrate. Video copy costs no CPU; audio is re-encoded
                # because MPEG-TS cannot carry every codec a publisher might send, and
                # one AAC track is a rounding error next to the video encoders.
                cmd = (
                    ["ffmpeg", "-y"] + hls_input
                    + [
                        "-c:v", "copy",
                        "-c:a", "aac", "-b:a", "192k",
                        "-copyts", "-start_at_zero",
                        "-map", "0:v:0", "-map", "0:a:0?",
                    ]
                    + hls_output
                )
            else:
                cmd = (
                    ["ffmpeg", "-y"] + hls_input
                    + [
                        "-c:v", codec, "-preset", preset,
                        "-r", str(fps), "-b:v", bitrate, "-crf", str(crf),
                        "-s", f"{width}x{height}",
                        "-sc_threshold", "0",
                        "-g", str(keyframe_frames),
                        "-keyint_min", str(keyframe_frames),
                        "-force_key_frames", f"expr:gte(t,n_forced*{keyframe_seconds})",
                        "-copyts", "-start_at_zero",
                        "-c:a", "aac", "-b:a", audio_bitrate,
                    ]
                    + hls_output
                )

            try:
                proc = subprocess.Popen(
                    cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL
                )
                procs.append(proc)
            except FileNotFoundError:
                pass

        self._processes[stream_key] = procs
        self._signatures[stream_key] = self._signature(qualities)
        self._failed_starts[stream_key] = 0 if procs else self._failed_starts.get(stream_key, 0) + 1

    def stop_transcode(self, stream_key: str):
        procs = self._processes.pop(stream_key, [])
        for proc in procs:
            try:
                proc.terminate()
                proc.wait(timeout=5)
            except Exception:
                try:
                    proc.kill()
                except Exception:
                    pass

    def stop_all(self):
        for key in list(self._processes.keys()):
            self.stop_transcode(key)


transcode_mgr = TranscodeManager()