mirror of
https://github.com/Dispatcharr/Dispatcharr.git
synced 2026-07-18 00:55:50 +00:00
cf1aef6747
Added a method to reset internal buffers to prevent corruption during stream transitions.
545 lines
22 KiB
Python
545 lines
22 KiB
Python
"""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()
|
|
|
|
# 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
|
|
writes_done = 0
|
|
with self.lock:
|
|
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
|