kaledonya-clips/intake/app/previews.py

"""Small Discord copies of big clips.

Discord's crawler gives a linked video about 8 seconds to download and a home upload link of ~4 MB/s means a clip
over ~30 MB never gets a player. For those, the /v/ page points Discord at a 720p copy that fits in ~22 MB
(clips over 4 minutes: 480p in ~18 MB, so they download with room to spare);
the original stays untouched. Copies are made one at a time in the background and live in WORK_DIR/previews.
"""
import hashlib
import json
import logging
import queue
import re
import shutil
import subprocess
import threading
import time
from pathlib import Path
from urllib.parse import quote

log = logging.getLogger("uvicorn.error")  # uvicorn only prints its own loggers

TARGET_BYTES = 22 * 1000 * 1000
LONG_SECONDS = 240  # longer clips get a smaller, 480p copy
LONG_TARGET_BYTES = 18 * 1000 * 1000
AUDIO_KBPS = 128
MAX_VIDEO_KBPS = 4000
MIN_VIDEO_KBPS = 250  # below this a clip is too long to be worth watching in Discord; it gets no copy
KEEP_DAYS = 95  # clips delete after 90 days
KEY_RE = re.compile(r"^[0-9a-f]{24}$")


def target_bytes(duration: float) -> int:
    return LONG_TARGET_BYTES if duration > LONG_SECONDS else TARGET_BYTES


def video_kbps(duration: float) -> int | None:
    """Video bitrate that makes `duration` seconds land near target_bytes(), or None when it can't sensibly fit."""
    if duration <= 0:
        return None
    kbps = min(MAX_VIDEO_KBPS, int(target_bytes(duration) * 8 / 1000 / duration) - AUDIO_KBPS)
    return kbps if kbps >= MIN_VIDEO_KBPS else None


def encode_cmd(src: str, dst: Path, kbps: int, duration: float) -> list[str]:
    height = 480 if duration > LONG_SECONDS else 720
    return [
        "ffmpeg", "-v", "error", "-y", "-threads", "4", "-i", src,
        "-map", "0:v:0", "-map", "0:a:0?",  # OBS clips carry a second audio track; Discord plays one
        "-vf", f"scale=-2:'min({height},ih)'", "-fpsmax", "30",
        "-c:v", "libx264", "-preset", "veryfast", "-pix_fmt", "yuv420p",
        "-b:v", f"{kbps}k", "-maxrate", f"{kbps * 6 // 5}k", "-bufsize", f"{kbps * 2}k",
        "-c:a", "aac", "-b:a", f"{AUDIO_KBPS}k", "-ac", "2",
        "-movflags", "+faststart", "-progress", "pipe:1", "-nostats", str(dst),
    ]


def follow_progress(lines, duration: float, report) -> None:
    """Turn ffmpeg's `-progress` lines into fractions done (0..1) for the site's upload dialog."""
    for line in lines:
        key, _, value = line.strip().partition("=")
        if key == "out_time_us" and value.isdigit() and duration > 0:
            report(min(1.0, round(int(value) / 1e6 / duration, 3)))


def duration_of(src: str) -> float:
    cmd = ["ffprobe", "-v", "error", "-show_entries", "format=duration", "-of", "json", src]
    try:
        out = subprocess.run(cmd, capture_output=True, timeout=30).stdout
        return float(json.loads(out)["format"]["duration"])
    except (OSError, subprocess.TimeoutExpired, ValueError, KeyError):
        return 0.0


def niced(cmd: list[str]) -> list[str]:
    """Run at low CPU priority where `nice` exists, so uploads and everything else on the box come first."""
    return ["nice", "-n", "10", *cmd] if shutil.which("nice") else cmd


class Previews:
    def __init__(self, directory: Path, zipline_url: str, over_bytes: int):
        self.dir = directory
        self.zipline_url = zipline_url.rstrip("/")
        self.over_bytes = over_bytes
        self.jobs: queue.Queue[tuple[str, int]] = queue.Queue()
        self.pending: set[str] = set()
        self.lock = threading.Lock()
        self.worker: threading.Thread | None = None
        self.current: str | None = None  # key being encoded right now, and how far along it is
        self.progress = 0.0

    @staticmethod
    def key(name: str, length: int) -> str:
        # the size is part of the key so a deleted clip's name reused for another video never gets the old copy
        return hashlib.sha256(f"{name}\0{length}".encode()).hexdigest()[:24]

    def wanted(self, length: int) -> bool:
        return length > self.over_bytes

    def file(self, key: str) -> Path | None:
        path = self.dir / f"{key}.mp4"
        return path if KEY_RE.match(key) and path.is_file() else None

    def ready(self, name: str, length: int) -> str | None:
        """The key of a finished copy, or None (asking for one in the background if it's missing)."""
        key = self.key(name, length)
        if self.file(key):
            return key
        self.request(name, length)
        return None

    def status(self, name: str, length: int) -> dict:
        """For the site's upload dialog. state: "small" (plays as is), "ready" (copy made), "none" (too long, no
        copy) or "making" (Discord would still get the big original), with how far the encode is and whether
        it's still waiting behind another clip."""
        done = lambda state: {"state": state, "progress": 1.0, "queued": False}
        if not self.wanted(length):
            return done("small")
        key = self.key(name, length)
        if self.file(key):
            return done("ready")
        if (self.dir / f"{key}.failed").exists():
            return done("none")
        self.request(name, length)
        with self.lock:
            now = self.current == key
            return {"state": "making", "progress": self.progress if now else 0.0, "queued": not now}

    def renamed(self, old: str, new: str, length: int) -> None:
        """Carry a clip's copy over to its new name; without it a renamed clip's link sends Discord the big original."""
        if not self.wanted(length):
            return
        old_key, new_key = self.key(old, length), self.key(new, length)
        for ext in (".mp4", ".failed"):
            try:
                (self.dir / f"{old_key}{ext}").replace(self.dir / f"{new_key}{ext}")
            except OSError:
                pass
        self.request(new, length)  # still being made under the old name: start the new one now, not on the first view

    def request(self, name: str, length: int) -> None:
        key = self.key(name, length)
        if not self.wanted(length) or self.file(key) or (self.dir / f"{key}.failed").exists():
            return
        with self.lock:
            if key in self.pending:
                return
            self.pending.add(key)
            self.jobs.put((name, length))
            if self.worker is None or not self.worker.is_alive():
                self.worker = threading.Thread(target=self.run, name="previews", daemon=True)
                self.worker.start()

    def run(self) -> None:
        while True:
            try:
                name, length = self.jobs.get(timeout=60)
            except queue.Empty:
                with self.lock:
                    if self.jobs.empty():
                        self.worker = None
                        return
                continue
            key = self.key(name, length)
            with self.lock:
                self.current, self.progress = key, 0.0
            try:
                self.prune()
                self.make(name, key)
            except Exception:  # one bad clip must not kill the worker
                log.exception("preview for %s failed", name)
            finally:
                with self.lock:
                    self.pending.discard(key)
                    self.current = None

    def make(self, name: str, key: str) -> None:
        self.dir.mkdir(parents=True, exist_ok=True)
        src = f"{self.zipline_url}/{quote(name)}"
        duration = duration_of(src)
        kbps = video_kbps(duration)
        tmp, dst = self.dir / f"{key}.part.mp4", self.dir / f"{key}.mp4"
        ok = False
        if kbps:
            started = time.monotonic()
            try:
                proc = subprocess.Popen(niced(encode_cmd(src, tmp, kbps, duration)), stdout=subprocess.PIPE,
                                        stderr=subprocess.DEVNULL, text=True)
                watchdog = threading.Timer(3 * 3600, proc.kill)
                watchdog.start()
                try:
                    follow_progress(proc.stdout, duration, self.set_progress)
                    code = proc.wait()
                finally:
                    watchdog.cancel()
                ok = code == 0 and tmp.is_file() and tmp.stat().st_size < TARGET_BYTES * 13 // 10
            except OSError:
                ok = False
            log.info("preview %s for %s: %s in %.0fs", key, name, "ok" if ok else "failed", time.monotonic() - started)
        if ok:
            tmp.replace(dst)
        else:
            tmp.unlink(missing_ok=True)
            (self.dir / f"{key}.failed").touch()  # don't burn CPU on it again

    def set_progress(self, fraction: float) -> None:
        with self.lock:
            self.progress = fraction

    def prune(self) -> None:
        cutoff = time.time() - KEEP_DAYS * 86400
        for path in self.dir.glob("*"):
            try:
                if path.stat().st_mtime < cutoff:
                    path.unlink()
            except OSError:
                pass