WsPool refactoring, TCP fallback backoff, removed direct connections outside of a pool

proxy/bridge.py

@@ -3,7 +3,7 @@ import logging
 import struct
 import time
 
-from typing import List, Optional
+from typing import Dict, List, Optional, Set, Tuple
 from urllib.parse import urlencode
 
 from .utils import *
@@ -20,6 +20,18 @@ from .cf_h2 import bridge_h2
 log = logging.getLogger('tg-mtproto-proxy')
 _st_I_le = struct.Struct('<I')
 
+TCP_BACKOFF_INITIAL = 30.0
+TCP_BACKOFF_MAX = 3600.0
+_tcp_failures: Dict[Tuple[str, int], int] = {}
+_tcp_retry_after: Dict[Tuple[str, int], float] = {}
+_tcp_connecting: Set[Tuple[str, int]] = set()
+
+
+def reset_tcp_backoff() -> None:
+    _tcp_failures.clear()
+    _tcp_retry_after.clear()
+    _tcp_connecting.clear()
+
 ZERO_64 = b'\x00' * 64
 
 
@@ -171,8 +183,6 @@ async def do_fallback(reader, writer, relay_init, label,
             if ok:
                 return True
         elif method == 'tcp' and fallback_dst:
-            log.info("[%s] DC%d%s -> TCP fallback to %s:443",
-                     label, dc, media_tag, fallback_dst)
             ok = await _tcp_fallback(
                 reader, writer, fallback_dst, 443,
                 relay_init, label, ctx)
@@ -271,19 +281,45 @@ async def _cfproxy_fallback(reader, writer, relay_init, label,
 
 
 async def _tcp_fallback(reader, writer, dst, port, relay_init, label, ctx: CryptoCtx):
-    try:
-        rr, rw = await asyncio.wait_for(
-            asyncio.open_connection(dst, port), timeout=10)
-    except Exception as exc:
-        log.warning("[%s] TCP fallback to %s:%d failed: %s",
-                    label, dst, port, repr(exc))
+    key = (dst, port)
+    if key in _tcp_connecting or time.monotonic() < _tcp_retry_after.get(key, 0):
+        log.debug("[%s] TCP fallback to %s:%d skipped (backoff or connecting)",
+                  label, dst, port)
         return False
 
-    stats.connections_tcp_fallback += 1
-    rw.write(relay_init)
-    await rw.drain()
-    await _bridge_tcp_reencrypt(reader, writer, rr, rw, label, ctx)
-    return True
+    _tcp_connecting.add(key)
+    rw = None
+    try:
+        try:
+            log.info("[%s] TCP fallback to %s:%d", label, dst, port)
+            rr, rw = await asyncio.wait_for(
+                asyncio.open_connection(dst, port), timeout=10)
+            rw.write(relay_init)
+            await rw.drain()
+        except Exception as exc:
+            failures = _tcp_failures.get(key, 0) + 1
+            _tcp_failures[key] = failures
+            delay = min(TCP_BACKOFF_INITIAL * 2 ** min(failures - 1, 7),
+                        TCP_BACKOFF_MAX)
+            _tcp_retry_after[key] = time.monotonic() + delay
+            log.warning("[%s] TCP fallback to %s:%d failed: %r; retry in %.0fs",
+                        label, dst, port, exc, delay)
+            return False
+        finally:
+            _tcp_connecting.discard(key)
+
+        _tcp_failures.pop(key, None)
+        _tcp_retry_after.pop(key, None)
+        stats.connections_tcp_fallback += 1
+        await _bridge_tcp_reencrypt(reader, writer, rr, rw, label, ctx)
+        return True
+    finally:
+        if rw is not None:
+            rw.close()
+            try:
+                await rw.wait_closed()
+            except (OSError, ConnectionError):
+                pass
 
 
 async def bridge_ws_reencrypt(reader, writer, ws: RawWebSocket, label,

proxy/pool.py

@@ -10,12 +10,11 @@ from typing import Dict, List, Optional, Tuple, Set
 from .raw_websocket import RawWebSocket, WsHandshakeError
 from .stats import stats
 from .config import proxy_config
-from .utils import ws_domains, DC_DEFAULT_IPS
+from .utils import ws_domains, DC_DEFAULT_IPS, WS_PATH, WS_PATH_TEST
 
 log = logging.getLogger('tg-mtproto-proxy')
 
-# TODO: domains handling is broken: wrong is_media flag causes tcp_reset after handshake,
-# but initial connection is still established no matter what is is_media flag is set to
+
 class _WsPool:
     WS_POOL_MAX_AGE = 120.0
     WS_POOL_CHECK_INTERVAL = 5.0
@@ -23,17 +22,20 @@ class _WsPool:
     REFILL_BACKOFF_MAX = 3600.0
     
     def __init__(self):
-        self._idle: Dict[Tuple[int, bool], deque] = {}
-        self._refilling: Set[Tuple[int, bool]] = set()
-        self._rotating: Dict[Tuple[int, bool], asyncio.Task] = {}
-        self._refill_failures: Dict[Tuple[int, bool], int] = {}
-        self._refill_after: Dict[Tuple[int, bool], float] = {}
-        self.try_fronting_first = True
+        self._idle: Dict[Tuple[int, bool, bool], deque] = {}
+        self._refilling: Dict[Tuple[int, bool, bool], asyncio.Task] = {}
+        self._rotating: Dict[Tuple[int, bool, bool], asyncio.Task] = {}
+        self._refill_failures: Dict[Tuple[int, bool, bool], int] = {}
+        self._refill_after: Dict[Tuple[int, bool, bool], float] = {}
+        self.try_fronting_first = False
 
-    async def get(self, dc: int, is_media: bool,
-                  target_ip: str, domains: List[str]
+    async def get(self, dc: int, is_media: bool, *, is_test_dc: bool = False
                   ) -> Optional[RawWebSocket]:
-        key = (dc, is_media)
+        target_ip = proxy_config.dc_redirects.get(dc)
+        if not target_ip or proxy_config.pool_size <= 0:
+            return None
+        key = (dc, is_media, is_test_dc)
+        domains = ws_domains(dc, is_media)
         now = time.monotonic()
 
         bucket = self._idle.get(key)
@@ -43,14 +45,12 @@ class _WsPool:
         while bucket:
             ws, created = 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.pool_hits += 1
-            log.debug("WS pool hit DC%d%s (age=%.1fs, left=%d)",
-                      dc, 'm' if is_media else '', age, len(bucket))
-            self.report_success(dc, is_media)
+            log.debug("WS pool hit DC%d%s%s (age=%.1fs, left=%d)",
+                      dc, 't' if is_test_dc else '', 'm' if is_media else '', age, len(bucket))
             self._schedule_refill(key, target_ip, domains)
             return ws
 
@@ -58,20 +58,25 @@ class _WsPool:
         self._schedule_refill(key, target_ip, 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, key, target_ip, domains):
+        if proxy_config.pool_size <= 0:
+            return
+        self._schedule_rotation(key, target_ip, domains)
         if (key in self._refilling
                 or time.monotonic() < self._refill_after.get(key, 0)):
             return
-        self._refilling.add(key)
-        asyncio.create_task(self._refill(key, target_ip, domains))
-
-    def report_success(self, dc: int, is_media: bool) -> None:
-        key = (dc, is_media)
-        self._refill_failures.pop(key, None)
-        self._refill_after.pop(key, None)
+        self._refilling[key] = asyncio.create_task(
+            self._refill(key, target_ip, domains))
 
     async def _refill(self, key, target_ip, domains):
-        dc, is_media = key
+        dc, is_media, is_test_dc = key
+        tasks = []
+        adopted = set()
         try:
             bucket = self._idle.setdefault(key, deque())
             needed = proxy_config.pool_size - len(bucket)
@@ -79,19 +84,23 @@ class _WsPool:
                 return
             connected = 0
             tasks = [asyncio.create_task(
-                self._connect_one(target_ip, domains))
+                self._connect_one(target_ip, domains,
+                                  WS_PATH_TEST if is_test_dc else WS_PATH))
                 for _ in range(needed)]
-            for t in tasks:
+            for t in asyncio.as_completed(tasks):
                 try:
                     ws = await t
                     if ws:
+                        if self._refilling.get(key) is not asyncio.current_task():
+                            return
                         bucket.append((ws, time.monotonic()))
+                        adopted.add(ws)
                         connected += 1
-                        self._schedule_rotation(key, target_ip, domains)
-                except Exception:
-                    pass
+                except Exception as exc:
+                    log.debug("WS pool connect failed: %r", exc)
             if connected:
-                self.report_success(dc, is_media)
+                self._refill_failures.pop(key, None)
+                self._refill_after.pop(key, None)
             else:
                 failures = self._refill_failures.get(key, 0) + 1
                 self._refill_failures[key] = failures
@@ -102,12 +111,20 @@ class _WsPool:
                 )
                 self._refill_after[key] = time.monotonic() + delay
                 log.info(
-                    "WS pool refill failed for DC%d%s, retry in %.0fs",
-                    dc, 'm' if is_media else '', delay)
-            log.debug("WS pool refilled DC%d%s: %d ready",
-                      dc, 'm' if is_media else '', len(bucket))
+                    "WS pool refill failed for DC%d%s%s, retry in %.0fs",
+                    dc, 't' if is_test_dc else '', 'm' if is_media else '', delay)
+            log.debug("WS pool refilled DC%d%s%s: %d ready",
+                      dc, 't' if is_test_dc else '', 'm' if is_media else '', len(bucket))
         finally:
-            self._refilling.discard(key)
+            for task in tasks:
+                if not task.done():
+                    task.cancel()
+            results = await asyncio.gather(*tasks, return_exceptions=True)
+            for ws in results:
+                if ws is not None and not isinstance(ws, BaseException) and ws not in adopted:
+                    await self._quiet_close(ws)
+            if self._refilling.get(key) is asyncio.current_task():
+                self._refilling.pop(key, None)
 
     def _schedule_rotation(self, key, target_ip, domains):
         if key in self._rotating:
@@ -116,28 +133,16 @@ class _WsPool:
             self._rotate(key, target_ip, domains))
 
     async def _rotate(self, key, target_ip, domains):
