diff --git a/proxy/pool.py b/proxy/pool.py index 7f6d1ea..97a9189 100644 --- a/proxy/pool.py +++ b/proxy/pool.py @@ -233,16 +233,24 @@ class _WsPool: class _CfWorkerPool: WS_POOL_MAX_AGE = 100.0 + WS_POOL_CHECK_INTERVAL = 5.0 PER_DC_LIMIT = 1 + REFILL_BACKOFF_INITIAL = 1.0 + REFILL_BACKOFF_MAX = 3600.0 def __init__(self): self._idle: Dict[int, deque] = {} self._refilling: Dict[int, asyncio.Task] = {} + self._rotating: Dict[int, asyncio.Task] = {} + self._refill_failures: Dict[int, int] = {} + self._refill_after: Dict[int, float] = {} self._exhausted_until: Dict[str, float] = {} async def get(self, dc: int, fallback_dst: str, worker_domains: List[str] ) -> Optional[Tuple[RawWebSocket, str]]: + if proxy_config.pool_size <= 0 or not worker_domains: + return None now = time.monotonic() bucket = self._idle.get(dc) @@ -252,8 +260,7 @@ class _CfWorkerPool: while bucket: ws, created, worker_domain = bucket.popleft() age = now - created - if (age > self.WS_POOL_MAX_AGE or ws._closed - or ws.writer.transport.is_closing()): + if self._is_stale(ws, created, now): asyncio.create_task(self._quiet_close(ws)) continue stats.cf_pool_hits += 1 @@ -264,10 +271,20 @@ class _CfWorkerPool: return ws, worker_domain stats.cf_pool_misses += 1 + self._schedule_refill(dc, fallback_dst, worker_domains) return None + def _is_stale(self, ws, created, now): + return (now - created >= self.WS_POOL_MAX_AGE or ws._closed + or ws.writer.transport.is_closing() or ws.reader.at_eof() + or ws.reader.exception() is not None) + def _schedule_refill(self, dc, fallback_dst, worker_domains): - if dc in self._refilling: + if proxy_config.pool_size <= 0 or not worker_domains: + return + self._schedule_rotation(dc, fallback_dst, worker_domains) + if (dc in self._refilling + or time.monotonic() < self._refill_after.get(dc, 0)): return self._refilling[dc] = asyncio.create_task(self._refill( dc, fallback_dst, list(worker_domains))) @@ -279,23 +296,80 @@ class _CfWorkerPool: needed = target_size - len(bucket) if needed <= 0: return - + connected = 0 for _ in range(needed): - connected = await self._connect_one( + result = await self._connect_one( worker_domains, fallback_dst, dc) - if connected is None: - break - ws, worker_domain = connected if self._refilling.get(dc) is not asyncio.current_task(): - await self._quiet_close(ws) + if result is not None: + await self._quiet_close(result[0]) return + if result is None: + break + ws, worker_domain = result bucket.append((ws, time.monotonic(), worker_domain)) + connected += 1 + if connected: + self._refill_failures.pop(dc, None) + self._refill_after.pop(dc, None) + else: + failures = self._refill_failures.get(dc, 0) + 1 + self._refill_failures[dc] = failures + delay = min( + self.REFILL_BACKOFF_INITIAL + * (2 ** min(failures - 1, 12)), + self.REFILL_BACKOFF_MAX, + ) + self._refill_after[dc] = time.monotonic() + delay + log.info("CF worker pool refill failed for DC%d, retry in %.0fs", + dc, delay) log.debug("CF worker pool refilled DC%d: %d ready", dc, len(bucket)) finally: if self._refilling.get(dc) is asyncio.current_task(): self._refilling.pop(dc, None) + def _schedule_rotation(self, dc, fallback_dst, worker_domains): + if dc in self._rotating: + return + self._rotating[dc] = asyncio.create_task( + self._rotate(dc, fallback_dst, list(worker_domains))) + + async def _rotate(self, dc, fallback_dst, worker_domains): + try: + while proxy_config.pool_size > 0: + bucket = self._idle.setdefault(dc, deque()) + now = time.monotonic() + expired = [] + ready = deque() + while bucket: + ws, created, worker_domain = bucket.popleft() + if self._is_stale(ws, created, now): + expired.append(ws) + else: + ready.append((ws, created, worker_domain)) + bucket.extend(ready) + + if expired: + for ws in expired: + asyncio.create_task(self._quiet_close(ws)) + log.debug("CF worker pool rotated DC%d: %d stale, %d ready", + dc, len(expired), len(bucket)) + + if len(bucket) < min(proxy_config.pool_size, self.PER_DC_LIMIT): + self._schedule_refill(dc, fallback_dst, worker_domains) + + wake_at = now + self.WS_POOL_CHECK_INTERVAL + if bucket: + wake_at = min(wake_at, min( + created + self.WS_POOL_MAX_AGE for _, created, _ in bucket)) + if dc not in self._refilling and self._refill_after.get(dc, 0) > now: + wake_at = min(wake_at, self._refill_after[dc]) + await asyncio.sleep(max(0, wake_at - time.monotonic())) + finally: + if self._rotating.get(dc) is asyncio.current_task(): + self._rotating.pop(dc, None) + async def _connect_one(self, worker_domains, fallback_dst, dc): query = urlencode({ 'dst': fallback_dst, @@ -352,7 +426,8 @@ class _CfWorkerPool: if dc not in proxy_config.dc_redirects } - if not cf_fallbacks or not proxy_config.cfproxy_worker_domains: + if (proxy_config.pool_size <= 0 or not cf_fallbacks + or not proxy_config.cfproxy_worker_domains): return worker_domains = list(proxy_config.cfproxy_worker_domains) @@ -363,7 +438,7 @@ class _CfWorkerPool: def reset(self): loop = asyncio.get_running_loop() - for task in self._refilling.values(): + for task in list(self._rotating.values()) + list(self._refilling.values()): if not task.done() and task.get_loop() is loop: task.cancel() for bucket in self._idle.values(): @@ -374,10 +449,13 @@ class _CfWorkerPool: log.debug("CF worker pool close failed: %r", exc) self._idle.clear() self._refilling.clear() + self._rotating.clear() + self._refill_failures.clear() + self._refill_after.clear() self._exhausted_until.clear() async def close(self): - tasks = list(self._refilling.values()) + tasks = list(self._rotating.values()) + list(self._refilling.values()) self.reset() await asyncio.gather(*tasks, return_exceptions=True) diff --git a/proxy/raw_websocket.py b/proxy/raw_websocket.py index 0f66db2..fd2abf6 100644 --- a/proxy/raw_websocket.py +++ b/proxy/raw_websocket.py @@ -225,20 +225,21 @@ class RawWebSocket: return None async def close(self): - if self._closed: - return - self._closed = True try: - self.writer.write( - self._build_frame(self.OP_CLOSE, b'', mask=True)) - await self.writer.drain() - except Exception: - pass - try: - self.writer.close() - await self.writer.wait_closed() - except Exception: - pass + if not self._closed: + self._closed = True + try: + self.writer.write( + self._build_frame(self.OP_CLOSE, b'', mask=True)) + await self.writer.drain() + except Exception: + pass + finally: + try: + self.writer.close() + await self.writer.wait_closed() + except Exception: + pass _WS_CLOSE_REASONS = { 1000: 'normal', 1001: 'going_away', 1002: 'protocol_error', diff --git a/proxy/tg_ws_proxy.py b/proxy/tg_ws_proxy.py index 89e7fc0..5933e41 100644 --- a/proxy/tg_ws_proxy.py +++ b/proxy/tg_ws_proxy.py @@ -42,6 +42,8 @@ log = logging.getLogger('tg-mtproto-proxy') LISTENER_CHECK_INTERVAL = 5.0 LISTENER_RESTART_DELAY = 1.0 +BAD_HANDSHAKE_DRAIN_TIMEOUT = 10.0 +BAD_HANDSHAKE_DRAIN_LIMIT = 64 * 1024 cf_h2_pool: Optional[CfH2Pool] = None @@ -251,6 +253,15 @@ def _build_crypto_ctx(client_dec_prekey_iv, secret, relay_init): return CryptoCtx(clt_decryptor, clt_encryptor, tg_encryptor, tg_decryptor) +async def _drain_bad_handshake(reader): + remaining = BAD_HANDSHAKE_DRAIN_LIMIT + while remaining > 0: + chunk = await reader.read(min(4096, remaining)) + if not chunk: + break + remaining -= len(chunk) + + async def _handle_client(reader, writer, secret: bytes): stats.connections_total += 1 stats.connections_active += 1 @@ -272,8 +283,9 @@ async def _handle_client(reader, writer, secret: bytes): stats.connections_bad += 1 log.warning("[%s] bad handshake (wrong secret or proto)", label) try: - while await clt_reader.read(4096): - pass + await asyncio.wait_for( + _drain_bad_handshake(clt_reader), + timeout=BAD_HANDSHAKE_DRAIN_TIMEOUT) except Exception: pass return diff --git a/tests/test_cf_h2.py b/tests/test_cf_h2.py index 446249b..d4472ce 100644 --- a/tests/test_cf_h2.py +++ b/tests/test_cf_h2.py @@ -812,7 +812,8 @@ class HttpMultiplexTest(unittest.IsolatedAsyncioTestCase): self.assertEqual(len(older), 1) self.assertEqual(len(latest), 4) for first, second in zip(latest[1:], latest[2:]): - self.assertGreaterEqual(second - first, .015) + # Account for subtraction roundoff at the threshold. + self.assertGreaterEqual(second - first, .015 - 1e-9) self.assertFalse(channel.pending) self.assertFalse(channel.replay_pending) self.assertEqual(self.lane.inflight, 0) diff --git a/tests/test_pool.py b/tests/test_pool.py index 38c521e..f39fd81 100644 --- a/tests/test_pool.py +++ b/tests/test_pool.py @@ -9,7 +9,7 @@ from unittest import mock from proxy.config import proxy_config from proxy.pool import _CfWorkerPool, _WsPool from proxy.raw_websocket import RawWebSocket, WsHandshakeError -from proxy.utils import ws_domains +from proxy.utils import DC_DEFAULT_IPS, ws_domains class _StopRotation(Exception): @@ -242,6 +242,154 @@ class WsPoolTest(unittest.IsolatedAsyncioTestCase): class CfWorkerPoolTest(unittest.IsolatedAsyncioTestCase): + async def asyncSetUp(self): + self.pool = _CfWorkerPool() + self.pool.WS_POOL_CHECK_INTERVAL = .01 + for name, value in [ + ('pool_size', 4), + ('dc_redirects', {dc: ip for dc, ip in DC_DEFAULT_IPS.items() if dc != 2}), + ('cfproxy_worker_domains', ['worker.example'])]: + patcher = mock.patch.object(proxy_config, name, value) + patcher.start() + self.addCleanup(patcher.stop) + self.addAsyncCleanup(self.pool.close) + + async def wait_for(self, predicate): + async def wait(): + while not predicate(): + await asyncio.sleep(.001) + await asyncio.wait_for(wait(), 1) + + async def test_failed_warmup_recovers_without_new_client_requests(self): + ready = _open_ws() + self.pool.REFILL_BACKOFF_INITIAL = .03 + with mock.patch.object(self.pool, '_connect_one', + side_effect=[None, (ready, 'worker.example')]) as dial: + await self.pool.warmup() + await self.wait_for(lambda: 2 in self.pool._refill_after) + self.assertFalse(self.pool._refilling) + await self.wait_for(lambda: bool(self.pool._idle.get(2))) + self.assertEqual(dial.await_count, 2) + self.assertIs(self.pool._idle[2][0][0], ready) + self.assertNotIn(2, self.pool._refill_failures) + self.assertNotIn(2, self.pool._refill_after) + + async def test_stale_connections_are_replaced(self): + for state in ('expired', 'eof', 'exception', 'closed', 'closing'): + with self.subTest(state=state): + stale, ready = _open_ws(), _open_ws() + created = time.monotonic() + if state == 'expired': + created -= self.pool.WS_POOL_MAX_AGE + elif state == 'eof': + stale.reader.feed_eof() + elif state == 'exception': + stale.reader.set_exception(ConnectionResetError()) + elif state == 'closed': + stale._closed = True + else: + stale.writer.transport.is_closing = lambda: True + self.pool._idle[2] = deque([(stale, created, 'worker.example')]) + with mock.patch.object(self.pool, '_connect_one', + return_value=(ready, 'worker.example')): + self.assertIsNone(await self.pool.get(2, '192.0.2.2', ['worker.example'])) + await self.wait_for(lambda: bool(self.pool._idle[2])) + self.assertIs(self.pool._idle[2][0][0], ready) + stale.writer.close.assert_called_once() + await self.pool.close() + + async def test_rotation_replaces_idle_connection_without_client_requests(self): + stale, ready = _open_ws(), _open_ws() + self.pool.WS_POOL_MAX_AGE = .05 + self.pool._idle[2] = deque([(stale, time.monotonic(), 'worker.example')]) + with mock.patch.object(self.pool, '_connect_one', + return_value=(ready, 'worker.example')) as dial: + await self.pool.warmup() + await self.wait_for(lambda: stale.writer.close.called) + self.assertIs(self.pool._idle[2][0][0], ready) + stale.writer.close.assert_called_once() + dial.assert_awaited_once() + + async def test_concurrent_misses_start_only_one_background_connection(self): + started = asyncio.Event() + + async def connect(*args): + started.set() + await asyncio.Future() + + with mock.patch.object(self.pool, '_connect_one', side_effect=connect) as dial: + for _ in range(10): + self.assertIsNone(await self.pool.get(2, '192.0.2.2', ['worker.example'])) + await asyncio.wait_for(started.wait(), 1) + self.assertEqual(dial.await_count, 1) + + async def test_failed_refills_use_capped_exponential_backoff(self): + ready = _open_ws() + self.pool.REFILL_BACKOFF_INITIAL = 1 + self.pool.REFILL_BACKOFF_MAX = 2 + with mock.patch.object(self.pool, '_connect_one', + side_effect=[None, None, None, (ready, 'worker.example')]), \ + mock.patch.object(self.pool, '_schedule_rotation'): + for failures, delay in enumerate((1, 2, 2), 1): + self.pool._refill_after.pop(2, None) + started = time.monotonic() + self.pool._schedule_refill(2, '192.0.2.2', ['worker.example']) + await asyncio.gather(*self.pool._refilling.values()) + self.assertFalse(self.pool._refilling) + self.assertEqual(self.pool._refill_failures[2], failures) + self.assertGreaterEqual(self.pool._refill_after[2], started + delay) + self.assertLessEqual(self.pool._refill_after[2], time.monotonic() + delay) + self.pool._refill_after.pop(2) + self.pool._schedule_refill(2, '192.0.2.2', ['worker.example']) + await asyncio.gather(*self.pool._refilling.values()) + self.assertIs(self.pool._idle[2][0][0], ready) + self.assertNotIn(2, self.pool._refill_failures) + self.assertNotIn(2, self.pool._refill_after) + + async def test_misses_during_backoff_do_not_retry_and_shutdown_cancels_retry(self): + self.pool.REFILL_BACKOFF_INITIAL = 60 + with mock.patch.object(self.pool, '_connect_one', return_value=None) as dial: + await self.pool.warmup() + await self.wait_for(lambda: 2 in self.pool._refill_after) + for _ in range(10): + self.assertIsNone(await self.pool.get(2, '192.0.2.2', ['worker.example'])) + await asyncio.sleep(.01) + self.assertEqual(dial.await_count, 1) + self.assertFalse(self.pool._refilling) + self.assertTrue(self.pool._rotating) + tasks = list(self.pool._rotating.values()) + await self.pool.close() + self.assertFalse(self.pool._refilling) + self.assertFalse(self.pool._rotating) + self.assertFalse(self.pool._refill_after) + self.assertFalse(self.pool._refill_failures) + self.assertTrue(all(task.done() for task in tasks)) + + async def test_hit_does_not_reset_failed_refill_backoff(self): + ready = _open_ws() + self.pool._idle[2] = deque([(ready, time.monotonic(), 'worker.example')]) + self.pool._refill_failures[2] = 3 + retry_at = time.monotonic() + 60 + self.pool._refill_after[2] = retry_at + with mock.patch.object(self.pool, '_schedule_rotation'): + self.assertEqual(await self.pool.get(2, '192.0.2.2', ['worker.example']), + (ready, 'worker.example')) + self.assertEqual(self.pool._refill_failures[2], 3) + self.assertEqual(self.pool._refill_after[2], retry_at) + self.assertFalse(self.pool._refilling) + + async def test_zero_size_and_missing_worker_domains_do_not_start_refills(self): + with mock.patch.object(self.pool, '_connect_one') as dial: + with mock.patch.object(proxy_config, 'pool_size', 0): + await self.pool.warmup() + self.assertIsNone(await self.pool.get(2, '192.0.2.2', ['worker.example'])) + self.assertFalse(self.pool._refilling) + self.assertFalse(self.pool._rotating) + self.assertIsNone(await self.pool.get(2, '192.0.2.2', [])) + self.assertFalse(self.pool._refilling) + self.assertFalse(self.pool._rotating) + dial.assert_not_called() + async def test_close_cancels_refill_and_closes_unclaimed_connections(self): pool = _CfWorkerPool() ready = _open_ws() @@ -263,6 +411,9 @@ class CfWorkerPoolTest(unittest.IsolatedAsyncioTestCase): await pool.close() self.assertFalse(pool._idle) self.assertFalse(pool._refilling) + self.assertFalse(pool._rotating) + self.assertFalse(pool._refill_after) + self.assertFalse(pool._refill_failures) ready.writer.close.assert_called_once() idle.writer.close.assert_called_once() diff --git a/tests/test_raw_websocket.py b/tests/test_raw_websocket.py index 7b22d31..a8d1bbc 100644 --- a/tests/test_raw_websocket.py +++ b/tests/test_raw_websocket.py @@ -143,6 +143,53 @@ class ParseCloseTest(unittest.TestCase): self.assertEqual(RawWebSocket._parse_close(b''), (None, '')) +class CloseTest(unittest.IsolatedAsyncioTestCase): + async def test_peer_close_still_closes_transport_without_another_close_frame(self): + reader = asyncio.StreamReader() + reader.feed_data(_raw_frame(RawWebSocket.OP_CLOSE, b'\x03\xe8')) + writer = Mock(drain=AsyncMock(), wait_closed=AsyncMock()) + ws = RawWebSocket(reader, writer) + + self.assertIsNone(await ws.recv()) + await ws.close() + + writer.close.assert_called_once() + writer.wait_closed.assert_awaited_once() + self.assertEqual(writer.write.call_count, 1) + + async def test_repeated_close_sends_only_one_close_frame(self): + writer = Mock(drain=AsyncMock(), wait_closed=AsyncMock()) + ws = RawWebSocket(asyncio.StreamReader(), writer) + + await ws.close() + await ws.close() + + self.assertTrue(ws._closed) + self.assertEqual(writer.write.call_count, 1) + writer.close.assert_called() + + async def test_cancelling_close_during_drain_still_closes_transport(self): + draining = asyncio.Event() + + async def drain(): + draining.set() + await asyncio.Future() + + writer = Mock(drain=AsyncMock(side_effect=drain), wait_closed=AsyncMock()) + ws = RawWebSocket(asyncio.StreamReader(), writer) + task = asyncio.create_task(ws.close()) + try: + await asyncio.wait_for(draining.wait(), 1) + task.cancel() + with self.assertRaises(asyncio.CancelledError): + await task + writer.close.assert_called_once() + writer.wait_closed.assert_awaited_once() + finally: + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + class HandshakeErrorTest(unittest.TestCase): def test_redirect_status_codes(self): for code in (301, 302, 303, 307, 308):