LogalDeveloper/Owlbot
0
Fork 0
Code Issues 13 Actions Packages Releases 3 Wiki Activity

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

Browse Source
This commit is contained in:
LogalDeveloper
2026-03-26 10:57:51 -04:00
parent 47605c0d9b
commit 3cf93da28d
3 changed files with 256 additions and 209 deletions
Download Patch File Download Diff File Expand all files Collapse all files
owlbot/builtin_modules/clips/cache.py
+225 -150
View File
@@ -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."
)