Mirrors/Dispatcharr
1
0
Fork 1
mirror of https://github.com/Dispatcharr/Dispatcharr.git synced 2026-07-18 00:55:50 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 4

Merge pull request #983 from CodeBormen/fix/962-VOD-proxy-connection-counter-leak-on-client-disconnect

Browse source 
Fix: VOD proxy connection counter leak on client disconnect
This commit is contained in:
SergeantPanda 2026-02-19 11:32:20 -06:00 • committed by GitHub
commit 1503127f61
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 185 additions and 28 deletions
Download patch file Download diff file Expand all files Collapse all files

49
apps/proxy/vod_proxy/connection_manager.py
View file

@ -459,13 +459,12 @@ class VODConnectionManager:
return False
try:
# Check profile connection limits using standardized key
if not self._check_profile_limits(m3u_profile):
# Atomically check and reserve a profile connection slot (INCR-first)
if not self._check_and_reserve_profile_slot(m3u_profile):
logger.warning(f"Profile {m3u_profile.name} connection limit exceeded")
return False
connection_key = self._get_connection_key(content_type, content_uuid, client_id)
profile_connections_key = self._get_profile_connections_key(m3u_profile.id)
content_connections_key = self._get_content_connections_key(content_type, content_uuid)
# Check if connection already exists to prevent duplicate counting
@ -473,6 +472,9 @@ class VODConnectionManager:
logger.info(f"Connection already exists for {client_id} - {content_type} {content_name}")
# Update activity but don't increment profile counter
self.redis_client.hset(connection_key, "last_activity", str(time.time()))
# Roll back the reservation — connection already counted
if m3u_profile.max_streams > 0:
self.redis_client.decr(self._get_profile_connections_key(m3u_profile.id))
return True
# Connection data
@ -499,8 +501,7 @@ class VODConnectionManager:
pipe.hset(connection_key, mapping=connection_data)
pipe.expire(connection_key, self.connection_ttl)
# Increment profile connections using standardized method
pipe.incr(profile_connections_key)
# Profile counter already incremented atomically above — no pipe.incr needed
# Add to content connections set
pipe.sadd(content_connections_key, client_id)
@ -513,6 +514,9 @@ class VODConnectionManager:
return True
except Exception as e:
# Roll back the profile reservation on failure
if m3u_profile.max_streams > 0:
self.redis_client.decr(self._get_profile_connections_key(m3u_profile.id))
logger.error(f"Error creating VOD connection: {e}")
return False
@ -531,6 +535,41 @@ class VODConnectionManager:
logger.error(f"Error checking profile limits: {e}")
return False
def _check_and_reserve_profile_slot(self, m3u_profile: M3UAccountProfile) -> bool:
"""
Atomically check and reserve a connection slot for the given profile.
Uses an INCR-first-then-check pattern to eliminate the TOCTOU race
condition where separate GET > check > INCR operations could allow
concurrent requests to both pass the capacity check.
For profiles with max_streams=0 (unlimited), no reservation is needed.
Returns:
bool: True if slot was reserved (or unlimited), False if at capacity
"""
if m3u_profile.max_streams == 0: # Unlimited
return True
try:
profile_connections_key = self._get_profile_connections_key(m3u_profile.id)
# Atomically increment first — single Redis command eliminates race window
new_count = self.redis_client.incr(profile_connections_key)
if new_count <= m3u_profile.max_streams:
logger.info(f"[PROFILE-RESERVE] Profile {m3u_profile.id} slot reserved: {new_count}/{m3u_profile.max_streams}")
return True
# Over capacity — roll back the increment
self.redis_client.decr(profile_connections_key)
logger.info(f"[PROFILE-RESERVE] Profile {m3u_profile.id} at capacity: {new_count - 1}/{m3u_profile.max_streams}")
return False
except Exception as e:
logger.error(f"Error reserving profile slot: {e}")
return False
def update_connection_activity(self, content_type: str, content_uuid: str,
client_id: str, bytes_sent: int = 0,
position_seconds: int = 0) -> bool:

164
apps/proxy/vod_proxy/multi_worker_connection_manager.py
View file

