diff --git a/apps/proxy/vod_proxy/connection_manager.py b/apps/proxy/vod_proxy/connection_manager.py index aafc75cd..c4c2130f 100644 --- a/apps/proxy/vod_proxy/connection_manager.py +++ b/apps/proxy/vod_proxy/connection_manager.py @@ -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: diff --git a/apps/proxy/vod_proxy/multi_worker_connection_manager.py b/apps/proxy/vod_proxy/multi_worker_connection_manager.py index 1534f761..74c5db41 100644 --- a/apps/proxy/vod_proxy/multi_worker_connection_manager.py +++ b/apps/proxy/vod_proxy/multi_worker_connection_manager.py @@ -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: