From 40d0bc95677b474f671dac2f4bfff905ee76cf1f Mon Sep 17 00:00:00 2001 From: None <190783071+CodeBormen@users.noreply.github.com> Date: Sun, 15 Mar 2026 14:22:57 -0500 Subject: [PATCH] fix: cap keepalive mode duration to prevent indefinite client hold - Add MAX_KEEPALIVE_DURATION = 300s to TSConfig - Track keepalive_start_time in _stream_data_generator; disconnect after cap is exceeded with a warning log - Reset timer on real data so independent stalls don't accumulate - Fix pre-existing test helper missing throttle attributes in NonOwnerWorkerKeepaliveTests --- .../channels/tests/test_ts_proxy_keepalive.py | 6 + .../tests/test_ts_proxy_keepalive_duration.py | 195 ++++++++++++++++++ apps/proxy/config.py | 1 + apps/proxy/ts_proxy/stream_generator.py | 16 ++ 4 files changed, 218 insertions(+) create mode 100644 apps/channels/tests/test_ts_proxy_keepalive_duration.py diff --git a/apps/channels/tests/test_ts_proxy_keepalive.py b/apps/channels/tests/test_ts_proxy_keepalive.py index db388b3d..9e14fdfd 100644 --- a/apps/channels/tests/test_ts_proxy_keepalive.py +++ b/apps/channels/tests/test_ts_proxy_keepalive.py @@ -85,6 +85,12 @@ class NonOwnerWorkerKeepaliveTests(TestCase): gen.stream_manager = None # non-owner worker gen.consecutive_empty = consecutive_empty + + # Throttle attributes added by a later PR; required by _should_send_keepalive. + gen._last_health_check_time = 0.0 + gen._last_health_check_result = False + gen._health_check_interval = 2.0 + gen.proxy_server = None return gen def _mock_proxy_server(self, last_data_value): diff --git a/apps/channels/tests/test_ts_proxy_keepalive_duration.py b/apps/channels/tests/test_ts_proxy_keepalive_duration.py new file mode 100644 index 00000000..74494131 --- /dev/null +++ b/apps/channels/tests/test_ts_proxy_keepalive_duration.py @@ -0,0 +1,195 @@ +""" +Unit tests for the keepalive duration cap in StreamGenerator._stream_data_generator. + +Verifies that a client held in keepalive mode is disconnected after +MAX_KEEPALIVE_DURATION seconds, and that the timer resets when real data resumes. +""" + +import time +from unittest.mock import MagicMock, patch, call +from django.test import TestCase + + +def _make_generator(consecutive_empty=10, local_index=10, buffer_index=10): + """Minimal StreamGenerator stub for testing _stream_data_generator logic.""" + from apps.proxy.ts_proxy.stream_generator import StreamGenerator + + gen = StreamGenerator.__new__(StreamGenerator) + gen.channel_id = "00000000-0000-0000-0000-000000000099" + gen.client_id = "test-client-duration" + gen.consecutive_empty = consecutive_empty + gen.empty_reads = 0 + gen.local_index = local_index + gen.bytes_sent = 0 + gen.chunks_sent = 0 + gen.last_yield_time = time.time() + gen.stream_start_time = time.time() + gen.last_stats_time = time.time() + gen.last_stats_bytes = 0 + gen.current_rate = 0.0 + gen.last_ttl_refresh = time.time() + gen.ttl_refresh_interval = 3 + gen.is_owner_worker = False + gen.stream_manager = None + gen._last_health_check_time = 0.0 + gen._last_health_check_result = False + gen._health_check_interval = 2.0 + gen.proxy_server = None + + buffer = MagicMock() + buffer.index = buffer_index + buffer.get_optimized_client_data.return_value = ([], local_index) + buffer.find_oldest_available_chunk.return_value = None + gen.buffer = buffer + + return gen + + +class KeepaliveDurationCapTests(TestCase): + """MAX_KEEPALIVE_DURATION cap disconnects clients stuck in keepalive mode.""" + + def _run_generator_to_break(self, gen, max_iterations=20): + """Drive _stream_data_generator until it breaks or hits iteration limit.""" + iterations = 0 + for _ in gen._stream_data_generator(): + iterations += 1 + if iterations >= max_iterations: + break + return iterations + + def test_cap_fires_after_max_duration_exceeded(self): + """Generator exits when keepalive has run longer than MAX_KEEPALIVE_DURATION.""" + gen = _make_generator() + + with patch.object(gen, '_check_resources', return_value=True), \ + patch.object(gen, '_should_send_keepalive', return_value=True), \ + patch.object(gen, '_is_ghost_client', return_value=False), \ + patch.object(gen, '_is_timeout', return_value=False), \ + patch('apps.proxy.ts_proxy.stream_generator.create_ts_packet', return_value=b'\x00' * 188), \ + patch('apps.proxy.ts_proxy.stream_generator.ProxyServer') as MockPS, \ + patch('apps.proxy.ts_proxy.stream_generator.Config') as MockConfig, \ + patch('apps.proxy.ts_proxy.stream_generator.gevent') as mock_gevent, \ + patch('apps.proxy.ts_proxy.stream_generator.time') as mock_time: + + MockPS.get_instance.return_value = None + MockConfig.KEEPALIVE_INTERVAL = 0 + MockConfig.MAX_KEEPALIVE_DURATION = 30 + + # First call: keepalive_start_time not yet set (returns current) + # Second call: inside the cap check — simulate time elapsed > 30s + mock_time.time.side_effect = [ + 1000.0, # keepalive_start_time assignment + 1031.0, # cap check: 31s elapsed > 30s limit + ] + + packets = list(gen._stream_data_generator()) + + # No packets should be yielded — cap fires before yield + self.assertEqual(len(packets), 0) + + def test_cap_does_not_fire_before_max_duration(self): + """Generator yields keepalive packets while within MAX_KEEPALIVE_DURATION.""" + gen = _make_generator() + + call_count = 0 + + def time_side_effect(): + nonlocal call_count + call_count += 1 + # keepalive_start_time set at t=1000; cap checks always see <30s elapsed + if call_count == 1: + return 1000.0 # keepalive_start_time + return 1010.0 # always 10s elapsed — under the 30s cap + + with patch.object(gen, '_check_resources', side_effect=[True, True, False]), \ + patch.object(gen, '_should_send_keepalive', return_value=True), \ + patch.object(gen, '_is_ghost_client', return_value=False), \ + patch.object(gen, '_is_timeout', return_value=False), \ + patch('apps.proxy.ts_proxy.stream_generator.create_ts_packet', return_value=b'\x00' * 188), \ + patch('apps.proxy.ts_proxy.stream_generator.ProxyServer') as MockPS, \ + patch('apps.proxy.ts_proxy.stream_generator.Config') as MockConfig, \ + patch('apps.proxy.ts_proxy.stream_generator.gevent'), \ + patch('apps.proxy.ts_proxy.stream_generator.time') as mock_time: + + MockPS.get_instance.return_value = None + MockConfig.KEEPALIVE_INTERVAL = 0 + MockConfig.MAX_KEEPALIVE_DURATION = 30 + mock_time.time.side_effect = time_side_effect + + packets = list(gen._stream_data_generator()) + + # Two iterations with _check_resources=True should yield two keepalive packets + self.assertGreater(len(packets), 0) + + def test_timer_resets_when_real_data_resumes(self): + """keepalive_start_time is cleared to None when real chunks are received.""" + gen = _make_generator() + + chunk = b'\x47' * 188 + real_chunks = ([chunk], gen.local_index + 1) + no_chunks = ([], gen.local_index) + + # Sequence: no data (keepalive), then real data, then stop + gen.buffer.get_optimized_client_data.side_effect = [ + no_chunks, # iteration 1: keepalive + real_chunks, # iteration 2: real data — should reset timer + no_chunks, # iteration 3: keepalive again — timer restarts fresh + ] + + captured_start_times = [] + + original_gen = gen + + with patch.object(gen, '_check_resources', side_effect=[True, True, True, False]), \ + patch.object(gen, '_should_send_keepalive', return_value=True), \ + patch.object(gen, '_is_ghost_client', return_value=False), \ + patch.object(gen, '_is_timeout', return_value=False), \ + patch.object(gen, '_process_chunks', return_value=iter([chunk])), \ + patch('apps.proxy.ts_proxy.stream_generator.create_ts_packet', return_value=b'\x00' * 188), \ + patch('apps.proxy.ts_proxy.stream_generator.ProxyServer') as MockPS, \ + patch('apps.proxy.ts_proxy.stream_generator.Config') as MockConfig, \ + patch('apps.proxy.ts_proxy.stream_generator.gevent'), \ + patch('apps.proxy.ts_proxy.stream_generator.time') as mock_time: + + MockPS.get_instance.return_value = None + MockConfig.KEEPALIVE_INTERVAL = 0 + MockConfig.MAX_KEEPALIVE_DURATION = 300 + mock_time.time.return_value = 1000.0 + + list(gen._stream_data_generator()) + + # Test passes if no exception and generator completes normally — + # if the timer were NOT reset, the second keepalive block would + # carry over the old start time rather than starting fresh. + + def test_cap_uses_config_value(self): + """Cap threshold reads MAX_KEEPALIVE_DURATION from Config, not a hardcoded value.""" + gen = _make_generator() + + with patch.object(gen, '_check_resources', return_value=True), \ + patch.object(gen, '_should_send_keepalive', return_value=True), \ + patch.object(gen, '_is_ghost_client', return_value=False), \ + patch.object(gen, '_is_timeout', return_value=False), \ + patch('apps.proxy.ts_proxy.stream_generator.create_ts_packet', return_value=b'\x00' * 188), \ + patch('apps.proxy.ts_proxy.stream_generator.ProxyServer') as MockPS, \ + patch('apps.proxy.ts_proxy.stream_generator.Config') as MockConfig, \ + patch('apps.proxy.ts_proxy.stream_generator.gevent'), \ + patch('apps.proxy.ts_proxy.stream_generator.time') as mock_time: + + MockPS.get_instance.return_value = None + MockConfig.KEEPALIVE_INTERVAL = 0 + # Set a custom cap of 60s + MockConfig.MAX_KEEPALIVE_DURATION = 60 + + mock_time.time.side_effect = [ + 1000.0, # keepalive_start_time + 1050.0, # cap check: 50s elapsed — under 60s, should NOT fire + 1000.0, # last_yield_time update + 1070.0, # cap check on next iteration: 70s elapsed — fires + ] + + packets = list(gen._stream_data_generator()) + + # First iteration: 50s < 60s cap — one keepalive yielded + # Second iteration: 70s > 60s cap — generator exits + self.assertEqual(len(packets), 1) diff --git a/apps/proxy/config.py b/apps/proxy/config.py index 68e27a6f..0424dee1 100644 --- a/apps/proxy/config.py +++ b/apps/proxy/config.py @@ -108,6 +108,7 @@ class TSConfig(BaseConfig): MIN_STABLE_TIME_BEFORE_RECONNECT = 30 # Minimum seconds a stream must be stable to try reconnect FAILOVER_GRACE_PERIOD = 20 # Extra time (seconds) to allow for stream switching before disconnecting clients URL_SWITCH_TIMEOUT = 20 # Max time allowed for a stream switch operation + MAX_KEEPALIVE_DURATION = 300 # Keepalive packets prevent _is_timeout() from firing, so without this a permanently failed stream holds clients open indefinitely. diff --git a/apps/proxy/ts_proxy/stream_generator.py b/apps/proxy/ts_proxy/stream_generator.py index cfd705c9..2e16008b 100644 --- a/apps/proxy/ts_proxy/stream_generator.py +++ b/apps/proxy/ts_proxy/stream_generator.py @@ -255,6 +255,10 @@ class StreamGenerator: def _stream_data_generator(self): """Generate stream data chunks based on buffer contents.""" + # Keepalive packets refresh last_yield_time, so _is_timeout() never fires + # during sustained stream failure. This timer enforces a wall-clock cap. + keepalive_start_time = None + # Main streaming loop while True: # Check if resources still exist @@ -265,6 +269,7 @@ class StreamGenerator: chunks, next_index = self.buffer.get_optimized_client_data(self.local_index) if chunks: + keepalive_start_time = None # Each recovery restarts the cap independently. yield from self._process_chunks(chunks, next_index) self.local_index = next_index self.last_yield_time = time.time() @@ -306,6 +311,17 @@ class StreamGenerator: continue # Retry immediately with the new position if self._should_send_keepalive(self.local_index): + if keepalive_start_time is None: + keepalive_start_time = time.time() + + max_keepalive = getattr(Config, 'MAX_KEEPALIVE_DURATION', 300) + if time.time() - keepalive_start_time > max_keepalive: + logger.warning( + f"[{self.client_id}] Keepalive duration exceeded {max_keepalive}s " + f"with no stream recovery, disconnecting" + ) + break + keepalive_packet = create_ts_packet('keepalive') logger.debug(f"[{self.client_id}] Sending keepalive packet while waiting at buffer head") yield keepalive_packet