Spaces:
Configuration error
Configuration error
| """ | |
| Base prebuffer class with shared functionality for HLS and DASH prebuffering. | |
| This module provides cross-process download coordination using file-based locking | |
| to prevent duplicate downloads across multiple uvicorn workers. Both player requests | |
| and background prebuffer tasks use the same coordination mechanism. | |
| """ | |
| import asyncio | |
| import logging | |
| import time | |
| import psutil | |
| from abc import ABC | |
| from dataclasses import dataclass, field | |
| from typing import Dict, Optional | |
| from mediaflow_proxy.utils.cache_utils import ( | |
| get_cached_segment, | |
| set_cached_segment, | |
| SEGMENT_DOWNLOAD_LOCK, | |
| ) | |
| from mediaflow_proxy.utils.http_utils import download_file_with_retry | |
| logger = logging.getLogger(__name__) | |
| class PrebufferStats: | |
| """Statistics for prebuffer performance tracking.""" | |
| cache_hits: int = 0 | |
| cache_misses: int = 0 | |
| segments_prebuffered: int = 0 | |
| bytes_prebuffered: int = 0 | |
| prefetch_triggered: int = 0 | |
| downloads_coordinated: int = 0 # Times we waited for existing download | |
| last_reset: float = field(default_factory=time.time) | |
| def hit_rate(self) -> float: | |
| """Calculate cache hit rate percentage.""" | |
| total = self.cache_hits + self.cache_misses | |
| return (self.cache_hits / total * 100) if total > 0 else 0.0 | |
| def reset(self) -> None: | |
| """Reset statistics.""" | |
| self.cache_hits = 0 | |
| self.cache_misses = 0 | |
| self.segments_prebuffered = 0 | |
| self.bytes_prebuffered = 0 | |
| self.prefetch_triggered = 0 | |
| self.downloads_coordinated = 0 | |
| self.last_reset = time.time() | |
| def to_dict(self) -> dict: | |
| """Convert stats to dictionary for logging.""" | |
| return { | |
| "cache_hits": self.cache_hits, | |
| "cache_misses": self.cache_misses, | |
| "hit_rate": f"{self.hit_rate:.1f}%", | |
| "segments_prebuffered": self.segments_prebuffered, | |
| "bytes_prebuffered_mb": f"{self.bytes_prebuffered / 1024 / 1024:.2f}", | |
| "prefetch_triggered": self.prefetch_triggered, | |
| "downloads_coordinated": self.downloads_coordinated, | |
| "uptime_seconds": int(time.time() - self.last_reset), | |
| } | |
| class BasePrebuffer(ABC): | |
| """ | |
| Base class for prebuffer systems with cross-process download coordination. | |
| This class provides: | |
| - Cross-process coordination using file-based locks to prevent duplicate downloads | |
| - Memory usage monitoring | |
| - Cache statistics tracking | |
| - Shared download and caching logic | |
| The file-based locking ensures that even with multiple uvicorn workers, | |
| only one worker downloads any given segment at a time. | |
| Subclasses should implement protocol-specific logic (HLS playlist parsing, | |
| DASH MPD handling, etc.) while inheriting the core download coordination. | |
| """ | |
| def __init__( | |
| self, | |
| max_cache_size: int, | |
| prebuffer_segments: int, | |
| max_memory_percent: float, | |
| emergency_threshold: float, | |
| segment_ttl: int = 60, | |
| ): | |
| """ | |
| Initialize the base prebuffer. | |
| Args: | |
| max_cache_size: Maximum number of segments to track | |
| prebuffer_segments: Number of segments to pre-buffer ahead | |
| max_memory_percent: Maximum memory usage percentage before skipping prebuffer | |
| emergency_threshold: Memory threshold for emergency cleanup | |
| segment_ttl: TTL for cached segments in seconds | |
| """ | |
| self.max_cache_size = max_cache_size | |
| self.prebuffer_segment_count = prebuffer_segments | |
| self.max_memory_percent = max_memory_percent | |
| self.emergency_threshold = emergency_threshold | |
| self.segment_ttl = segment_ttl | |
| # Cross-process lock for download coordination | |
| self._cross_process_lock = SEGMENT_DOWNLOAD_LOCK | |
| # Statistics (per-worker, not shared - but that's fine for monitoring) | |
| self.stats = PrebufferStats() | |
| # Stats logging task | |
| self._stats_task: Optional[asyncio.Task] = None | |
| self._stats_interval = 60 # Log stats every 60 seconds | |
| # Lock cleanup task | |
| self._cleanup_task: Optional[asyncio.Task] = None | |
| self._cleanup_interval = 300 # Cleanup stale locks every 5 minutes | |
| def _get_memory_usage_percent(self) -> float: | |
| """Get current memory usage percentage.""" | |
| try: | |
| memory = psutil.virtual_memory() | |
| return memory.percent | |
| except Exception as e: | |
| logger.warning(f"Failed to get memory usage: {e}") | |
| return 0.0 | |
| def _check_memory_threshold(self) -> bool: | |
| """Check if memory usage exceeds the emergency threshold.""" | |
| return self._get_memory_usage_percent() > self.emergency_threshold | |
| def _should_skip_for_memory(self) -> bool: | |
| """Check if we should skip prebuffering due to high memory usage.""" | |
| return self._get_memory_usage_percent() > self.max_memory_percent | |
| def record_cache_hit(self) -> None: | |
| """Record a cache hit for statistics.""" | |
| self.stats.cache_hits += 1 | |
| self._ensure_stats_logging() | |
| def record_cache_miss(self) -> None: | |
| """Record a cache miss for statistics.""" | |
| self.stats.cache_misses += 1 | |
| self._ensure_stats_logging() | |
| def _ensure_stats_logging(self) -> None: | |
| """Ensure the stats logging and cleanup tasks are running.""" | |
| if self._stats_task is None or self._stats_task.done(): | |
| self._stats_task = asyncio.create_task(self._periodic_stats_logging()) | |
| if self._cleanup_task is None or self._cleanup_task.done(): | |
| self._cleanup_task = asyncio.create_task(self._periodic_lock_cleanup()) | |
| async def _periodic_stats_logging(self) -> None: | |
| """Periodically log prebuffer statistics.""" | |
| while True: | |
| try: | |
| await asyncio.sleep(self._stats_interval) | |
| # Only log if there's been activity | |
| if self.stats.cache_hits > 0 or self.stats.cache_misses > 0: | |
| self.log_stats() | |
| except asyncio.CancelledError: | |
| return | |
| except Exception as e: | |
| logger.warning(f"Error in stats logging: {e}") | |
| async def _periodic_lock_cleanup(self) -> None: | |
| """Periodically clean up stale lock files.""" | |
| while True: | |
| try: | |
| await asyncio.sleep(self._cleanup_interval) | |
| await self._cross_process_lock.cleanup_stale_locks() | |
| except asyncio.CancelledError: | |
| return | |
| except Exception as e: | |
| logger.warning(f"Error in lock cleanup: {e}") | |
| async def get_or_download( | |
| self, | |
| url: str, | |
| headers: Dict[str, str], | |
| timeout: float = 10.0, | |
| ) -> Optional[bytes]: | |
| """ | |
| Get a segment from cache or download it, with cross-process coordination. | |
| This is the primary method for getting segments. It: | |
| 1. Checks cache first (immediate return if hit) | |
| 2. Acquires cross-process lock to prevent duplicate downloads across workers | |
| 3. Double-checks cache after acquiring lock | |
| 4. Downloads and caches if needed | |
| The file-based locking ensures that even with multiple uvicorn workers, | |
| only one worker downloads any given segment at a time. | |
| Args: | |
| url: URL of the segment to get | |
| headers: Headers to use for the request | |
| timeout: Maximum time to wait for lock acquisition (seconds). | |
| Keep this short (10s) for player requests - if lock is held | |
| too long, fall back to direct streaming. | |
| Returns: | |
| Segment data if successful, None if failed or timed out | |
| """ | |
| self._ensure_stats_logging() | |
| # Check cache first (file-based cache is shared across workers) | |
| cached = await get_cached_segment(url) | |
| if cached: | |
| self.record_cache_hit() | |
| logger.info(f"[get_or_download] CACHE HIT ({len(cached)} bytes): {url}") | |
| return cached | |
| # Cache miss - need to coordinate download across workers | |
| logger.info(f"[get_or_download] CACHE MISS: {url}") | |
| try: | |
| # Acquire cross-process lock - only one worker downloads at a time | |
| async with self._cross_process_lock.acquire(url, timeout=timeout): | |
| # Double-check cache after acquiring lock | |
| # Another worker may have completed the download while we waited | |
| cached = await get_cached_segment(url) | |
| if cached: | |
| # Count this as a cache hit since we didn't download | |
| self.record_cache_hit() | |
| self.stats.downloads_coordinated += 1 | |
| logger.info(f"[get_or_download] Found in cache after lock (coordinated): {url}") | |
| return cached | |
| # We're the one who needs to download - count as miss now | |
| self.record_cache_miss() | |
| # We're the first - download and cache | |
| logger.info(f"[get_or_download] Downloading: {url}") | |
| content = await self._download_and_cache(url, headers) | |
| return content | |
| except asyncio.TimeoutError: | |
| logger.warning(f"[get_or_download] Lock TIMEOUT ({timeout}s), falling back to streaming: {url}") | |
| return None | |
| except Exception as e: | |
| logger.warning(f"[get_or_download] Error during download coordination: {e}") | |
| return None | |
| async def _download_and_cache( | |
| self, | |
| url: str, | |
| headers: Dict[str, str], | |
| ) -> Optional[bytes]: | |
| """ | |
| Download a segment and cache it. | |
| This method should only be called while holding the cross-process lock. | |
| Args: | |
| url: URL to download | |
| headers: Headers for the request | |
| Returns: | |
| Downloaded content if successful, None otherwise | |
| """ | |
| try: | |
| content = await download_file_with_retry(url, headers) | |
| if content: | |
| logger.info(f"[_download_and_cache] Downloaded {len(content)} bytes, caching: {url}") | |
| await set_cached_segment(url, content, ttl=self.segment_ttl) | |
| self.stats.segments_prebuffered += 1 | |
| self.stats.bytes_prebuffered += len(content) | |
| return content | |
| else: | |
| logger.warning(f"[_download_and_cache] Download returned empty: {url}") | |
| return None | |
| except Exception as e: | |
| logger.warning(f"[_download_and_cache] Failed to download: {url} - {e}") | |
| return None | |
| async def try_get_cached(self, url: str) -> Optional[bytes]: | |
| """ | |
| Check cache only, don't download. | |
| Use this for background prebuffer tasks that shouldn't block | |
| if segment isn't available yet. | |
| Args: | |
| url: URL to check in cache | |
| Returns: | |
| Cached data if available, None otherwise | |
| """ | |
| return await get_cached_segment(url) | |
| async def prebuffer_segment(self, url: str, headers: Dict[str, str]) -> None: | |
| """ | |
| Prebuffer a single segment in the background. | |
| This method uses cross-process locking to prevent duplicate downloads | |
| across multiple workers. | |
| Args: | |
| url: URL of segment to prebuffer | |
| headers: Headers for the request | |
| """ | |
| if self._should_skip_for_memory(): | |
| logger.debug("Skipping prebuffer due to high memory usage") | |
| return | |
| # Check if already cached | |
| cached = await get_cached_segment(url) | |
| if cached: | |
| logger.debug(f"[prebuffer_segment] Already cached, skipping: {url}") | |
| return | |
| try: | |
| # Try to acquire lock with short timeout for prebuffering | |
| # If lock is held by another process, skip this segment | |
| async with self._cross_process_lock.acquire(url, timeout=1.0): | |
| # Double-check cache after acquiring lock | |
| cached = await get_cached_segment(url) | |
| if cached: | |
| logger.debug(f"[prebuffer_segment] Found in cache after lock: {url}") | |
| return | |
| # Download and cache | |
| logger.info(f"[prebuffer_segment] Downloading: {url}") | |
| await self._download_and_cache(url, headers) | |
| except asyncio.TimeoutError: | |
| # Another process is downloading, skip this segment | |
| logger.debug(f"[prebuffer_segment] Lock busy, skipping: {url}") | |
| except Exception as e: | |
| logger.warning(f"[prebuffer_segment] Error: {e}") | |
| async def prebuffer_segments_batch( | |
| self, | |
| urls: list, | |
| headers: Dict[str, str], | |
| max_concurrent: int = 2, | |
| ) -> None: | |
| """ | |
| Prebuffer multiple segments with concurrency control. | |
| Args: | |
| urls: List of segment URLs to prebuffer | |
| headers: Headers for requests | |
| max_concurrent: Maximum concurrent downloads (default 2 to avoid | |
| lock contention with player requests) | |
| """ | |
| if self._should_skip_for_memory(): | |
| logger.warning("Skipping prebuffer due to high memory usage") | |
| return | |
| semaphore = asyncio.Semaphore(max_concurrent) | |
| async def limited_prebuffer(url: str): | |
| async with semaphore: | |
| await self.prebuffer_segment(url, headers) | |
| # Start all prebuffer tasks | |
| tasks = [limited_prebuffer(url) for url in urls] | |
| await asyncio.gather(*tasks, return_exceptions=True) | |
| def log_stats(self) -> None: | |
| """Log current prebuffer statistics.""" | |
| logger.info(f"Prebuffer Stats: {self.stats.to_dict()}") | |