Refactored HLS cache to use set-based segment tracking and ffmpeg concat protocol.

owlbot/builtin_modules/clips/cache.py

@@ -22,11 +22,9 @@ from __future__ import annotations
 
 import asyncio
 import contextlib
-import os
 import shutil
 import tempfile
-from collections import deque
-from dataclasses import dataclass
+from contextlib import asynccontextmanager
 from pathlib import Path
 from typing import TYPE_CHECKING
 from urllib.parse import urljoin
@@ -34,54 +32,60 @@ from urllib.parse import urljoin
 import m3u8  # type: ignore[import-untyped]
 
 if TYPE_CHECKING:
+    from collections.abc import AsyncIterator
+
     from owlbot.api import ModuleContext
 
 from .types import get_state
 
-
-@dataclass
-class ChunkInfo:
-    """Metadata for a single cached HLS segment."""
-
-    sequence: int
-    duration: float
-    path: Path
+_PLAYLIST_MAX_RETRIES = 10
+_PLAYLIST_RETRY_DELAY = 3.0
 
 
 class ChunkCache:
     """Manages a rolling window of cached HLS stream segments.
 
-    Tracks downloaded chunks and prunes old ones when the total duration
-    exceeds the configured cache window.
+    Uses a set to track downloaded sequence numbers. All segments in a variant
+    playlist share the same target duration, so total cache duration is
+    derived from the set size and the target duration.
     """
 
-    def __init__(self, cache_dir: Path, cache_duration: int) -> None:
+    def __init__(
+        self, cache_dir: Path, cache_duration: int, min_clip_length: int
+    ) -> None:
         """Initialize the chunk cache.
 
         :param cache_dir: Directory to store cached .ts files.
         :param cache_duration: Maximum cache window in seconds.
+        :param min_clip_length: Minimum clip duration in seconds.
         """
         self._cache_dir = cache_dir
         self._cache_duration = cache_duration
-        self._chunks: deque[ChunkInfo] = deque()
-        self._total_duration = 0.0
+        self._min_clip_length = min_clip_length
+        self._sequences: set[int] = set()
+        self._target_duration: float = 0.0
+        self._prune_guard_count: int = 0
 
-    @property
-    def chunks(self) -> list[ChunkInfo]:
-        """Return the list of cached chunks."""
-        return list(self._chunks)
+    def __len__(self) -> int:
+        """Return the number of cached segments."""
+        return len(self._sequences)
 
     @property
     def total_duration(self) -> float:
         """Return total duration of cached chunks in seconds."""
-        return self._total_duration
+        return len(self._sequences) * self._target_duration
 
     @property
-    def last_sequence(self) -> int | None:
-        """Return the sequence number of the most recent chunk, or None."""
-        if not self._chunks:
-            return None
-        return self._chunks[-1].sequence
+    def sequences(self) -> set[int]:
+        """Return the set of cached sequence numbers."""
+        return self._sequences.copy()
+
+    def set_target_duration(self, duration: float) -> None:
+        """Set the target segment duration from the HLS playlist.
+
+        :param duration: Target segment duration in seconds.
+        """
+        self._target_duration = duration
 
     def segment_path(self, sequence: int) -> Path:
         """Return the file path for a segment by sequence number.
@@ -91,49 +95,84 @@ class ChunkCache:
         """
         return self._cache_dir / f"segment-{sequence}.ts"
 
-    def add_chunk(self, sequence: int, duration: float, path: Path) -> None:
-        """Add a downloaded chunk to the cache.
+    def has_chunk(self, sequence: int) -> bool:
+        """Check if a segment is in the cache.
 
         :param sequence: HLS media sequence number.
-        :param duration: Segment duration in seconds.
-        :param path: Path to the downloaded .ts file.
+        :return: True if the segment is cached.
         """
-        self._chunks.append(ChunkInfo(sequence=sequence, duration=duration, path=path))
-        self._total_duration += duration
+        return sequence in self._sequences
 
-    def prune(self) -> None:
-        """Remove oldest chunks until total duration is within the cache window."""
-        while self._total_duration > self._cache_duration and len(self._chunks) > 1:
-            old = self._chunks.popleft()
-            self._total_duration -= old.duration
-            if old.path.exists():
-                old.path.unlink()
+    def add_chunk(self, sequence: int) -> None:
+        """Add a segment to the cache.
 
-    def snapshot(self, work_dir: Path) -> tuple[list[Path], float]:
-        """Hardlink current chunks into a working directory.
-
-        :param work_dir: Directory to hardlink files into.
-        :return: Tuple of (list of hardlinked file paths, total duration).
+        :param sequence: HLS media sequence number.
         """
-        files: list[Path] = []
-        duration = 0.0
+        self._sequences.add(sequence)
 
-        for chunk in self._chunks:
-            if not chunk.path.exists():
-                continue
-            dest = work_dir / chunk.path.name
-            os.link(chunk.path, dest)
-            files.append(dest)
-            duration += chunk.duration
+    def prune(self) -> bool:
+        """Remove oldest segments until total duration is within the cache window.
 
-        return files, duration
+        :return: True if any segments were pruned, False otherwise.
+        """
+        if self._prune_guard_count > 0:
+            return False
+        pruned = False
+        while self.total_duration > self._cache_duration and len(self._sequences) > 1:
+            oldest = min(self._sequences)
+            self._sequences.discard(oldest)
+            path = self.segment_path(oldest)
+            if path.exists():
+                path.unlink()
+            pruned = True
+        return pruned
+
+    @asynccontextmanager
+    async def suppress_pruning(self) -> AsyncIterator[None]:
+        """Context manager that suppresses pruning while reading from the cache.
+
+        Use as:
+            async with cache.suppress_pruning():
+                # ffmpeg reads from cache here
+
+        Supports nested calls via reference counting.
+        """
+        self._prune_guard_count += 1
+        try:
+            yield
+        finally:
+            self._prune_guard_count -= 1
+
+    @property
+    def enough_for_clipping(self) -> bool:
+        """Check if cached data is sufficient for clipping.
+
+        Requires at least 2x the configured min clip length to ensure enough
+        buffered data.
+
+        :return: True if total cached duration >= 2 * min_clip_length.
+        """
+        return self.total_duration >= self._min_clip_length * 2
+
+    def sorted_sequences(self) -> list[int]:
+        """Return cached sequence numbers in ascending order.
+
+        :return: Sorted list of sequence numbers.
+        """
+        return sorted(self._sequences)
+
+    def concat_protocol_string(self) -> str:
+        """Build the ffmpeg concat protocol input string.
+
+        :return: Pipe-separated list of cached segment paths.
+        """
+        return "|".join(str(self.segment_path(seq)) for seq in self.sorted_sequences())
 
     def cleanup(self) -> None:
         """Remove the cache directory and all its contents."""
         if self._cache_dir.exists():
             shutil.rmtree(self._cache_dir)
-        self._chunks.clear()
-        self._total_duration = 0.0
+        self._sequences.clear()
 
 
 async def start_caching(ctx: ModuleContext) -> None:
@@ -152,9 +191,14 @@ async def start_caching(ctx: ModuleContext) -> None:
     owncast_url = ctx.owncast_client.base_url
     master_url = f"{owncast_url.rstrip('/')}/hls/stream.m3u8"
     cache_duration = int(ctx.config.get("cache_duration"))
+    min_clip_length = int(ctx.config.get("min_clip_length"))
     cache_dir_path = Path(tempfile.mkdtemp(prefix="owlbot-clips-"))
 
-    cache = ChunkCache(cache_dir=cache_dir_path, cache_duration=cache_duration)
+    cache = ChunkCache(
+        cache_dir=cache_dir_path,
+        cache_duration=cache_duration,
+        min_clip_length=min_clip_length,
+    )
     module_state.cache = cache
 
     task = asyncio.create_task(
@@ -182,44 +226,76 @@ async def stop_caching(ctx: ModuleContext) -> None:
         ctx.logger.info("HLS cache stopped and cleaned up.")
 
 
-async def _resolve_variant_url(ctx: ModuleContext, master_url: str) -> str | None:
-    """Fetch and parse the master playlist, return the highest bandwidth variant URL.
+async def _fetch_playlist(
+    ctx: ModuleContext,
+    url: str,
+    label: str,
+) -> m3u8.M3U8 | None:
+    """Fetch and parse an HLS playlist with retry logic.
+
+    Attempts up to `_PLAYLIST_MAX_RETRIES` fetches with a
+    `_PLAYLIST_RETRY_DELAY`-second backoff between failures.
+    Propagates `~asyncio.CancelledError` immediately.
 
     :param ctx: The module context.
-    :param master_url: URL to the master m3u8 playlist.
-    :return: URL to the highest bandwidth variant playlist, or None on failure.
+    :param url: URL of the m3u8 playlist.
+    :param label: Human-readable label for log messages.
+    :return: Parsed playlist, or ``None`` if all attempts are exhausted.
+    """
+    for attempt in range(1, _PLAYLIST_MAX_RETRIES + 1):
+        reason: str | None = None
+        exc_info = False
+        try:
+            async with ctx.http.get(url) as resp:
+                if resp.status != 200:
+                    reason = f"HTTP {resp.status}"
+                else:
+                    content = await resp.text()
+                    return m3u8.loads(content)
+        except asyncio.CancelledError:
+            raise
+        except Exception as e:
+            reason = str(e)
+            exc_info = True
+        ctx.logger.warning(
+            f"{label} fetch failed (attempt {attempt}/{_PLAYLIST_MAX_RETRIES}): "
+            f"{reason}. Retrying in {_PLAYLIST_RETRY_DELAY:.0f}s.",
+            exc_info=exc_info,
+        )
+        await asyncio.sleep(_PLAYLIST_RETRY_DELAY)
+    ctx.logger.error(f"{label} failed after {_PLAYLIST_MAX_RETRIES} attempts.")
+    return None
+
+
+async def _download_segment(
+    ctx: ModuleContext,
+    cache: ChunkCache,
+    segment_url: str,
+    seq: int,
+) -> bool:
+    """Download a single HLS segment and add it to the cache.
+
+    :param ctx: The module context.
+    :param cache: The ChunkCache to add the segment to.
+    :param segment_url: URL to download the .ts segment from.
+    :param seq: HLS media sequence number.
+    :return: True if the segment was downloaded successfully, False otherwise.
     """
     try:
-        async with ctx.http.get(master_url) as resp:
-            if resp.status != 200:
-                ctx.logger.warning(f"Master playlist returned HTTP {resp.status}.")
-                return None
-            content = await resp.text()
+        async with ctx.http.get(segment_url) as seg_resp:
+            if seg_resp.status != 200:
+                return False
+            data = await seg_resp.read()
     except Exception:
-        ctx.logger.warning("Failed to fetch master playlist.", exc_info=True)
-        return None
+        ctx.logger.debug(f"Network error downloading segment {seq}.", exc_info=True)
+        return False
 
-    playlist = m3u8.loads(content)
+    chunk_path = cache.segment_path(seq)
+    await asyncio.to_thread(chunk_path.write_bytes, data)
 
-    if not playlist.playlists:
-        ctx.logger.warning("No variant playlists found in master playlist.")
-        return None
-
-    # Select highest bandwidth variant.
-    best = max(
-        playlist.playlists,
-        key=lambda p: p.stream_info.bandwidth or 0,
-    )
-
-    variant_uri: str = best.uri
-    # Resolve relative URI against master URL.
-    variant_url = urljoin(master_url, variant_uri)
-    ctx.logger.debug(
-        f"Selected variant: {variant_url} "
-        f"(bandwidth={best.stream_info.bandwidth}, "
-        f"resolution={best.stream_info.resolution})"
-    )
-    return variant_url
+    cache.add_chunk(seq)
+    ctx.logger.debug(f"Cached segment {seq} ({len(data)} bytes).")
+    return True
 
 
 async def _polling_loop(
@@ -233,30 +309,44 @@ async def _polling_loop(
     :param cache: The ChunkCache instance to populate.
     :param master_url: URL to the master m3u8 playlist.
     """
-    max_retries = 10
-    variant_url: str | None = None
-    for attempt in range(1, max_retries + 1):
-        variant_url = await _resolve_variant_url(ctx, master_url)
-        if variant_url is not None:
-            break
-        ctx.logger.warning(
-            f"Variant URL resolution failed (attempt {attempt}/{max_retries}). "
-            f"Retrying in 3s."
-        )
-        await asyncio.sleep(3.0)
-
-    if variant_url is None:
-        ctx.logger.error(
-            f"Could not resolve variant URL after {max_retries} attempts. "
-            f"Caching disabled."
-        )
+    master_playlist = await _fetch_playlist(ctx, master_url, "Master playlist")
+    if master_playlist is None or not master_playlist.playlists:
+        ctx.logger.error("Could not fetch master playlist. Caching disabled.")
         module_state = get_state(ctx)
         cache.cleanup()
         module_state.cache = None
         module_state.cache_task = None
         return
 
-    poll_interval = 2.0  # Default, updated from playlist
+    best = max(
+        master_playlist.playlists,
+        key=lambda p: p.stream_info.bandwidth or 0,
+    )
+    variant_url = urljoin(master_url, best.uri)
+    ctx.logger.debug(
+        f"Selected variant: {variant_url} "
+        f"(bandwidth={best.stream_info.bandwidth}, "
+        f"resolution={best.stream_info.resolution})"
+    )
+
+    variant_playlist = await _fetch_playlist(ctx, variant_url, "Variant playlist")
+    if variant_playlist is None:
+        ctx.logger.error("Could not fetch variant playlist. Caching disabled.")
+        module_state = get_state(ctx)
+        cache.cleanup()
+        module_state.cache = None
+        module_state.cache_task = None
+        return
+
+    poll_interval = (
+        float(variant_playlist.target_duration)
+        if variant_playlist.target_duration
+        else 2.0
+    )
+    cache.set_target_duration(poll_interval)
+    ctx.logger.debug(f"Initial poll interval set to {poll_interval:.1f}s.")
+
+    pending_retries: set[int] = set()
 
     while True:
         try:
@@ -269,63 +359,48 @@ async def _polling_loop(
 
             playlist = m3u8.loads(content)
 
-            # Update poll interval from target duration.
-            if playlist.target_duration:
-                new_interval = float(playlist.target_duration)
-                if new_interval != poll_interval:
-                    ctx.logger.debug(
-                        f"Poll interval updated: {poll_interval:.1f}s -> "
-                        f"{new_interval:.1f}s."
-                    )
-                    poll_interval = new_interval
-
             media_sequence = playlist.media_sequence or 0
 
+            playlist_sequences: dict[int, str] = {}
+            for i, segment in enumerate(playlist.segments):
+                playlist_sequences[media_sequence + i] = segment.uri
+
+            abandoned: set[int] = set()
+            for seq in list(pending_retries):
+                if seq not in playlist_sequences:
+                    abandoned.add(seq)
+                    continue
+                segment_url = urljoin(variant_url, playlist_sequences[seq])
+                if await _download_segment(ctx, cache, segment_url, seq):
+                    pending_retries.discard(seq)
+                    ctx.logger.info(f"Segment {seq} recovered after retry.")
+
+            for seq in abandoned:
+                pending_retries.discard(seq)
+                ctx.logger.warning(f"Segment {seq} no longer in playlist, giving up.")
+
             for i, segment in enumerate(playlist.segments):
                 seq = media_sequence + i
 
-                # Skip already-downloaded segments.
-                last = cache.last_sequence
-                if last is not None and seq <= last:
+                if cache.has_chunk(seq):
+                    continue
+
+                if seq in pending_retries:
                     continue
 
-                # Download the segment.
                 segment_url = urljoin(variant_url, segment.uri)
-                try:
-                    async with ctx.http.get(segment_url) as seg_resp:
-                        if seg_resp.status != 200:
-                            ctx.logger.warning(
-                                f"Segment {seq} returned HTTP {seg_resp.status}."
-                            )
-                            continue
-                        data = await seg_resp.read()
-                except Exception:
+                if not await _download_segment(ctx, cache, segment_url, seq):
+                    pending_retries.add(seq)
                     ctx.logger.warning(
-                        f"Failed to download segment {seq}.", exc_info=True
+                        f"Segment {seq} download failed, queued for retry."
                     )
-                    continue
 
-                chunk_path = cache.segment_path(seq)
-                await asyncio.to_thread(chunk_path.write_bytes, data)
-
-                cache.add_chunk(
-                    sequence=seq,
-                    duration=segment.duration,
-                    path=chunk_path,
-                )
-                ctx.logger.debug(
-                    f"Cached segment {seq} "
-                    f"({segment.duration:.2f}s, {len(data)} bytes)."
-                )
-
-            # Prune old segments.
-            pre_count = len(cache.chunks)
-            cache.prune()
-            pruned = pre_count - len(cache.chunks)
+            pre_prune_count = len(cache)
+            pruned = cache.prune()
             if pruned:
                 ctx.logger.debug(
-                    f"Pruned {pruned} segment(s). "
-                    f"Cache: {len(cache.chunks)} chunks, "
+                    f"Pruned {pre_prune_count - len(cache)} segment(s). "
+                    f"Cache: {len(cache)} chunks, "
                     f"{cache.total_duration:.1f}s total."
                 )
 

owlbot/builtin_modules/clips/commands.py

@@ -31,47 +31,36 @@ from .types import EditorSession, get_state
 async def clip_command(ctx: CommandContext) -> None:
     """Create a clip from the current stream.
 
-    Snapshots the HLS cache, generates a preview MP4, and sends the user
-    a link to the clip editor.
+    Uses the concat protocol to read directly from the HLS cache while
+    suppressing pruning, then sends the user a link to the clip editor.
 
     :param ctx: The command context.
     """
     module = ctx.module
     module_state = get_state(module)
 
-    # Check if stream is live / cache is available.
-    min_clip_length = int(ctx.module.config.get("min_clip_length"))
-    min_cache_duration = min_clip_length * 2
-    if (
-        module_state.cache is None
-        or module_state.cache.total_duration < min_cache_duration
-    ):
+    if module_state.cache is None:
+        ctx.logger.debug("Clip command rejected: cache is None.")
+        await ctx.owncast_client.send_message(
+            "Clipping is unavailable for this stream."
+        )
+        return
+
+    if not module_state.cache.enough_for_clipping:
         ctx.logger.debug("Clip command rejected: not enough cached data.")
         await ctx.owncast_client.send_message(
             "Not enough stream data is available yet. Please try again later."
         )
         return
 
-    # Snapshot chunks into a working directory.
+    cache = module_state.cache
+
     work_dir = Path(tempfile.mkdtemp(prefix="owlbot-clip-work-"))
-    snapshot_files, snapshot_duration = module_state.cache.snapshot(work_dir)
-
-    ctx.logger.debug(
-        f"Snapshot for {ctx.user.display_name}: {len(snapshot_files)} chunks, "
-        f"{snapshot_duration:.1f}s into {work_dir}."
-    )
-
-    if not snapshot_files:
-        shutil.rmtree(work_dir, ignore_errors=True)
-        await ctx.owncast_client.send_message("No stream data available for clipping.")
-        return
-
-    # Generate preview and send editor link.
-    manager = module_state.manager
     preview_path = work_dir / "preview.mp4"
 
     try:
-        duration = await manager.generate_preview(snapshot_files, preview_path)
+        manager = module_state.manager
+        duration = await manager.generate_preview_from_cache(preview_path, cache)
     except Exception:
         module.logger.error("Failed to generate clip preview.", exc_info=True)
         shutil.rmtree(work_dir, ignore_errors=True)
@@ -80,12 +69,6 @@ async def clip_command(ctx: CommandContext) -> None:
         )
         return
 
-    # Clean up hardlinked chunks (preview is self-contained now).
-    for f in snapshot_files:
-        if f.exists():
-            f.unlink()
-
-    # Create editor session.
     token = secrets.token_urlsafe(32)
     session_expiry = int(module.config.get("session_expiry"))
     session = EditorSession(

owlbot/builtin_modules/clips/processing.py

@@ -27,6 +27,8 @@ if TYPE_CHECKING:
     import logging
     from pathlib import Path
 
+    from .cache import ChunkCache
+
 
 class ProcessingManager:
     """Manages ffmpeg/ffprobe subprocess execution with concurrency control.
@@ -44,55 +46,42 @@ class ProcessingManager:
         self._ffmpeg_semaphore = asyncio.Semaphore(1)
         self._ffprobe_semaphore = asyncio.Semaphore(1)
 
-    async def generate_preview(
+    async def generate_preview_from_cache(
         self,
-        chunk_files: list[Path],
         output_path: Path,
+        cache: ChunkCache,
     ) -> float:
-        """Generate a preview MP4 from a list of HLS chunk files.
+        """Generate a preview MP4 from the cached HLS segments.
 
-        Acquires the ffmpeg semaphore for the concat step, then the ffprobe
-        semaphore for the duration probe.
+        Suppresses cache pruning while building the concat string and running
+        ffmpeg, then probes the output duration via ffprobe.
 
-        :param chunk_files: Ordered list of .ts chunk file paths.
         :param output_path: Where to write the output MP4.
+        :param cache: The ChunkCache to read segments from.
         :return: Duration of the generated preview in seconds.
-        :raises RuntimeError: If ffmpeg fails or no chunks provided.
+        :raises RuntimeError: If no segments are cached.
         """
-        if not chunk_files:
-            msg = "ffmpeg: no chunk files provided for preview generation."
-            raise RuntimeError(msg)
+        async with cache.suppress_pruning():
+            concat_string = cache.concat_protocol_string()
+            if not concat_string:
+                msg = "ffmpeg: no cached segments available."
+                raise RuntimeError(msg)
 
-        concat_path = output_path.parent / "concat.txt"
-
-        async with self._ffmpeg_semaphore:
-            concat_path.write_text(
-                "\n".join(f"file '{f}'" for f in chunk_files),
-                encoding="utf-8",
-            )
-            try:
+            async with self._ffmpeg_semaphore:
                 self._logger.debug(
-                    f"Generating preview from {len(chunk_files)} chunks "
-                    f"-> {output_path}."
+                    f"Generating preview via concat protocol -> {output_path}."
                 )
                 await self._run_ffprocess(
                     "ffmpeg",
                     "-y",
-                    "-f",
-                    "concat",
-                    "-safe",
-                    "0",
                     "-i",
-                    str(concat_path),
+                    f"concat:{concat_string}",
                     "-c",
                     "copy",
                     "-movflags",
                     "+faststart",
                     str(output_path),
                 )
-            finally:
-                if concat_path.exists():
-                    concat_path.unlink()
 
         async with self._ffprobe_semaphore:
             duration = await self._probe_duration(output_path)