-        dc, is_media = key
+        dc, is_media, is_test_dc = key
         try:
-            while True:
-                bucket = self._idle.get(key)
-                if not bucket:
-                    return
-
-                expires_at = min(
-                    created + self.WS_POOL_MAX_AGE
-                    for _, created in bucket)
-                await asyncio.sleep(min(
-                    self.WS_POOL_CHECK_INTERVAL,
-                    max(0, expires_at - time.monotonic())))
-
+            while proxy_config.pool_size > 0:
+                bucket = self._idle.setdefault(key, deque())
                 now = time.monotonic()
                 expired = []
                 ready = deque()
                 while bucket:
                     ws, created = bucket.popleft()
-                    if (now - created >= self.WS_POOL_MAX_AGE
-                            or ws._closed
-                            or ws.writer.transport.is_closing()):
+                    if self._is_stale(ws, created, now):
                         expired.append(ws)
                     else:
                         ready.append((ws, created))
@@ -147,49 +152,45 @@ class _WsPool:
                     for ws in expired:
                         asyncio.create_task(self._quiet_close(ws))
                     log.debug(
-                        "WS pool rotated DC%d%s: %d stale, %d ready",
-                        dc, 'm' if is_media else '', len(expired), len(bucket))
+                        "WS pool rotated DC%d%s%s: %d stale, %d ready",
+                        dc, 't' if is_test_dc else '', 'm' if is_media else '',
+                        len(expired), len(bucket))
 
                 if len(bucket) < proxy_config.pool_size:
                     self._schedule_refill(key, target_ip, 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 key not in self._refilling and self._refill_after.get(key, 0) > now:
+                    wake_at = min(wake_at, self._refill_after[key])
+                await asyncio.sleep(max(0, wake_at - time.monotonic()))
         finally:
             if self._rotating.get(key) is asyncio.current_task():
                 self._rotating.pop(key, None)
 
-    async def _connect_one(self, target_ip, domains) -> Optional[RawWebSocket]:
+    async def _connect_one(self, target_ip, domains, path=WS_PATH) -> Optional[RawWebSocket]:
         for domain in domains:
-            if self.try_fronting_first:
-                ws = await self._connect_fronted(target_ip, domain)
-                if ws:
-                    return ws
-            try:
-                ws = await RawWebSocket.connect(
-                    target_ip, domain, timeout=8)
-                self.try_fronting_first = False
-                return ws
-            except (asyncio.TimeoutError, ConnectionResetError):
-                if self.try_fronting_first:
-                    return None
-                return await self._connect_fronted(target_ip, domain)
-            except WsHandshakeError as exc:
-                if exc.is_redirect:
+            modes = (True, False) if self.try_fronting_first else (False, True)
+            for fronted in modes:
+                try:
+                    ws = await RawWebSocket.connect(
+                        target_ip, domain, timeout=7 if fronted else 8, path=path,
+                        sni="sprinthost.ru" if fronted else None)
+                except Exception as exc:
+                    stats.ws_errors += 1
+                    log.debug("WS pool connect %s%s via %s (fronting=%s): %r",
+                              domain, path, target_ip, fronted, exc)
                     continue
-                return None
-            except Exception:
-                return None
+                self.try_fronting_first = fronted
+                if fronted:
+                    stats.connections_fronting += 1
+                log.debug("WS pool connected %s%s via %s (fronting=%s)",
+                          domain, path, target_ip, fronted)
+                return ws
         return None
 
-    async def _connect_fronted(self, target_ip, domain) -> Optional[RawWebSocket]:
-        try:
-            ws = await RawWebSocket.connect(
-                target_ip, domain, timeout=7, sni="sprinthost.ru")
-        except Exception:
-            return None
-
-        stats.connections_fronting += 1
-        self.try_fronting_first = True
-        return ws
-
     async def _quiet_close(self, ws):
         try:
             await ws.close()
@@ -202,14 +203,21 @@ class _WsPool:
                 continue
             for is_media in (False, True):
                 domains = ws_domains(dc, is_media)
-                self._schedule_refill((dc, is_media), target_ip, domains)
+                key = (dc, is_media, proxy_config.force_test_dc)
+                self._schedule_refill(key, target_ip, domains)
         log.info("WS pool warmup started for %d DC(s)", len(proxy_config.dc_redirects))
 
     def reset(self):
         loop = asyncio.get_running_loop()
-        for task in self._rotating.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():
+            for ws, _ in bucket:
+                try:
+                    ws.writer.close()
+                except Exception as exc:
+                    log.debug("WS pool close failed: %r", exc)
         self._idle.clear()
         self._refilling.clear()
         self._rotating.clear()
@@ -217,6 +225,11 @@ class _WsPool:
         self._refill_after.clear()
         self.try_fronting_first = False
 
+    async def close(self):
+        tasks = list(self._rotating.values()) + list(self._refilling.values())
+        self.reset()
+        await asyncio.gather(*tasks, return_exceptions=True)
+
 
 class _CfWorkerPool:
     WS_POOL_MAX_AGE = 100.0

proxy/raw_websocket.py

@@ -104,26 +104,26 @@ class RawWebSocket:
             timeout=min(timeout, 10),
         )
         
