Fixes #1517

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
+            if self._is_stale(ws, created, now):
-                    or ws.writer.transport.is_closing()):
                 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)
 

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(
+            if not self._closed:
-                self._build_frame(self.OP_CLOSE, b'', mask=True))
+                self._closed = True
-            await self.writer.drain()
+                try:
-        except Exception:
+                    self.writer.write(
-            pass
+                        self._build_frame(self.OP_CLOSE, b'', mask=True))
-        try:
+                    await self.writer.drain()
-            self.writer.close()
+                except Exception:
-            await self.writer.wait_closed()
+                    pass
-        except Exception:
+        finally:
-            pass
+            try:
+                self.writer.close()
+                await self.writer.wait_closed()
+            except Exception:
+                pass
 
     _WS_CLOSE_REASONS = {
         1000: 'normal', 1001: 'going_away', 1002: 'protocol_error',

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):
+                await asyncio.wait_for(
-                    pass
+                    _drain_bad_handshake(clt_reader),
+                    timeout=BAD_HANDSHAKE_DRAIN_TIMEOUT)
             except Exception:
                 pass
             return

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)

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()
 

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):