Spaces:
Configuration error
Configuration error
File size: 15,069 Bytes
b3ca762 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 | """
DASH Pre-buffer system for reducing latency and improving streaming performance.
This module extends BasePrebuffer with DASH-specific functionality including
MPD parsing integration, profile handling, and init segment management.
"""
import asyncio
import logging
import time
from typing import Dict, Optional, List
from mediaflow_proxy.utils.base_prebuffer import BasePrebuffer
from mediaflow_proxy.utils.cache_utils import (
get_cached_mpd,
get_cached_init_segment,
)
from mediaflow_proxy.configs import settings
logger = logging.getLogger(__name__)
class DASHPreBuffer(BasePrebuffer):
"""
Pre-buffer system for DASH streams.
Extends BasePrebuffer with DASH-specific features:
- MPD manifest parsing and profile handling
- Init segment prebuffering
- Live stream segment tracking
- Profile-based segment prefetching
Uses event-based download coordination from BasePrebuffer to prevent
duplicate downloads between player requests and background prebuffering.
"""
def __init__(
self,
max_cache_size: Optional[int] = None,
prebuffer_segments: Optional[int] = None,
):
"""
Initialize the DASH pre-buffer system.
Args:
max_cache_size: Maximum number of segments to cache (uses config if None)
prebuffer_segments: Number of segments to pre-buffer ahead (uses config if None)
"""
super().__init__(
max_cache_size=max_cache_size or settings.dash_prebuffer_cache_size,
prebuffer_segments=prebuffer_segments or settings.dash_prebuffer_segments,
max_memory_percent=settings.dash_prebuffer_max_memory_percent,
emergency_threshold=settings.dash_prebuffer_emergency_threshold,
segment_ttl=settings.dash_segment_cache_ttl,
)
self.inactivity_timeout = settings.dash_prebuffer_inactivity_timeout
# DASH-specific state
# Track active streams for prefetching: mpd_url -> stream_info
self.active_streams: Dict[str, dict] = {}
self.prefetch_tasks: Dict[str, asyncio.Task] = {}
# Additional stats for DASH
self.init_segments_prebuffered = 0
# Cleanup task
self._cleanup_task: Optional[asyncio.Task] = None
def log_stats(self) -> None:
"""Log current prebuffer statistics with DASH-specific info."""
stats = self.stats.to_dict()
stats["init_segments_prebuffered"] = self.init_segments_prebuffered
stats["active_streams"] = len(self.active_streams)
logger.info(f"DASH Prebuffer Stats: {stats}")
async def prebuffer_dash_manifest(
self,
mpd_url: str,
headers: Dict[str, str],
) -> None:
"""
Pre-buffer segments from a DASH manifest using existing MPD parsing.
Args:
mpd_url: URL of the DASH manifest
headers: Headers to use for requests
"""
try:
# First get the basic MPD info without segments
parsed_mpd = await get_cached_mpd(mpd_url, headers, parse_drm=False)
if not parsed_mpd:
logger.warning(f"Failed to get parsed MPD for prebuffering: {mpd_url}")
return
is_live = parsed_mpd.get("isLive", False)
base_profiles = parsed_mpd.get("profiles", [])
if not base_profiles:
logger.warning(f"No profiles found in MPD for prebuffering: {mpd_url}")
return
# Now get segments for each profile by parsing with profile_id
profiles_with_segments = []
for profile in base_profiles:
profile_id = profile.get("id")
if profile_id:
parsed_with_segments = await get_cached_mpd(
mpd_url, headers, parse_drm=False, parse_segment_profile_id=profile_id
)
# Find the matching profile with segments
for p in parsed_with_segments.get("profiles", []):
if p.get("id") == profile_id:
profiles_with_segments.append(p)
break
# Store stream info for ongoing prefetching
self.active_streams[mpd_url] = {
"headers": headers,
"is_live": is_live,
"profiles": profiles_with_segments,
"last_access": time.time(),
}
# Prebuffer init segments and media segments
await self._prebuffer_profiles(profiles_with_segments, headers, is_live)
# Start cleanup task if not running
self._ensure_cleanup_task_running()
logger.info(
f"Pre-buffered DASH manifest: {mpd_url} (live={is_live}, profiles={len(profiles_with_segments)})"
)
except Exception as e:
logger.warning(f"Failed to pre-buffer DASH manifest {mpd_url}: {e}")
async def _prebuffer_profiles(
self,
profiles: List[dict],
headers: Dict[str, str],
is_live: bool = False,
) -> None:
"""
Pre-buffer init segments and media segments for all profiles.
For live streams, prebuffers from the END of the segment list.
For VOD, prebuffers from the beginning.
Args:
profiles: List of parsed profiles with resolved URLs
headers: Headers to use for requests
is_live: Whether this is a live stream
"""
if self._should_skip_for_memory():
logger.warning("Memory usage too high, skipping prebuffer")
return
# Collect all segment URLs to prebuffer
segment_urls = []
init_urls = []
for profile in profiles:
# Collect init segment URL
init_url = profile.get("initUrl")
if init_url:
init_urls.append(init_url)
# Get segments to prebuffer
segments = profile.get("segments", [])
if not segments:
continue
# For live streams, prebuffer from the END (most recent)
if is_live:
segments_to_buffer = segments[-self.prebuffer_segment_count :]
else:
segments_to_buffer = segments[: self.prebuffer_segment_count]
for segment in segments_to_buffer:
segment_url = segment.get("media")
if segment_url:
segment_urls.append(segment_url)
# Prebuffer init segments (using special init cache)
for init_url in init_urls:
asyncio.create_task(self._prebuffer_init_segment(init_url, headers))
# Prebuffer media segments using base class method
if segment_urls:
await self.prebuffer_segments_batch(segment_urls, headers)
async def _prebuffer_init_segment(
self,
init_url: str,
headers: Dict[str, str],
) -> None:
"""
Prebuffer an init segment using the init segment cache.
Args:
init_url: URL of the init segment
headers: Headers for the request
"""
try:
# get_cached_init_segment handles both caching and downloading
content = await get_cached_init_segment(init_url, headers)
if content:
self.init_segments_prebuffered += 1
self.stats.bytes_prebuffered += len(content)
logger.debug(f"Prebuffered init segment ({len(content)} bytes)")
except Exception as e:
logger.warning(f"Failed to prebuffer init segment: {e}")
async def prefetch_upcoming_segments(
self,
mpd_url: str,
current_segment_url: str,
headers: Dict[str, str],
profile_id: Optional[str] = None,
) -> None:
"""
Prefetch upcoming segments based on current playback position.
Called when a segment is requested to prefetch the next N segments.
Args:
mpd_url: URL of the MPD manifest
current_segment_url: URL of the currently requested segment
headers: Headers to use for requests
profile_id: Optional profile ID to limit prefetching to
"""
self.stats.prefetch_triggered += 1
try:
# First check if we have cached profiles with segments
if mpd_url in self.active_streams:
# Update last access time
self.active_streams[mpd_url]["last_access"] = time.time()
profiles = self.active_streams[mpd_url].get("profiles", [])
else:
# Get parsed MPD
parsed_mpd = await get_cached_mpd(mpd_url, headers, parse_drm=False)
if not parsed_mpd:
return
profiles = parsed_mpd.get("profiles", [])
for profile in profiles:
pid = profile.get("id")
if profile_id and pid != profile_id:
continue
segments = profile.get("segments", [])
# If no segments, try to get them by parsing with profile_id
if not segments and pid:
parsed_with_segments = await get_cached_mpd(
mpd_url, headers, parse_drm=False, parse_segment_profile_id=pid
)
for p in parsed_with_segments.get("profiles", []):
if p.get("id") == pid:
segments = p.get("segments", [])
break
# Find current segment index
current_index = -1
for i, segment in enumerate(segments):
if segment.get("media") == current_segment_url:
current_index = i
break
if current_index < 0:
continue
# Collect next N segment URLs
segment_urls = []
end_index = min(current_index + 1 + self.prebuffer_segment_count, len(segments))
for i in range(current_index + 1, end_index):
segment_url = segments[i].get("media")
if segment_url:
segment_urls.append(segment_url)
if segment_urls:
logger.debug(f"Prefetching {len(segment_urls)} upcoming segments from index {current_index + 1}")
# Run prefetch in background
asyncio.create_task(self.prebuffer_segments_batch(segment_urls, headers, max_concurrent=3))
except Exception as e:
logger.warning(f"Failed to prefetch upcoming segments: {e}")
async def prefetch_for_live_playlist(
self,
profiles: List[dict],
headers: Dict[str, str],
) -> None:
"""
Prefetch segments for a live playlist refresh.
Called from process_playlist to ensure upcoming segments are cached.
Args:
profiles: List of profiles with resolved segment URLs
headers: Headers to use for requests
"""
segment_urls = []
for profile in profiles:
segments = profile.get("segments", [])
if not segments:
continue
# For live, prefetch the last N segments (most recent)
segments_to_prefetch = segments[-self.prebuffer_segment_count :]
for segment in segments_to_prefetch:
segment_url = segment.get("media")
if segment_url:
# Check if already cached before adding
cached = await self.try_get_cached(segment_url)
if not cached:
segment_urls.append(segment_url)
if segment_urls:
logger.debug(f"Live playlist prefetch: {len(segment_urls)} segments")
asyncio.create_task(self.prebuffer_segments_batch(segment_urls, headers, max_concurrent=3))
def _ensure_cleanup_task_running(self) -> None:
"""Ensure the cleanup task is running."""
if self._cleanup_task is None or self._cleanup_task.done():
self._cleanup_task = asyncio.create_task(self._cleanup_inactive_streams())
async def _cleanup_inactive_streams(self) -> None:
"""
Periodically check for and clean up inactive streams.
Runs in the background and removes streams that haven't been
accessed recently.
"""
while True:
try:
await asyncio.sleep(30) # Check every 30 seconds
if not self.active_streams:
logger.debug("No active DASH streams to monitor, stopping cleanup")
return
current_time = time.time()
streams_to_remove = []
for mpd_url, stream_info in self.active_streams.items():
last_access = stream_info.get("last_access", 0)
time_since_access = current_time - last_access
if time_since_access > self.inactivity_timeout:
streams_to_remove.append(mpd_url)
logger.info(f"Cleaning up inactive DASH stream ({time_since_access:.0f}s idle)")
# Remove inactive streams
for mpd_url in streams_to_remove:
self.active_streams.pop(mpd_url, None)
task = self.prefetch_tasks.pop(mpd_url, None)
if task:
task.cancel()
if streams_to_remove:
logger.info(f"Cleaned up {len(streams_to_remove)} inactive DASH stream(s)")
except asyncio.CancelledError:
logger.debug("DASH cleanup task cancelled")
return
except Exception as e:
logger.warning(f"Error in DASH cleanup task: {e}")
def get_stats(self) -> dict:
"""Get current prebuffer statistics."""
stats = self.stats.to_dict()
stats["init_segments_prebuffered"] = self.init_segments_prebuffered
stats["active_streams"] = len(self.active_streams)
return stats
def clear_cache(self) -> None:
"""Clear active streams tracking and log final stats."""
self.log_stats()
self.active_streams.clear()
for task in self.prefetch_tasks.values():
task.cancel()
self.prefetch_tasks.clear()
# Cancel cleanup task
if self._cleanup_task and not self._cleanup_task.done():
self._cleanup_task.cancel()
self._cleanup_task = None
self.stats.reset()
self.init_segments_prebuffered = 0
logger.info("DASH pre-buffer state cleared")
async def close(self) -> None:
"""Close the pre-buffer system."""
self.clear_cache()
# Global DASH pre-buffer instance
dash_prebuffer = DASHPreBuffer()
|