-        set_sock_opts(writer.transport, proxy_config.buffer_size)
-
-        ws_key = base64.b64encode(os.urandom(16)).decode()
-
-        req = (
-            f'GET {path} HTTP/1.1\r\n'
-            f'Host: {domain}\r\n'
-            f'Upgrade: websocket\r\n'
-            f'Connection: Upgrade\r\n'
-            f'Sec-WebSocket-Key: {ws_key}\r\n'
-            f'Sec-WebSocket-Version: 13\r\n'
-            f'Sec-WebSocket-Protocol: binary\r\n'
-            f'\r\n'
-        )
-
-        writer.write(req.encode())
-        await writer.drain()
-
-        response_lines: list[str] = []
         try:
+            set_sock_opts(writer.transport, proxy_config.buffer_size)
+
+            ws_key = base64.b64encode(os.urandom(16)).decode()
+
+            req = (
+                f'GET {path} HTTP/1.1\r\n'
+                f'Host: {domain}\r\n'
+                f'Upgrade: websocket\r\n'
+                f'Connection: Upgrade\r\n'
+                f'Sec-WebSocket-Key: {ws_key}\r\n'
+                f'Sec-WebSocket-Version: 13\r\n'
+                f'Sec-WebSocket-Protocol: binary\r\n'
+                f'\r\n'
+            )
+
+            writer.write(req.encode())
+            await writer.drain()
+
+            response_lines: list[str] = []
             while True:
                 line = await asyncio.wait_for(reader.readline(),
                                               timeout=timeout)
@@ -131,34 +131,32 @@ class RawWebSocket:
                     break
                 response_lines.append(
                     line.decode('utf-8', errors='replace').strip())
-        except asyncio.TimeoutError:
+
+            if not response_lines:
+                raise WsHandshakeError(0, 'empty response')
+
+            first_line = response_lines[0]
+            parts = first_line.split(' ', 2)
+            try:
+                status_code = int(parts[1]) if len(parts) >= 2 else 0
+            except ValueError:
+                status_code = 0
+
+            if status_code == 101:
+                return RawWebSocket(reader, writer)
+
+            headers: dict[str, str] = {}
+            for hl in response_lines[1:]:
+                if ':' in hl:
+                    k, v = hl.split(':', 1)
+                    headers[k.strip().lower()] = v.strip()
+
+            raise WsHandshakeError(status_code, first_line, headers,
+                                    location=headers.get('location'))
+        except BaseException:
             writer.close()
             raise
 
-        if not response_lines:
-            writer.close()
-            raise WsHandshakeError(0, 'empty response')
-
-        first_line = response_lines[0]
-        parts = first_line.split(' ', 2)
-        try:
-            status_code = int(parts[1]) if len(parts) >= 2 else 0
-        except ValueError:
-            status_code = 0
-
-        if status_code == 101:
-            return RawWebSocket(reader, writer)
-
-        headers: dict[str, str] = {}
-        for hl in response_lines[1:]:
-            if ':' in hl:
-                k, v = hl.split(':', 1)
-                headers[k.strip().lower()] = v.strip()
-
-        writer.close()
-        raise WsHandshakeError(status_code, first_line, headers,
-                                location=headers.get('location'))
-
     async def send(self, data: bytes):
         if self._closed:
             raise ConnectionError("WebSocket closed")

proxy/tg_ws_proxy.py

@@ -10,7 +10,7 @@ import argparse
 import logging
 import socket as _socket
 
-from typing import Dict, Optional, Set, Tuple
+from typing import Optional, Set, Tuple
 
 
 if __name__ == '__main__' and (__package__ is None or __package__ == ''):
@@ -22,8 +22,8 @@ if __name__ == '__main__' and (__package__ is None or __package__ == ''):
 from .utils import *
 from .stats import stats
 from .config import proxy_config, parse_dc_ip_list, start_cfproxy_domain_refresh, coerce_domain_list
-from .bridge import MsgSplitter, CryptoCtx, do_fallback, bridge_ws_reencrypt
-from .raw_websocket import RawWebSocket, WsHandshakeError, set_sock_opts
+from .bridge import MsgSplitter, CryptoCtx, do_fallback, bridge_ws_reencrypt, reset_tcp_backoff
+from .raw_websocket import set_sock_opts
 from .fake_tls import proxy_to_masking_domain, verify_client_hello, build_server_hello, FakeTlsStream, TLS_RECORD_HANDSHAKE
 from .balancer import balancer
 from .pool import ws_pool, cf_worker_pool
@@ -34,14 +34,8 @@ from ._aes import Cipher, algorithms, modes
 
 log = logging.getLogger('tg-mtproto-proxy')
 
-IP_FAIL_COOLDOWN = 3600.0
-DC_FAIL_COOLDOWN = 60.0
-WS_FAIL_TIMEOUT = 2.0
 LISTENER_CHECK_INTERVAL = 5.0
 LISTENER_RESTART_DELAY = 1.0
-ws_blacklist: Set[str] = set()
-dc_fail_until: Dict[str, float] = {}
-ip_fail_until: Dict[str, float] = {}
 cf_h2_pool: Optional[CfH2Pool] = None
 
 
@@ -298,119 +292,11 @@ async def _handle_client(reader, writer, secret: bytes):
         relay_init = _generate_relay_init(proto_tag, dc_idx)
         ctx = _build_crypto_ctx(client_dec_prekey_iv, secret, relay_init)
 
-        dc_key = f'{dc}{"t" if is_test_dc else ""}{"m" if is_media else ""}'
         media_tag = " media" if is_media else ""
