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

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 '