Refactored HLS cache to use set-based segment tracking and ffmpeg concat protocol.
CI / Formatting (push) Successful in 5s
CI / Linting (push) Successful in 5s
CI / Tests (Python 3.12) (push) Successful in 20s
CI / Tests (Python 3.13) (push) Successful in 19s
CI / Tests (Python 3.14) (push) Successful in 17s
CI / Type Checking (push) Successful in 9s
CI / Spelling (push) Successful in 7s

This commit is contained in:
2026-03-26 10:57:51 -04:00
parent 47605c0d9b
commit 3cf93da28d
3 changed files with 256 additions and 209 deletions
+225 -150
View File
@@ -22,11 +22,9 @@ from __future__ import annotations
import asyncio import asyncio
import contextlib import contextlib
import os
import shutil import shutil
import tempfile import tempfile
from collections import deque from contextlib import asynccontextmanager
from dataclasses import dataclass
from pathlib import Path from pathlib import Path
from typing import TYPE_CHECKING from typing import TYPE_CHECKING
from urllib.parse import urljoin from urllib.parse import urljoin
@@ -34,54 +32,60 @@ from urllib.parse import urljoin
import m3u8 # type: ignore[import-untyped] import m3u8 # type: ignore[import-untyped]
if TYPE_CHECKING: if TYPE_CHECKING:
from collections.abc import AsyncIterator
from owlbot.api import ModuleContext from owlbot.api import ModuleContext
from .types import get_state from .types import get_state
_PLAYLIST_MAX_RETRIES = 10
@dataclass _PLAYLIST_RETRY_DELAY = 3.0
class ChunkInfo:
"""Metadata for a single cached HLS segment."""
sequence: int
duration: float
path: Path
class ChunkCache: class ChunkCache:
"""Manages a rolling window of cached HLS stream segments. """Manages a rolling window of cached HLS stream segments.
Tracks downloaded chunks and prunes old ones when the total duration Uses a set to track downloaded sequence numbers. All segments in a variant
exceeds the configured cache window. 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. """Initialize the chunk cache.
:param cache_dir: Directory to store cached .ts files. :param cache_dir: Directory to store cached .ts files.
:param cache_duration: Maximum cache window in seconds. :param cache_duration: Maximum cache window in seconds.
:param min_clip_length: Minimum clip duration in seconds.
""" """
self._cache_dir = cache_dir self._cache_dir = cache_dir
self._cache_duration = cache_duration self._cache_duration = cache_duration
self._chunks: deque[ChunkInfo] = deque() self._min_clip_length = min_clip_length
self._total_duration = 0.0 self._sequences: set[int] = set()
self._target_duration: float = 0.0
self._prune_guard_count: int = 0
@property def __len__(self) -> int:
def chunks(self) -> list[ChunkInfo]: """Return the number of cached segments."""
"""Return the list of cached chunks.""" return len(self._sequences)
return list(self._chunks)
@property @property
def total_duration(self) -> float: def total_duration(self) -> float:
"""Return total duration of cached chunks in seconds.""" """Return total duration of cached chunks in seconds."""
return self._total_duration return len(self._sequences) * self._target_duration
@property @property
def last_sequence(self) -> int | None: def sequences(self) -> set[int]:
"""Return the sequence number of the most recent chunk, or None.""" """Return the set of cached sequence numbers."""
if not self._chunks: return self._sequences.copy()
return None
return self._chunks[-1].sequence 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: def segment_path(self, sequence: int) -> Path:
"""Return the file path for a segment by sequence number. """Return the file path for a segment by sequence number.
@@ -91,49 +95,84 @@ class ChunkCache:
""" """
return self._cache_dir / f"segment-{sequence}.ts" return self._cache_dir / f"segment-{sequence}.ts"
def add_chunk(self, sequence: int, duration: float, path: Path) -> None: def has_chunk(self, sequence: int) -> bool:
"""Add a downloaded chunk to the cache. """Check if a segment is in the cache.
:param sequence: HLS media sequence number. :param sequence: HLS media sequence number.
:param duration: Segment duration in seconds. :return: True if the segment is cached.
:param path: Path to the downloaded .ts file.
""" """
self._chunks.append(ChunkInfo(sequence=sequence, duration=duration, path=path)) return sequence in self._sequences
self._total_duration += duration
def prune(self) -> None: def add_chunk(self, sequence: int) -> None:
"""Remove oldest chunks until total duration is within the cache window.""" """Add a segment to the cache.
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 snapshot(self, work_dir: Path) -> tuple[list[Path], float]: :param sequence: HLS media sequence number.
"""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).
""" """
files: list[Path] = [] self._sequences.add(sequence)
duration = 0.0
for chunk in self._chunks: def prune(self) -> bool:
if not chunk.path.exists(): """Remove oldest segments until total duration is within the cache window.
continue
dest = work_dir / chunk.path.name
os.link(chunk.path, dest)
files.append(dest)
duration += chunk.duration
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: def cleanup(self) -> None:
"""Remove the cache directory and all its contents.""" """Remove the cache directory and all its contents."""
if self._cache_dir.exists(): if self._cache_dir.exists():
shutil.rmtree(self._cache_dir) shutil.rmtree(self._cache_dir)
self._chunks.clear() self._sequences.clear()
self._total_duration = 0.0
async def start_caching(ctx: ModuleContext) -> None: 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 owncast_url = ctx.owncast_client.base_url
master_url = f"{owncast_url.rstrip('/')}/hls/stream.m3u8" master_url = f"{owncast_url.rstrip('/')}/hls/stream.m3u8"
cache_duration = int(ctx.config.get("cache_duration")) 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_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 module_state.cache = cache
task = asyncio.create_task( task = asyncio.create_task(
@@ -182,44 +226,76 @@ async def stop_caching(ctx: ModuleContext) -> None:
ctx.logger.info("HLS cache stopped and cleaned up.") ctx.logger.info("HLS cache stopped and cleaned up.")
async def _resolve_variant_url(ctx: ModuleContext, master_url: str) -> str | None: async def _fetch_playlist(
"""Fetch and parse the master playlist, return the highest bandwidth variant URL. 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 ctx: The module context.
:param master_url: URL to the master m3u8 playlist. :param url: URL of the m3u8 playlist.
:return: URL to the highest bandwidth variant playlist, or None on failure. :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: try:
async with ctx.http.get(master_url) as resp: async with ctx.http.get(segment_url) as seg_resp:
if resp.status != 200: if seg_resp.status != 200:
ctx.logger.warning(f"Master playlist returned HTTP {resp.status}.") return False
return None data = await seg_resp.read()
content = await resp.text()
except Exception: except Exception:
ctx.logger.warning("Failed to fetch master playlist.", exc_info=True) ctx.logger.debug(f"Network error downloading segment {seq}.", exc_info=True)
return None 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: cache.add_chunk(seq)
ctx.logger.warning("No variant playlists found in master playlist.") ctx.logger.debug(f"Cached segment {seq} ({len(data)} bytes).")
return None return True
# 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
async def _polling_loop( async def _polling_loop(
@@ -233,30 +309,44 @@ async def _polling_loop(
:param cache: The ChunkCache instance to populate. :param cache: The ChunkCache instance to populate.
:param master_url: URL to the master m3u8 playlist. :param master_url: URL to the master m3u8 playlist.
""" """
max_retries = 10 master_playlist = await _fetch_playlist(ctx, master_url, "Master playlist")
variant_url: str | None = None if master_playlist is None or not master_playlist.playlists:
for attempt in range(1, max_retries + 1): ctx.logger.error("Could not fetch master playlist. Caching disabled.")
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."
)
module_state = get_state(ctx) module_state = get_state(ctx)
cache.cleanup() cache.cleanup()
module_state.cache = None module_state.cache = None
module_state.cache_task = None module_state.cache_task = None
return 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: while True:
try: try:
@@ -269,63 +359,48 @@ async def _polling_loop(
playlist = m3u8.loads(content) 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 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): for i, segment in enumerate(playlist.segments):
seq = media_sequence + i seq = media_sequence + i
# Skip already-downloaded segments. if cache.has_chunk(seq):
last = cache.last_sequence continue
if last is not None and seq <= last:
if seq in pending_retries:
continue continue
# Download the segment.
segment_url = urljoin(variant_url, segment.uri) segment_url = urljoin(variant_url, segment.uri)
try: if not await _download_segment(ctx, cache, segment_url, seq):
async with ctx.http.get(segment_url) as seg_resp: pending_retries.add(seq)
if seg_resp.status != 200:
ctx.logger.warning( ctx.logger.warning(
f"Segment {seq} returned HTTP {seg_resp.status}." f"Segment {seq} download failed, queued for retry."
)
continue
data = await seg_resp.read()
except Exception:
ctx.logger.warning(
f"Failed to download segment {seq}.", exc_info=True
)
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_prune_count = len(cache)
pre_count = len(cache.chunks) pruned = cache.prune()
cache.prune()
pruned = pre_count - len(cache.chunks)
if pruned: if pruned:
ctx.logger.debug( ctx.logger.debug(
f"Pruned {pruned} segment(s). " f"Pruned {pre_prune_count - len(cache)} segment(s). "
f"Cache: {len(cache.chunks)} chunks, " f"Cache: {len(cache)} chunks, "
f"{cache.total_duration:.1f}s total." f"{cache.total_duration:.1f}s total."
) )
+14 -31
View File
@@ -31,47 +31,36 @@ from .types import EditorSession, get_state
async def clip_command(ctx: CommandContext) -> None: async def clip_command(ctx: CommandContext) -> None:
"""Create a clip from the current stream. """Create a clip from the current stream.
Snapshots the HLS cache, generates a preview MP4, and sends the user Uses the concat protocol to read directly from the HLS cache while
a link to the clip editor. suppressing pruning, then sends the user a link to the clip editor.
:param ctx: The command context. :param ctx: The command context.
""" """
module = ctx.module module = ctx.module
module_state = get_state(module) module_state = get_state(module)
# Check if stream is live / cache is available. if module_state.cache is None:
min_clip_length = int(ctx.module.config.get("min_clip_length")) ctx.logger.debug("Clip command rejected: cache is None.")
min_cache_duration = min_clip_length * 2 await ctx.owncast_client.send_message(
if ( "Clipping is unavailable for this stream."
module_state.cache is None )
or module_state.cache.total_duration < min_cache_duration return
):
if not module_state.cache.enough_for_clipping:
ctx.logger.debug("Clip command rejected: not enough cached data.") ctx.logger.debug("Clip command rejected: not enough cached data.")
await ctx.owncast_client.send_message( await ctx.owncast_client.send_message(
"Not enough stream data is available yet. Please try again later." "Not enough stream data is available yet. Please try again later."
) )
return return
# Snapshot chunks into a working directory. cache = module_state.cache
work_dir = Path(tempfile.mkdtemp(prefix="owlbot-clip-work-")) 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" preview_path = work_dir / "preview.mp4"
try: 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: except Exception:
module.logger.error("Failed to generate clip preview.", exc_info=True) module.logger.error("Failed to generate clip preview.", exc_info=True)
shutil.rmtree(work_dir, ignore_errors=True) shutil.rmtree(work_dir, ignore_errors=True)
@@ -80,12 +69,6 @@ async def clip_command(ctx: CommandContext) -> None:
) )
return 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) token = secrets.token_urlsafe(32)
session_expiry = int(module.config.get("session_expiry")) session_expiry = int(module.config.get("session_expiry"))
session = EditorSession( session = EditorSession(
+15 -26
View File
@@ -27,6 +27,8 @@ if TYPE_CHECKING:
import logging import logging
from pathlib import Path from pathlib import Path
from .cache import ChunkCache
class ProcessingManager: class ProcessingManager:
"""Manages ffmpeg/ffprobe subprocess execution with concurrency control. """Manages ffmpeg/ffprobe subprocess execution with concurrency control.
@@ -44,55 +46,42 @@ class ProcessingManager:
self._ffmpeg_semaphore = asyncio.Semaphore(1) self._ffmpeg_semaphore = asyncio.Semaphore(1)
self._ffprobe_semaphore = asyncio.Semaphore(1) self._ffprobe_semaphore = asyncio.Semaphore(1)
async def generate_preview( async def generate_preview_from_cache(
self, self,
chunk_files: list[Path],
output_path: Path, output_path: Path,
cache: ChunkCache,
) -> float: ) -> 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 Suppresses cache pruning while building the concat string and running
semaphore for the duration probe. 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 output_path: Where to write the output MP4.
:param cache: The ChunkCache to read segments from.
:return: Duration of the generated preview in seconds. :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: async with cache.suppress_pruning():
msg = "ffmpeg: no chunk files provided for preview generation." concat_string = cache.concat_protocol_string()
if not concat_string:
msg = "ffmpeg: no cached segments available."
raise RuntimeError(msg) raise RuntimeError(msg)
concat_path = output_path.parent / "concat.txt"
async with self._ffmpeg_semaphore: async with self._ffmpeg_semaphore:
concat_path.write_text(
"\n".join(f"file '{f}'" for f in chunk_files),
encoding="utf-8",
)
try:
self._logger.debug( self._logger.debug(
f"Generating preview from {len(chunk_files)} chunks " f"Generating preview via concat protocol -> {output_path}."
f"-> {output_path}."
) )
await self._run_ffprocess( await self._run_ffprocess(
"ffmpeg", "ffmpeg",
"-y", "-y",
"-f",
"concat",
"-safe",
"0",
"-i", "-i",
str(concat_path), f"concat:{concat_string}",
"-c", "-c",
"copy", "copy",
"-movflags", "-movflags",
"+faststart", "+faststart",
str(output_path), str(output_path),
) )
finally:
if concat_path.exists():
concat_path.unlink()
async with self._ffprobe_semaphore: async with self._ffprobe_semaphore:
duration = await self._probe_duration(output_path) duration = await self._probe_duration(output_path)