updated to lastest version

mediaflow_proxy/utils/dash_prebuffer.py

@@ -0,0 +1,373 @@
+import logging
+import psutil
+from typing import Dict, Optional, List
+from urllib.parse import urljoin
+import xmltodict
+from mediaflow_proxy.utils.http_utils import create_httpx_client
+from mediaflow_proxy.configs import settings
+
+logger = logging.getLogger(__name__)
+
+
+class DASHPreBuffer:
+    """
+    Pre-buffer system for DASH streams to reduce latency and improve streaming performance.
+    """
+    
+    def __init__(self, max_cache_size: Optional[int] = None, prebuffer_segments: Optional[int] = None):
+        """
+        Initialize the DASH pre-buffer system.
+        
+        Args:
+            max_cache_size (int): Maximum number of segments to cache (uses config if None)
+            prebuffer_segments (int): Number of segments to pre-buffer ahead (uses config if None)
+        """
+        self.max_cache_size = max_cache_size or settings.dash_prebuffer_cache_size
+        self.prebuffer_segments = prebuffer_segments or settings.dash_prebuffer_segments
+        self.max_memory_percent = settings.dash_prebuffer_max_memory_percent
+        self.emergency_threshold = settings.dash_prebuffer_emergency_threshold
+        
+        # Cache for different types of DASH content
+        self.segment_cache: Dict[str, bytes] = {}
+        self.init_segment_cache: Dict[str, bytes] = {}
+        self.manifest_cache: Dict[str, dict] = {}
+        
+        # Track segment URLs for each adaptation set
+        self.adaptation_segments: Dict[str, List[str]] = {}
+        self.client = create_httpx_client()
+    
+    def _get_memory_usage_percent(self) -> float:
+        """
+        Get current memory usage percentage.
+        
+        Returns:
+            float: 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.
+        
+        Returns:
+            bool: True if emergency cleanup is needed
+        """
+        memory_percent = self._get_memory_usage_percent()
+        return memory_percent > self.emergency_threshold
+    
+    def _emergency_cache_cleanup(self) -> None:
+        """
+        Perform emergency cache cleanup when memory usage is high.
+        """
+        if self._check_memory_threshold():
+            logger.warning("Emergency DASH cache cleanup triggered due to high memory usage")
+            
+            # Clear 50% of segment cache
+            segment_cache_size = len(self.segment_cache)
+            segment_keys_to_remove = list(self.segment_cache.keys())[:segment_cache_size // 2]
+            for key in segment_keys_to_remove:
+                del self.segment_cache[key]
+            
+            # Clear 50% of init segment cache
+            init_cache_size = len(self.init_segment_cache)
+            init_keys_to_remove = list(self.init_segment_cache.keys())[:init_cache_size // 2]
+            for key in init_keys_to_remove:
+                del self.init_segment_cache[key]
+            
+            logger.info(f"Emergency cleanup removed {len(segment_keys_to_remove)} segments and {len(init_keys_to_remove)} init segments from cache")
+    
+    async def prebuffer_dash_manifest(self, mpd_url: str, headers: Dict[str, str]) -> None:
+        """
+        Pre-buffer segments from a DASH manifest.
+        
+        Args:
+            mpd_url (str): URL of the DASH manifest
+            headers (Dict[str, str]): Headers to use for requests
+        """
+        try:
+            # Download and parse MPD manifest
+            response = await self.client.get(mpd_url, headers=headers)
+            response.raise_for_status()
+            mpd_content = response.text
+            
+            # Parse MPD XML
+            mpd_dict = xmltodict.parse(mpd_content)
+            
+            # Store manifest in cache
+            self.manifest_cache[mpd_url] = mpd_dict
+            
+            # Extract initialization segments and first few segments
+            await self._extract_and_prebuffer_segments(mpd_dict, mpd_url, headers)
+            
+            logger.info(f"Pre-buffered DASH manifest: {mpd_url}")
+            
+        except Exception as e:
+            logger.warning(f"Failed to pre-buffer DASH manifest {mpd_url}: {e}")
+    
+    async def _extract_and_prebuffer_segments(self, mpd_dict: dict, base_url: str, headers: Dict[str, str]) -> None:
+        """
+        Extract and pre-buffer segments from MPD manifest.
+        
+        Args:
+            mpd_dict (dict): Parsed MPD manifest
+            base_url (str): Base URL for resolving relative URLs
+            headers (Dict[str, str]): Headers to use for requests
+        """
+        try:
+            # Extract Period and AdaptationSet information
+            mpd = mpd_dict.get('MPD', {})
+            periods = mpd.get('Period', [])
+            if not isinstance(periods, list):
+                periods = [periods]
+            
+            for period in periods:
+                adaptation_sets = period.get('AdaptationSet', [])
+                if not isinstance(adaptation_sets, list):
+                    adaptation_sets = [adaptation_sets]
+                
+                for adaptation_set in adaptation_sets:
+                    # Extract initialization segment
+                    init_segment = adaptation_set.get('SegmentTemplate', {}).get('@initialization')
+                    if init_segment:
+                        init_url = urljoin(base_url, init_segment)
+                        await self._download_init_segment(init_url, headers)
+                    
+                    # Extract segment template
+                    segment_template = adaptation_set.get('SegmentTemplate', {})
+                    if segment_template:
+                        await self._prebuffer_template_segments(segment_template, base_url, headers)
+                    
+                    # Extract segment list
+                    segment_list = adaptation_set.get('SegmentList', {})
+                    if segment_list:
+                        await self._prebuffer_list_segments(segment_list, base_url, headers)
+                        
+        except Exception as e:
+            logger.warning(f"Failed to extract segments from MPD: {e}")
+    
+    async def _download_init_segment(self, init_url: str, headers: Dict[str, str]) -> None:
+        """
+        Download and cache initialization segment.
+        
+        Args:
+            init_url (str): URL of the initialization segment
+            headers (Dict[str, str]): Headers to use for request
+        """
+        try:
+            # Check memory usage before downloading
+            memory_percent = self._get_memory_usage_percent()
+            if memory_percent > self.max_memory_percent:
+                logger.warning(f"Memory usage {memory_percent}% exceeds limit {self.max_memory_percent}%, skipping init segment download")
+                return
+            
+            response = await self.client.get(init_url, headers=headers)
+            response.raise_for_status()
+            
+            # Cache the init segment
+            self.init_segment_cache[init_url] = response.content
+            
+            # Check for emergency cleanup
+            if self._check_memory_threshold():
+                self._emergency_cache_cleanup()
+            
+            logger.debug(f"Cached init segment: {init_url}")
+            
+        except Exception as e:
+            logger.warning(f"Failed to download init segment {init_url}: {e}")
+    
+    async def _prebuffer_template_segments(self, segment_template: dict, base_url: str, headers: Dict[str, str]) -> None:
+        """
+        Pre-buffer segments using segment template.
+        
+        Args:
+            segment_template (dict): Segment template from MPD
+            base_url (str): Base URL for resolving relative URLs
+            headers (Dict[str, str]): Headers to use for requests
+        """
+        try:
+            media_template = segment_template.get('@media')
+            if not media_template:
+                return
+            
+            # Extract template parameters
+            start_number = int(segment_template.get('@startNumber', 1))
+            duration = float(segment_template.get('@duration', 0))
+            timescale = float(segment_template.get('@timescale', 1))
+            
+            # Pre-buffer first few segments
+            for i in range(self.prebuffer_segments):
+                segment_number = start_number + i
+                segment_url = media_template.replace('$Number$', str(segment_number))
+                full_url = urljoin(base_url, segment_url)
+                
+                await self._download_segment(full_url, headers)
+                
+        except Exception as e:
+            logger.warning(f"Failed to pre-buffer template segments: {e}")
+    
+    async def _prebuffer_list_segments(self, segment_list: dict, base_url: str, headers: Dict[str, str]) -> None:
+        """
+        Pre-buffer segments from segment list.
+        
+        Args:
+            segment_list (dict): Segment list from MPD
+            base_url (str): Base URL for resolving relative URLs
+            headers (Dict[str, str]): Headers to use for requests
+        """
+        try:
+            segments = segment_list.get('SegmentURL', [])
+            if not isinstance(segments, list):
+                segments = [segments]
+            
+            # Pre-buffer first few segments
+            for segment in segments[:self.prebuffer_segments]:
+                segment_url = segment.get('@src')
+                if segment_url:
+                    full_url = urljoin(base_url, segment_url)
+                    await self._download_segment(full_url, headers)
+                    
+        except Exception as e:
+            logger.warning(f"Failed to pre-buffer list segments: {e}")
+    
+    async def _download_segment(self, segment_url: str, headers: Dict[str, str]) -> None:
+        """
+        Download a single segment and cache it.
+        
+        Args:
+            segment_url (str): URL of the segment to download
+            headers (Dict[str, str]): Headers to use for request
+        """
+        try:
+            # Check memory usage before downloading
+            memory_percent = self._get_memory_usage_percent()
+            if memory_percent > self.max_memory_percent:
+                logger.warning(f"Memory usage {memory_percent}% exceeds limit {self.max_memory_percent}%, skipping segment download")
+                return
+            
+            response = await self.client.get(segment_url, headers=headers)
+            response.raise_for_status()
+            
+            # Cache the segment
+            self.segment_cache[segment_url] = response.content
+            
+            # Check for emergency cleanup
+            if self._check_memory_threshold():
+                self._emergency_cache_cleanup()
+            # Maintain cache size
+            elif len(self.segment_cache) > self.max_cache_size:
+                # Remove oldest entries (simple FIFO)
+                oldest_key = next(iter(self.segment_cache))
+                del self.segment_cache[oldest_key]
+                
+            logger.debug(f"Cached DASH segment: {segment_url}")
+            
+        except Exception as e:
+            logger.warning(f"Failed to download DASH segment {segment_url}: {e}")
+    
+    async def get_segment(self, segment_url: str, headers: Dict[str, str]) -> Optional[bytes]:
+        """
+        Get a segment from cache or download it.
+        
+        Args:
+            segment_url (str): URL of the segment
+            headers (Dict[str, str]): Headers to use for request
+            
+        Returns:
+            Optional[bytes]: Cached segment data or None if not available
+        """
+        # Check segment cache first
+        if segment_url in self.segment_cache:
+            logger.debug(f"DASH cache hit for segment: {segment_url}")
+            return self.segment_cache[segment_url]
+        
+        # Check init segment cache
+        if segment_url in self.init_segment_cache:
+            logger.debug(f"DASH cache hit for init segment: {segment_url}")
+            return self.init_segment_cache[segment_url]
+        
+        # Check memory usage before downloading
+        memory_percent = self._get_memory_usage_percent()
+        if memory_percent > self.max_memory_percent:
+            logger.warning(f"Memory usage {memory_percent}% exceeds limit {self.max_memory_percent}%, skipping download")
+            return None
+        
+        # Download if not in cache
+        try:
+            response = await self.client.get(segment_url, headers=headers)
+            response.raise_for_status()
+            segment_data = response.content
+            
+            # Determine if it's an init segment or regular segment
+            if 'init' in segment_url.lower() or segment_url.endswith('.mp4'):
+                self.init_segment_cache[segment_url] = segment_data
+            else:
+                self.segment_cache[segment_url] = segment_data
+            
+            # Check for emergency cleanup
+            if self._check_memory_threshold():
+                self._emergency_cache_cleanup()
+            # Maintain cache size
+            elif len(self.segment_cache) > self.max_cache_size:
+                oldest_key = next(iter(self.segment_cache))
+                del self.segment_cache[oldest_key]
+            
+            logger.debug(f"Downloaded and cached DASH segment: {segment_url}")
+            return segment_data
+            
+        except Exception as e:
+            logger.warning(f"Failed to get DASH segment {segment_url}: {e}")
+            return None
+    
+    async def get_manifest(self, mpd_url: str, headers: Dict[str, str]) -> Optional[dict]:
+        """
+        Get MPD manifest from cache or download it.
+        
+        Args:
+            mpd_url (str): URL of the MPD manifest
+            headers (Dict[str, str]): Headers to use for request
+            
+        Returns:
+            Optional[dict]: Cached manifest data or None if not available
+        """
+        # Check cache first
+        if mpd_url in self.manifest_cache:
+            logger.debug(f"DASH cache hit for manifest: {mpd_url}")
+            return self.manifest_cache[mpd_url]
+        
+        # Download if not in cache
+        try:
+            response = await self.client.get(mpd_url, headers=headers)
+            response.raise_for_status()
+            mpd_content = response.text
+            mpd_dict = xmltodict.parse(mpd_content)
+            
+            # Cache the manifest
+            self.manifest_cache[mpd_url] = mpd_dict
+            
+            logger.debug(f"Downloaded and cached DASH manifest: {mpd_url}")
+            return mpd_dict
+            
+        except Exception as e:
+            logger.warning(f"Failed to get DASH manifest {mpd_url}: {e}")
+            return None
+    
+    def clear_cache(self) -> None:
+        """Clear the DASH cache."""
+        self.segment_cache.clear()
+        self.init_segment_cache.clear()
+        self.manifest_cache.clear()
+        self.adaptation_segments.clear()
+        logger.info("DASH pre-buffer cache cleared")
+    
+    async def close(self) -> None:
+        """Close the pre-buffer system."""
+        await self.client.aclose()
+
+
+# Global DASH pre-buffer instance
+dash_prebuffer = DASHPreBuffer()