Blame for mediaflow_proxy/utils/base_prebuffer.py

"""
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__)


@dataclass
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)

    @property
    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()}")