-        now = time.monotonic()
-        ws_path = WS_PATH_TEST if is_test_dc else WS_PATH
-        target = proxy_config.dc_redirects.get(dc)
-        is_any_cf_fallback = proxy_config.fallback_cfproxy or proxy_config.cfproxy_worker_domains
-        domains = ws_domains(dc, is_media)
-        ws = None
-
-        # Fallback if DC not in config, if WS blacklisted for this DC/is_media or if connect to ip is timed out
-        if (dc not in proxy_config.dc_redirects
-            or dc_key in ws_blacklist
-            or now < ip_fail_until.get(target, 0) and is_any_cf_fallback):
-
-            if dc not in proxy_config.dc_redirects:
-                log.info("[%s] DC%d not in config -> fallback",
-                         label, dc)
-            elif dc_key in ws_blacklist:
-                log.info("[%s] DC%d%s WS blacklisted -> fallback",
-                         label, dc, media_tag)
-            else:
-                # Try to get WS from pool first, might be accidental timeout
-                ws = await ws_pool.get(
-                    dc, is_media, target, domains
-                ) if not is_test_dc else None
-
-                if not ws:
-                    log.info("[%s] DC%d%s WS connect to %s was timed out -> fallback",
-                            label, dc, media_tag, target)
-                else:
-                    log.info("[%s] DC%d%s WS connect to %s was timed out, but pool hit -> using WS",
-                            label, dc, media_tag, target)
-
-            if not ws:
-                splitter = None
-                try:
-                    splitter = MsgSplitter(relay_init, proto_int)
-                except Exception:
-                    pass
-                ok = await do_fallback(
-                    clt_reader, clt_writer, relay_init, label,
-                    dc, is_test_dc, is_media, media_tag,
-                    ctx, splitter=splitter, h2_pool=cf_h2_pool, proto_tag=proto_tag)
-                if not ok:
-                    log.warning("[%s] DC%d%s no fallback available",
-                                label, dc, media_tag)
-                return
-
-        ws_timeout = WS_FAIL_TIMEOUT if now < dc_fail_until.get(dc_key, 0) else 5.0
-        ws_failed_redirect = False
-        ws_timed_out = False
-        all_redirects = True
-
-        ws = ws or await ws_pool.get(
-            dc, is_media, target, domains
-        ) if not is_test_dc else None
-        if ws:
-            log.info("[%s] DC%d%s -> pool hit via %s",
-                     label, dc, media_tag, target)
-        else:
-            for domain in domains:
-                url = f'wss://{domain}{ws_path}'
-                log.info("[%s] DC%d%s -> %s via %s",
-                         label, dc, media_tag, url, target)
-                try:
-                    ws = await RawWebSocket.connect(target, domain,
-                                                    timeout=ws_timeout,
-                                                    path=ws_path)
-                    all_redirects = False
-                    break
-                except WsHandshakeError as exc:
-                    stats.ws_errors += 1
-                    if exc.is_redirect:
-                        ws_failed_redirect = True
-                        log.warning("[%s] DC%d%s got %d from %s -> %s",
-                                    label, dc, media_tag,
-                                    exc.status_code, domain,
-                                    exc.location or '?')
-                        continue
-                    else:
-                        all_redirects = False
-                        log.warning("[%s] DC%d%s WS handshake: %s",
-                                    label, dc, media_tag, exc.status_line)
-                except asyncio.TimeoutError:
-                    stats.ws_errors += 1
-                    ws_timed_out = True
-                    log.warning("[%s] DC%d%s WS connect timed out via %s",
-                                label, dc, media_tag, domain)
-                    break
-                except Exception as exc:
-                    stats.ws_errors += 1
-                    all_redirects = False
-                    log.warning("[%s] DC%d%s WS connect failed: %s",
-                                label, dc, media_tag, repr(exc))
-
-        # WS failed -> fallback
+        ws = await ws_pool.get(dc, is_media, is_test_dc=is_test_dc)
         if ws is None:
