Proyecto LCX Dispatcharr multicuenta

apps/proxy/ts_proxy/stream_buffer.py

@@ -0,0 +1,546 @@
+"""Buffer management for TS streams"""
+
+import threading
+import logging
+import time
+from collections import deque
+from typing import Optional, Deque
+import random
+from apps.proxy.config import TSConfig as Config
+from .redis_keys import RedisKeys
+from .config_helper import ConfigHelper
+from .constants import TS_PACKET_SIZE
+from .utils import get_logger
+import gevent.event
+import gevent  # Make sure this import is at the top
+
+logger = get_logger()
+
+class StreamBuffer:
+    """Manages stream data buffering with optimized chunk storage"""
+
+    def __init__(self, channel_id=None, redis_client=None):
+        self.channel_id = channel_id
+        self.redis_client = redis_client
+        self.lock = threading.Lock()
+        self.index = 0
+        self.TS_PACKET_SIZE = TS_PACKET_SIZE
+
+        # STANDARDIZED KEYS: Use RedisKeys class instead of hardcoded patterns
+        self.buffer_index_key = RedisKeys.buffer_index(channel_id) if channel_id else ""
+        self.buffer_prefix = RedisKeys.buffer_chunk_prefix(channel_id) if channel_id else ""
+
+        self.chunk_ttl = ConfigHelper.redis_chunk_ttl()
+
+        # Initialize from Redis if available
+        if self.redis_client and channel_id:
+            try:
+                current_index = self.redis_client.get(self.buffer_index_key)
+                if current_index:
+                    self.index = int(current_index)
+                    logger.info(f"Initialized buffer from Redis with index {self.index}")
+            except Exception as e:
+                logger.error(f"Error initializing buffer from Redis: {e}")
+
+        self._write_buffer = bytearray()
+        self.target_chunk_size = ConfigHelper.get('BUFFER_CHUNK_SIZE', TS_PACKET_SIZE * 5644)  # ~1MB default
+
+        # Sorted-set key for chunk receive-timestamps (time-based positioning)
+        self.chunk_timestamps_key = RedisKeys.chunk_timestamps(channel_id) if channel_id else ""
+
+        # Register Lua scripts once — subsequent calls use EVALSHA (just the
+        # SHA hash) instead of sending the full script text on every invocation.
+        if self.redis_client:
+            self._find_oldest_chunk_sha = self.redis_client.register_script(
+                self._FIND_OLDEST_CHUNK_LUA
+            )
+            self._find_chunk_by_time_sha = self.redis_client.register_script(
+                self._FIND_CHUNK_BY_TIME_LUA
+            )
+        else:
+            self._find_oldest_chunk_sha = None
+            self._find_chunk_by_time_sha = None
+
+        # Track timers for proper cleanup
+        self.stopping = False
+        self.fill_timers = []
+        self.chunk_available = gevent.event.Event()
+
+    def add_chunk(self, chunk):
+        """Add data with optimized Redis storage and TS packet alignment"""
+        if not chunk:
+            return False
+
+        try:
+            # Accumulate partial packets between chunks
+            if not hasattr(self, '_partial_packet'):
+                self._partial_packet = bytearray()
+
+            # Lock the full operation to prevent race with reset_buffer_position
+            writes_done = 0
+            with self.lock:
+                # Combine with any previous partial packet
+                combined_data = bytearray(self._partial_packet) + bytearray(chunk)
+
+                # Calculate complete packets
+                complete_packets_size = (len(combined_data) // self.TS_PACKET_SIZE) * self.TS_PACKET_SIZE
+
+                if complete_packets_size == 0:
+                    # Not enough data for a complete packet
+                    self._partial_packet = combined_data
+                    return True
+
+                # Split into complete packets and remainder
+                complete_packets = combined_data[:complete_packets_size]
+                self._partial_packet = combined_data[complete_packets_size:]
+
+                # Add completed packets to write buffer
+                self._write_buffer.extend(complete_packets)
+
+                # Only write to Redis when we have enough data for an optimized chunk
+                while len(self._write_buffer) >= self.target_chunk_size:
+                    # Extract a full chunk
+                    chunk_data = self._write_buffer[:self.target_chunk_size]
+                    self._write_buffer = self._write_buffer[self.target_chunk_size:]
+
+                    # Write optimized chunk to Redis
+                    if self.redis_client:
+                        chunk_index = self.redis_client.incr(self.buffer_index_key)
+                        chunk_key = RedisKeys.buffer_chunk(self.channel_id, chunk_index)
+                        self.redis_client.setex(chunk_key, self.chunk_ttl, bytes(chunk_data))
+
+                        # Record receive timestamp for time-based client positioning
+                        if self.chunk_timestamps_key:
+                            now = time.time()
+                            self.redis_client.zadd(self.chunk_timestamps_key, {str(chunk_index): now})
+                            # Prune entries whose chunks have expired from Redis
+                            self.redis_client.zremrangebyscore(self.chunk_timestamps_key, '-inf', now - self.chunk_ttl)
+                            self.redis_client.expire(self.chunk_timestamps_key, self.chunk_ttl)
+
+                        # Update local tracking
+                        self.index = chunk_index
+                        writes_done += 1
+
+            if writes_done > 0:
+                logger.debug(f"Added {writes_done} chunks ({self.target_chunk_size} bytes each) to Redis for channel {self.channel_id} at index {self.index}")
+
+            self.chunk_available.set()  # Signal that new data is available
+            self.chunk_available.clear()  # Reset for next notification
+
+            return True
+
+        except Exception as e:
+            logger.error(f"Error adding chunk to buffer: {e}")
+            return False
+
+    def reset_buffer_position(self):
+        """
+        Reset internal buffers for a clean stream transition (failover).
+
+        Called by stream_manager.update_url() when switching between FFmpeg
+        processes. Without this, _partial_packet from the old FFmpeg gets
+        concatenated with the first bytes from the new FFmpeg, creating
+        corrupted TS packets that break audio decoder sync in the client.
+        """
+        try:
+            with self.lock:
+                old_write_size = len(self._write_buffer)
+                old_partial_size = len(getattr(self, '_partial_packet', b''))
+
+                self._write_buffer = bytearray()
+                if hasattr(self, '_partial_packet'):
+                    self._partial_packet = bytearray()
+
+                if old_write_size > 0 or old_partial_size > 0:
+                    logger.info(
+                        f"Reset buffer position for channel {self.channel_id}: "
+                        f"cleared {old_write_size} bytes from write buffer, "
+                        f"{old_partial_size} bytes from partial packet"
+                    )
+                else:
+                    logger.debug(
+                        f"Reset buffer position for channel {self.channel_id}: "
+                        f"buffers were already clean"
+                    )
+        except Exception as e:
+            logger.error(
+                f"Error resetting buffer position for channel {self.channel_id}: {e}"
+            )
+
+    def get_chunks(self, start_index=None):
+        """Get chunks from the buffer with detailed logging"""
+        try:
+            request_id = f"req_{random.randint(1000, 9999)}"
+            logger.debug(f"[{request_id}] get_chunks called with start_index={start_index}")
+
+            if not self.redis_client:
+                logger.error("Redis not available, cannot retrieve chunks")
+                return []
+
+            # If no start_index provided, use most recent chunks
+            if start_index is None:
+                start_index = max(0, self.index - 10)  # Start closer to current position
+                logger.debug(f"[{request_id}] No start_index provided, using {start_index}")
+
+            # Get current index from Redis
+            current_index = int(self.redis_client.get(self.buffer_index_key) or 0)
+
+            # Calculate range of chunks to retrieve
+            start_id = start_index + 1
+            chunks_behind = current_index - start_id
+
+            # Adaptive chunk retrieval based on how far behind
+            if chunks_behind > 100:
+                fetch_count = 15
+                logger.debug(f"[{request_id}] Client very behind ({chunks_behind} chunks), fetching {fetch_count}")
+            elif chunks_behind > 50:
+                fetch_count = 10
+                logger.debug(f"[{request_id}] Client moderately behind ({chunks_behind} chunks), fetching {fetch_count}")
+            elif chunks_behind > 20:
+                fetch_count = 5
+                logger.debug(f"[{request_id}] Client slightly behind ({chunks_behind} chunks), fetching {fetch_count}")
+            else:
+                fetch_count = 3
+                logger.debug(f"[{request_id}] Client up-to-date (only {chunks_behind} chunks behind), fetching {fetch_count}")
+
+            end_id = min(current_index + 1, start_id + fetch_count)
+
+            if start_id >= end_id:
+                logger.debug(f"[{request_id}] No new chunks to fetch (start_id={start_id}, end_id={end_id})")
+                return []
+
+            # Log the range we're retrieving
+            logger.debug(f"[{request_id}] Retrieving chunks {start_id} to {end_id-1} (total: {end_id-start_id})")
+
+            # Directly fetch from Redis using pipeline for efficiency
+            pipe = self.redis_client.pipeline()
+            for idx in range(start_id, end_id):
+                chunk_key = RedisKeys.buffer_chunk(self.channel_id, idx)
+                pipe.get(chunk_key)
+
+            results = pipe.execute()
+
+            # Process results
+            chunks = [result for result in results if result is not None]
+
+            # Count non-None results
+            found_chunks = len(chunks)
+            missing_chunks = len(results) - found_chunks
+
+            if missing_chunks > 0:
+                logger.debug(f"[{request_id}] Missing {missing_chunks}/{len(results)} chunks in Redis")
+
+            # Update local tracking
+            if chunks:
+                self.index = end_id - 1
+
+            # Final log message
+            chunk_sizes = [len(c) for c in chunks]
+            total_bytes = sum(chunk_sizes) if chunks else 0
+            logger.debug(f"[{request_id}] Returning {len(chunks)} chunks ({total_bytes} bytes)")
+
+            return chunks
+
+        except Exception as e:
+            logger.error(f"Error getting chunks from buffer: {e}", exc_info=True)
+            return []
+
+    def get_chunks_exact(self, start_index, count):
+        """Get exactly the requested number of chunks from given index"""
+        try:
+            if not self.redis_client:
+                logger.error("Redis not available, cannot retrieve chunks")
+                return []
+
+            # Calculate range to retrieve
+            start_id = start_index + 1
+            end_id = start_id + count
+
+            # Get current buffer position
+            current_index = int(self.redis_client.get(self.buffer_index_key) or 0)
+
+            # If requesting beyond current buffer, return what we have
+            if start_id > current_index:
+                return []
+
+            # Cap end at current buffer position
+            end_id = min(end_id, current_index + 1)
+
+            # Directly fetch from Redis using pipeline
+            pipe = self.redis_client.pipeline()
+            for idx in range(start_id, end_id):
+                chunk_key = RedisKeys.buffer_chunk(self.channel_id, idx)
+                pipe.get(chunk_key)
+
+            results = pipe.execute()
+
+            # Filter out None results
+            chunks = [result for result in results if result is not None]
+
+            # Update local index if needed
+            if chunks and start_id + len(chunks) - 1 > self.index:
+                self.index = start_id + len(chunks) - 1
+
+            return chunks
+
+        except Exception as e:
+            logger.error(f"Error getting exact chunks: {e}", exc_info=True)
+            return []
+
+    def stop(self):
+        """Stop the buffer and cancel all timers"""
+        # Set stopping flag first to prevent new timer creation
+        self.stopping = True
+
+        # Cancel all pending timers
+        timers_cancelled = 0
+        for timer in list(self.fill_timers):
+            try:
+                if timer and not timer.dead:  # Changed from timer.is_alive()
+                    timer.kill()  # Changed from timer.cancel()
+                    timers_cancelled += 1
+            except Exception as e:
+                logger.error(f"Error canceling timer: {e}")
+
+        if timers_cancelled:
+            logger.info(f"Cancelled {timers_cancelled} buffer timers for channel {self.channel_id}")
+
+        # Clear timer list
+        self.fill_timers.clear()
+
+        try:
+            # Flush any remaining data in the write buffer
+            if hasattr(self, '_write_buffer') and len(self._write_buffer) > 0:
+                # Ensure remaining data is aligned to TS packets
+                complete_size = (len(self._write_buffer) // 188) * 188
+
+                if complete_size > 0:
+                    final_chunk = self._write_buffer[:complete_size]
+
+                    # Write final chunk to Redis
+                    with self.lock:
+                        if self.redis_client:
+                            try:
+                                chunk_index = self.redis_client.incr(self.buffer_index_key)
+                                chunk_key = f"{self.buffer_prefix}{chunk_index}"
+                                self.redis_client.setex(chunk_key, self.chunk_ttl, bytes(final_chunk))
+                                self.index = chunk_index
+                                logger.info(f"Flushed final chunk of {len(final_chunk)} bytes to Redis")
+                            except Exception as e:
+                                logger.error(f"Error flushing final chunk: {e}")
+
+                # Clear buffers
+                self._write_buffer = bytearray()
+                if hasattr(self, '_partial_packet'):
+                    self._partial_packet = bytearray()
+
+            # Clean up the chunk timestamps sorted set
+            if self.redis_client and self.chunk_timestamps_key:
+                try:
+                    self.redis_client.delete(self.chunk_timestamps_key)
+                except Exception as e:
+                    logger.error(f"Error deleting chunk timestamps key: {e}")
+
+        except Exception as e:
+            logger.error(f"Error during buffer stop: {e}")
+
+    def get_optimized_client_data(self, client_index):
+        """Get optimal amount of data for client streaming based on position and target size"""
+        # Define limits
+        MIN_CHUNKS = 3                      # Minimum chunks to read for efficiency
+        MAX_CHUNKS = 20                     # Safety limit to prevent memory spikes
+        TARGET_SIZE = 1024 * 1024           # Target ~1MB per response (typical media buffer)
+        MAX_SIZE = 2 * 1024 * 1024          # Hard cap at 2MB
+
+        # Calculate how far behind we are
+        chunks_behind = self.index - client_index
+
+        # Determine optimal chunk count
+        if chunks_behind <= MIN_CHUNKS:
+            # Not much data, retrieve what's available
+            chunk_count = max(1, chunks_behind)
+        elif chunks_behind <= MAX_CHUNKS:
+            # Reasonable amount behind, catch up completely
+            chunk_count = chunks_behind
+        else:
+            # Way behind, retrieve MAX_CHUNKS to avoid memory pressure
+            chunk_count = MAX_CHUNKS
+
+        # Retrieve chunks
+        chunks = self.get_chunks_exact(client_index, chunk_count)
+
+        # Check if we got significantly fewer chunks than expected (likely due to expiration)
+        # Only check if we expected multiple chunks and got none or very few
+        if chunk_count > 3 and len(chunks) == 0 and chunks_behind > 10:
+            # Chunks are missing - likely expired from Redis
+            # Return empty list to signal client should skip forward
+            logger.debug(f"Chunks missing for client at index {client_index}, buffer at {self.index} ({chunks_behind} behind)")
+            return [], client_index
+
+        # Check total size
+        total_size = sum(len(c) for c in chunks)
+
+        # If we're under target and have more chunks available, get more
+        if total_size < TARGET_SIZE and chunks_behind > chunk_count:
+            # Calculate how many more chunks we can get
+            additional = min(MAX_CHUNKS - chunk_count, chunks_behind - chunk_count)
+            more_chunks = self.get_chunks_exact(client_index + chunk_count, additional)
+
+            # Check if adding more would exceed MAX_SIZE
+            additional_size = sum(len(c) for c in more_chunks)
+            if total_size + additional_size <= MAX_SIZE:
+                chunks.extend(more_chunks)
+                chunk_count += len(more_chunks)  # Fixed: count actual additional chunks retrieved
+
+        return chunks, client_index + chunk_count
+
+    # Lua script that runs an atomic binary search on the Redis server.
+    # Chunks expire in FIFO order (same TTL, sequential writes), so the
+    # alive range is contiguous: [oldest_surviving .. buffer_head].
+    # Binary search finds the boundary in O(log N) EXISTS calls with zero
+    # round-trips between steps and no TOCTOU races (Lua scripts are atomic).
+    #
+    # ARGV[1] = key prefix  (e.g. "ts_proxy:channel:<id>:buffer:chunk:")
+    # ARGV[2] = low index   (client_index + 1, first chunk the client needs)
+    # ARGV[3] = high index  (buffer head, most recent chunk)
+    #
+    # Returns: the index of the oldest existing chunk, or -1 if none exist.
+    _FIND_OLDEST_CHUNK_LUA = """
+    local prefix = ARGV[1]
+    local low    = tonumber(ARGV[2])
+    local high   = tonumber(ARGV[3])
+
+    if redis.call('EXISTS', prefix .. high) == 0 then
+        return -1
+    end
+
+    local result = high
+    while low <= high do
+        local mid = math.floor((low + high) / 2)
+        if redis.call('EXISTS', prefix .. mid) == 1 then
+            result = mid
+            high = mid - 1
+        else
+            low = mid + 1
+        end
+    end
+    return result
+    """
+
+    def find_oldest_available_chunk(self, client_index):
+        """Find the oldest (lowest-index) chunk that still exists in Redis.
+
+        Executes an atomic Lua binary search on the Redis server — one
+        round-trip, ~log2(N) EXISTS calls, no TOCTOU between steps.
+
+        The actual read attempt (get_optimized_client_data) is what
+        authoritatively detects expiration; this method is best-effort
+        positioning that self-corrects on the next iteration if the found
+        chunk also expires before the client can read it.
+
+        Args:
+            client_index: The client's current local_index (last consumed chunk).
+
+        Returns:
+            int or None: The local_index value the client should jump to
+                         (one before the first available chunk), or None if no
+                         chunks are available at all.
+        """
+        if not self.redis_client:
+            return None
+
+        low = client_index + 1   # First chunk the client needs
+        high = self.index        # Latest chunk written
+
+        if low > high:
+            return None
+
+        try:
+            # Uses EVALSHA under the hood — sends only the SHA hash,
+            # not the full script text, on every call after the first.
+            result = self._find_oldest_chunk_sha(
+                args=[
+                    RedisKeys.buffer_chunk_prefix(self.channel_id),
+                    low,
+                    high,
+                ],
+            )
+
+            if result == -1:
+                return None
+
+            # Return result - 1 so local_index points to one before the
+            # first available chunk (matching the "last consumed" convention).
+            return int(result) - 1
+
+        except Exception as e:
+            logger.error(f"Error running find_oldest_chunk Lua script for channel {self.channel_id}: {e}")
+            return None
+
+    # ------------------------------------------------------------------
+    # Lua script: atomic reverse-scan of the chunk_timestamps sorted set.
+    # Finds the chunk whose receive-timestamp is closest to (but <=) a
+    # target wall-clock time.  Returns the chunk index or -1.
+    #
+    # KEYS[1] = chunk_timestamps sorted-set key
+    # ARGV[1] = target timestamp  (time.time() - desired_seconds_behind)
+    # ------------------------------------------------------------------
+    _FIND_CHUNK_BY_TIME_LUA = """
+    local ts_key  = KEYS[1]
+    local target  = tonumber(ARGV[1])
+
+    -- ZREVRANGEBYSCORE returns members with score <= target, highest first.
+    local result = redis.call('ZREVRANGEBYSCORE', ts_key, target, '-inf', 'LIMIT', 0, 1)
+    if #result == 0 then
+        return -1
+    end
+    return tonumber(result[1])
+    """
+
+    def find_chunk_index_by_time(self, seconds_behind):
+        """Find the chunk index that was received approximately *seconds_behind*
+        seconds ago.
+
+        Uses an atomic Lua script against the chunk_timestamps sorted set so
+        no data can expire between the lookup and the read.
+
+        Returns:
+            int or None: The chunk index to position the client at (this is
+                         the *last consumed* convention, so the next read
+                         starts at index+1).  None if no suitable chunk
+                         exists.
+        """
+        if not self.redis_client or not self.chunk_timestamps_key:
+            return None
+
+        target_time = time.time() - seconds_behind
+
+        try:
+            result = self._find_chunk_by_time_sha(
+                keys=[self.chunk_timestamps_key],
+                args=[target_time],
+            )
+            if result is None or int(result) == -1:
+                # No chunk old enough — fall back to the oldest available chunk
+                oldest = self.redis_client.zrange(self.chunk_timestamps_key, 0, 0)
+                if oldest:
+                    return max(0, int(oldest[0]) - 1)  # "last consumed" convention
+                return None
+
+            # Return index - 1 so next read starts at that chunk
+            return max(0, int(result) - 1)
+
+        except Exception as e:
+            logger.error(f"Error in find_chunk_index_by_time for channel {self.channel_id}: {e}")
+            return None
+
+    # Add a new method to safely create timers
+    def schedule_timer(self, delay, callback, *args, **kwargs):
+        """Schedule a timer and track it for proper cleanup"""
+        if self.stopping:
+            return None
+
+        # Replace threading.Timer with gevent.spawn_later for better compatibility
+        timer = gevent.spawn_later(delay, callback, *args, **kwargs)
+        self.fill_timers.append(timer)
+        return timer