crawle-snapp/crawler/recording.py

"""Per-session Xvfb/FFmpeg recording with shared-volume retention."""

from __future__ import annotations

import fcntl
import os
import re
import signal
import subprocess
import threading
import time
import uuid
from datetime import UTC, datetime
from pathlib import Path

from config import config
from logger import get_logger

log = get_logger(__name__)

_DISPLAY_LOCK = threading.Lock()
_SAFE_NAME = re.compile(r"[^A-Za-z0-9_.-]+")
_WIDTH = 1920
_HEIGHT = 1080


def _safe_label(label: str | None) -> str:
    cleaned = _SAFE_NAME.sub("-", label or "session").strip("-._")
    return cleaned[:80] or "session"


def _stop_process(process: subprocess.Popen[bytes] | None, *, interrupt: bool = False) -> None:
    if process is None or process.poll() is not None:
        return
    try:
        process.send_signal(signal.SIGINT if interrupt else signal.SIGTERM)
        process.wait(timeout=10)
    except subprocess.TimeoutExpired:
        process.kill()
        process.wait(timeout=5)


def _prune_recordings(directory: Path) -> None:
    keep = max(1, config.recording_max_files)
    lock_path = directory / ".retention.lock"
    with lock_path.open("a+b") as lock_file:
        fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX)
        recordings = sorted(
            directory.glob("*.mp4"),
            key=lambda path: path.stat().st_mtime_ns,
            reverse=True,
        )
        for stale in recordings[keep:]:
            try:
                stale.unlink()
                log.info("Removed expired recording %s", stale.name)
            except FileNotFoundError:
                pass


class SessionRecorder:
    """Own one virtual display and one FFmpeg recorder."""

    def __init__(self, label: str | None = None) -> None:
        self.label = _safe_label(label)
        self.display: str | None = None
        self._xvfb: subprocess.Popen[bytes] | None = None
        self._ffmpeg: subprocess.Popen[bytes] | None = None
        self._temp_path: Path | None = None
        self._final_path: Path | None = None

    def start(self) -> bool:
        directory = Path(config.recordings_dir)
        directory.mkdir(parents=True, exist_ok=True)
        stamp = datetime.now(UTC).strftime("%Y%m%dT%H%M%S.%fZ")
        token = uuid.uuid4().hex[:8]
        self._final_path = directory / f"{self.label}-{stamp}-{token}.mp4"
        self._temp_path = directory / f".{self._final_path.name}.part.mp4"

        try:
            with _DISPLAY_LOCK:
                self._start_xvfb()
            self._start_ffmpeg()
        except Exception as exc:
            log.error("Could not start session recording: %s", exc)
            self.stop(publish=False)
            return False

        log.info(
            "Recording session to %s at %d FPS",
            self._final_path.name,
            config.recording_fps,
        )
        return True

    def _start_xvfb(self) -> None:
        for display_number in range(100, 1000):
            socket_path = Path(f"/tmp/.X11-unix/X{display_number}")
            if socket_path.exists():
                continue
            display = f":{display_number}"
            process = subprocess.Popen(
                [
                    "Xvfb",
                    display,
                    "-screen",
                    "0",
                    f"{_WIDTH}x{_HEIGHT}x24",
                    "-nolisten",
                    "tcp",
                    "-ac",
                ],
                stdout=subprocess.DEVNULL,
                stderr=subprocess.PIPE,
            )
            deadline = time.monotonic() + 5
            while time.monotonic() < deadline:
                if process.poll() is not None:
                    stderr = (process.stderr.read() if process.stderr else b"").decode(
                        "utf-8", errors="replace"
                    )
                    raise RuntimeError(f"Xvfb exited early: {stderr.strip()}")
                if socket_path.exists():
                    self.display = display
                    self._xvfb = process
                    return
                time.sleep(0.05)
            _stop_process(process)
        raise RuntimeError("no free X display was available")

    def _start_ffmpeg(self) -> None:
        if self.display is None or self._temp_path is None:
            raise RuntimeError("virtual display is not ready")
        self._ffmpeg = subprocess.Popen(
            [
                "ffmpeg",
                "-y",
                "-nostdin",
                "-loglevel",
                "error",
                "-f",
                "x11grab",
                "-framerate",
                str(config.recording_fps),
                "-video_size",
                f"{_WIDTH}x{_HEIGHT}",
                "-i",
                f"{self.display}.0",
                "-c:v",
                "libx264",
                "-preset",
                "ultrafast",
                "-threads",
                "1",
                "-crf",
                "28",
                "-pix_fmt",
                "yuv420p",
                str(self._temp_path),
            ],
            stdout=subprocess.DEVNULL,
            stderr=subprocess.PIPE,
        )
        time.sleep(0.2)
        if self._ffmpeg.poll() is not None:
            stderr = (self._ffmpeg.stderr.read() if self._ffmpeg.stderr else b"").decode(
                "utf-8", errors="replace"
            )
            raise RuntimeError(f"FFmpeg exited early: {stderr.strip()}")

    def stop(self, *, publish: bool = True) -> Path | None:
        _stop_process(self._ffmpeg, interrupt=True)
        _stop_process(self._xvfb)
        self._ffmpeg = None
        self._xvfb = None

        if self._temp_path is None or self._final_path is None:
            return None
        if not publish or not self._temp_path.exists() or self._temp_path.stat().st_size == 0:
            self._temp_path.unlink(missing_ok=True)
            return None

        os.replace(self._temp_path, self._final_path)
        _prune_recordings(self._final_path.parent)
        log.info("Recording saved: %s", self._final_path.name)
        return self._final_path