-            if ws_timed_out:
-                ip_fail_until[target] = now + IP_FAIL_COOLDOWN
-                log.info("[%s] DC%d%s WS connect to %s timed out, cooldown for %ds",
-                         label, dc, media_tag, target, int(IP_FAIL_COOLDOWN))
-
-            if ws_failed_redirect and all_redirects:
-                ws_blacklist.add(dc_key)
-                log.warning("[%s] DC%d%s blacklisted for WS (all 302)",
-                            label, dc, media_tag)
-            elif ws_failed_redirect:
-                dc_fail_until[dc_key] = now + DC_FAIL_COOLDOWN
-            else:
-                dc_fail_until[dc_key] = now + DC_FAIL_COOLDOWN
-                log.info("[%s] DC%d%s WS failed for %ds",
-                         label, dc, media_tag, int(DC_FAIL_COOLDOWN))
-
+            log.info("[%s] DC%d%s WS pool unavailable -> fallback",
+                     label, dc, media_tag)
             splitter_fb = None
             try:
                 splitter_fb = MsgSplitter(relay_init, proto_int)
@@ -421,12 +307,12 @@ async def _handle_client(reader, writer, secret: bytes):
                 dc, is_test_dc, is_media, media_tag,
                 ctx, splitter=splitter_fb, h2_pool=cf_h2_pool, proto_tag=proto_tag)
             if ok:
-                log.info("[%s] DC%d%s fallback closed",
-                         label, dc, media_tag)
+                log.info("[%s] DC%d%s fallback closed", label, dc, media_tag)
+            else:
+                log.warning("[%s] DC%d%s no fallback available", label, dc, media_tag)
             return
 
-        ip_fail_until.pop(target, None)
-        ws_pool.report_success(dc, is_media)
+        log.info("[%s] DC%d%s -> WS pool hit", label, dc, media_tag)
         stats.connections_ws += 1
 
         splitter = None
@@ -476,11 +362,9 @@ async def _run(stop_event: Optional[asyncio.Event] = None):
     global _server_instance, _server_stop_event, cf_h2_pool
     _server_stop_event = stop_event
 
+    reset_tcp_backoff()
     ws_pool.reset()
     cf_worker_pool.reset()
-    ws_blacklist.clear()
-    dc_fail_until.clear()
-    ip_fail_until.clear()
     _client_tasks.clear()
     cf_h2_pool = CfH2Pool() if proxy_config.h2_enabled else None
 
@@ -567,8 +451,7 @@ async def _run(stop_event: Optional[asyncio.Event] = None):
                 if now < next_stats:
                     continue
                 next_stats = now + 60
-                bl = ', '.join(f'DC{k}' for k in sorted(ws_blacklist)) or 'none'
-                log.info("stats: %s | ws_bl: %s", stats.summary(), bl)
+                log.info("stats: %s", stats.summary())
                 if cf_h2_pool is not None:
                     cf_h2_pool.log_stats()
         except asyncio.CancelledError:
@@ -650,6 +533,7 @@ async def _run(stop_event: Optional[asyncio.Event] = None):
             if task is not None:
                 await _quiet_cancel(task)
         await server.wait_closed()
+        await ws_pool.close()
         if cf_h2_pool is not None:
             await cf_h2_pool.close()
             cf_h2_pool = None
@@ -695,7 +579,7 @@ def main():
     ap.add_argument('--buf-kb', type=int, default=256, metavar='KB',
                     help='Socket send/recv buffer size in KB (default 256)')
     ap.add_argument('--pool-size', type=int, default=4, metavar='N',
-                    help='WS connection pool size per DC (default 4, min 0)')
+                    help='WS connection pool size per DC (default 4, 0 disables direct WS)')
     ap.add_argument('--cfproxy-domain', action='append', default=None,
                     metavar='DOMAIN',
                     help='User defined Cloudflare-proxied domain for H2 media and WS fallback '

proxy/utils.py

@@ -58,11 +58,11 @@ WS_PATH = '/apiws'
 WS_PATH_TEST = WS_PATH + '_test'
 
 
-def ws_domains(dc: int, is_media) -> List[str]:
+def ws_domains(dc: int, is_media: bool) -> List[str]:
     if dc == 203:
         dc = 2
     if not is_media:
-        return [f'kws{dc}.web.telegram.org', f'kws{dc}-1.web.telegram.org']
+        return [f'kws{dc}.web.telegram.org']
     return [f'kws{dc}-1.web.telegram.org', f'kws{dc}.web.telegram.org']