@ -416,8 +416,22 @@ class RedisBackedVODConnection:
logger.info(f"[{self.session_id}] Updated connection state: length={state.content_length}, type={state.content_type}")
# Save updated state
self._save_connection_state(state)
# Save updated state under lock to avoid overwriting concurrent
# active_streams changes (e.g., another stream's GeneratorExit decrement)
if self._acquire_lock():
try:
current = self._get_connection_state()
if current:
# Preserve the current active_streams value — it may have been
# modified by concurrent increment/decrement operations while
# waiting for the upstream HTTP response.
state.active_streams = current.active_streams
self._save_connection_state(state)
finally:
self._release_lock()
else:
# Fallback: save without lock but skip active_streams to avoid overwrite
logger.warning(f"[{self.session_id}] Could not acquire lock for get_stream state save")
self.local_response = response
return response
@ -466,35 +480,44 @@ class RedisBackedVODConnection:
return range_header
def increment_active_streams(self):
"""Increment active streams count in Redis"""
"""Increment active streams count in Redis. Returns new active_streams count, or 0 on failure."""
if not self._acquire_lock():
return False
logger.warning(f"[{self.session_id}] INCR-AS failed: could not acquire lock")
return 0
try:
state = self._get_connection_state()
if state:
old = state.active_streams
state.active_streams += 1
state.last_activity = time.time()
self._save_connection_state(state)
logger.debug(f"[{self.session_id}] Active streams incremented to {state.active_streams}")
return True
return False
logger.debug(f"[{self.session_id}] INCR-AS {old} -> {state.active_streams}")
return state.active_streams
logger.warning(f"[{self.session_id}] INCR-AS failed: no state")
return 0
finally:
self._release_lock()
def decrement_active_streams(self):
"""Decrement active streams count in Redis"""
if not self._acquire_lock():
logger.warning(f"[{self.session_id}] DECR-AS failed: could not acquire lock")
return False
try:
state = self._get_connection_state()
if state and state.active_streams > 0:
old = state.active_streams
state.active_streams -= 1
state.last_activity = time.time()
self._save_connection_state(state)
logger.debug(f"[{self.session_id}] Active streams decremented to {state.active_streams}")
logger.debug(f"[{self.session_id}] DECR-AS {old} -> {state.active_streams}")
return True
if not state:
logger.warning(f"[{self.session_id}] DECR-AS failed: no state")
else:
logger.warning(f"[{self.session_id}] DECR-AS failed: active_streams already {state.active_streams}")
return False
finally:
self._release_lock()
@ -674,6 +697,41 @@ class MultiWorkerVODConnectionManager:
logger.error(f"Error checking profile limits: {e}")
return False
def _check_and_reserve_profile_slot(self, m3u_profile) -> bool:
"""
Atomically check and reserve a connection slot for the given profile.
Uses an INCR-first-then-check pattern to eliminate the TOCTOU race
condition where separate GET > check > INCR operations could allow
concurrent requests to both pass the capacity check.
For profiles with max_streams=0 (unlimited), no reservation is needed.
Returns:
bool: True if slot was reserved (or unlimited), False if at capacity
"""
if m3u_profile.max_streams == 0: # Unlimited
return True
try:
profile_connections_key = self._get_profile_connections_key(m3u_profile.id)
# Atomically increment first — single Redis command eliminates race window
new_count = self.redis_client.incr(profile_connections_key)
if new_count <= m3u_profile.max_streams:
logger.info(f"[PROFILE-RESERVE] Profile {m3u_profile.id} slot reserved: {new_count}/{m3u_profile.max_streams}")
return True
# Over capacity — roll back the increment
self.redis_client.decr(profile_connections_key)
logger.info(f"[PROFILE-RESERVE] Profile {m3u_profile.id} at capacity: {new_count - 1}/{m3u_profile.max_streams}")
return False
except Exception as e:
logger.error(f"Error reserving profile slot: {e}")
return False
def _increment_profile_connections(self, m3u_profile):
"""Increment profile connection count"""
try:
@ -756,10 +814,11 @@ class MultiWorkerVODConnectionManager:
if not existing_state:
logger.info(f"[{client_id}] Worker {self.worker_id} - Creating new Redis-backed connection")
# Check profile limits before creating new connection
if not self._check_profile_limits(m3u_profile):
# Atomically check and reserve a profile connection slot (INCR-first)
if not self._check_and_reserve_profile_slot(m3u_profile):
logger.warning(f"[{client_id}] Profile {m3u_profile.name} connection limit exceeded")
return HttpResponse("Connection limit exceeded for profile", status=429)
profile_connections_incremented = True
# Apply timeshift parameters
modified_stream_url = self._apply_timeshift_parameters(stream_url, utc_start, utc_end, offset)
@ -802,16 +861,43 @@ class MultiWorkerVODConnectionManager:
worker_id=self.worker_id
):
logger.error(f"[{client_id}] Worker {self.worker_id} - Failed to create Redis connection")
# Roll back the profile slot reservation since connection failed
self._decrement_profile_connections(m3u_profile.id)
profile_connections_incremented = False
return HttpResponse("Failed to create connection", status=500)
# Increment profile connections after successful connection creation
self._increment_profile_connections(m3u_profile)
profile_connections_incremented = True
logger.info(f"[{client_id}] Worker {self.worker_id} - Created consolidated connection with session metadata")
else:
logger.info(f"[{client_id}] Worker {self.worker_id} - Using existing Redis-backed connection")
# Immediately increment active_streams to prevent cleanup race condition.
# Without this, stream's GeneratorExit can see active_streams=0
# and DECR the profile counter before the new generator starts.
if matching_session_id:
# Idle session reuse: active_streams already incremented at line 776
# Always need to re-reserve profile slot (GeneratorExit DECRed it)
if not self._check_and_reserve_profile_slot(m3u_profile):
logger.warning(f"[{client_id}] Profile {m3u_profile.name} connection limit exceeded on session reuse")
redis_connection.decrement_active_streams()
return HttpResponse("Connection limit exceeded for profile", status=429)
profile_connections_incremented = True
else:
# Concurrent/reconnect: increment active_streams now (not in generator)
new_count = redis_connection.increment_active_streams()
if new_count == 1:
# 0→1 transition: previous stream's GeneratorExit already DECRed
# the profile counter, need to re-reserve the slot
if not self._check_and_reserve_profile_slot(m3u_profile):
logger.warning(f"[{client_id}] Profile {m3u_profile.name} connection limit exceeded on reconnect")
redis_connection.decrement_active_streams()
return HttpResponse("Connection limit exceeded for profile", status=429)
profile_connections_incremented = True
elif new_count == 0:
logger.error(f"[{client_id}] Failed to increment active streams")
return HttpResponse("Failed to reserve stream", status=500)
# else: new_count > 1, another stream is already active and profile
# counter already reflects it — no INCR needed
# Transfer ownership to current worker and update session activity
if redis_connection._acquire_lock():
try:
@ -834,6 +920,12 @@ class MultiWorkerVODConnectionManager:
if upstream_response is None:
logger.warning(f"[{client_id}] Worker {self.worker_id} - Range not satisfiable")
if existing_state:
# Roll back the active_streams increment from the else branch
redis_connection.decrement_active_streams()
if profile_connections_incremented:
self._decrement_profile_connections(m3u_profile.id)
profile_connections_incremented = False
return HttpResponse("Requested Range Not Satisfiable", status=416)
# Get connection headers
@ -846,13 +938,14 @@ class MultiWorkerVODConnectionManager:
try:
logger.info(f"[{client_id}] Worker {self.worker_id} - Starting Redis-backed stream")
# Increment active streams (unless we already did it for session reuse)
if not matching_session_id:
# New session - increment active streams
# Increment active streams only for brand-new connections.
# For existing connections (session reuse or concurrent requests),
# active_streams was already incremented in the else branch above
# to prevent cleanup race conditions with GeneratorExit.
if not existing_state:
redis_connection.increment_active_streams()
else:
# Reused session - we already incremented when reserving the session
logger.debug(f"[{client_id}] Using pre-reserved session - active streams already incremented")
logger.debug(f"[{client_id}] Active streams already incremented in connection reuse path")
bytes_sent = 0
chunk_count = 0
@ -898,11 +991,19 @@ class MultiWorkerVODConnectionManager:
# Schedule smart cleanup if no active streams after normal completion
if not redis_connection.has_active_streams():
# Decrement profile counter immediately — don't defer to daemon thread
state = redis_connection._get_connection_state()
profile_id = state.m3u_profile_id if state else m3u_profile.id
if profile_id:
self._decrement_profile_connections(profile_id)
logger.info(f"[{client_id}] Profile counter decremented for profile {profile_id} on normal completion")
def delayed_cleanup():
time.sleep(1) # Wait 1 second
# Smart cleanup: check active streams and ownership
logger.info(f"[{client_id}] Worker {self.worker_id} - Checking for smart cleanup after normal completion")
redis_connection.cleanup(connection_manager=self, current_worker_id=self.worker_id)
# No connection_manager — profile already decremented above
redis_connection.cleanup(current_worker_id=self.worker_id)
import threading
cleanup_thread = threading.Thread(target=delayed_cleanup)
@ -917,11 +1018,19 @@ class MultiWorkerVODConnectionManager:
# Schedule smart cleanup if no active streams
if not redis_connection.has_active_streams():
# Decrement profile counter immediately — don't defer to daemon thread
state = redis_connection._get_connection_state()
profile_id = state.m3u_profile_id if state else m3u_profile.id
if profile_id:
self._decrement_profile_connections(profile_id)
logger.info(f"[{client_id}] Profile counter decremented for profile {profile_id} on client disconnect")
def delayed_cleanup():
time.sleep(1) # Wait 1 second
# Smart cleanup: check active streams and ownership
logger.info(f"[{client_id}] Worker {self.worker_id} - Checking for smart cleanup after client disconnect")
redis_connection.cleanup(connection_manager=self, current_worker_id=self.worker_id)
# No connection_manager — profile already decremented above
redis_connection.cleanup(current_worker_id=self.worker_id)
import threading
cleanup_thread = threading.Thread(target=delayed_cleanup)
@ -933,8 +1042,17 @@ class MultiWorkerVODConnectionManager:
if not decremented:
redis_connection.decrement_active_streams()
decremented = True
# Smart cleanup on error - immediate cleanup since we're in error state
redis_connection.cleanup(connection_manager=self, current_worker_id=self.worker_id)
# Decrement profile counter immediately if no other active streams
if not redis_connection.has_active_streams():
state = redis_connection._get_connection_state()
profile_id = state.m3u_profile_id if state else m3u_profile.id
if profile_id:
self._decrement_profile_connections(profile_id)
logger.info(f"[{client_id}] Profile counter decremented for profile {profile_id} on stream error")
# Smart cleanup on error - immediate cleanup since we're in error state
# No connection_manager — profile already decremented above
redis_connection.cleanup(current_worker_id=self.worker_id)
yield b"Error: Stream interrupted"
finally: