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

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',

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