try ws_pool if connection is timed out

proxy/pool.py

@@ -19,9 +19,8 @@ log = logging.getLogger('tg-mtproto-proxy')
 class _WsPool:
     WS_POOL_MAX_AGE = 120.0
     WS_POOL_CHECK_INTERVAL = 5.0
-    REFILL_BACKOFF_INITIAL = 60.0
+    REFILL_BACKOFF_INITIAL = 1.0
     REFILL_BACKOFF_MAX = 3600.0
-    # TODO: Lower refill backoff timer
     
     def __init__(self):
         self._idle: Dict[Tuple[int, bool], deque] = {}
@@ -32,8 +31,7 @@ class _WsPool:
         self.try_fronting_first = False  # TODO: https://github.com/Flowseal/tg-ws-proxy/issues/1232
 
     async def get(self, dc: int, is_media: bool,
-                  target_ip: str, domains: List[str],
-                  *, allow_refill: bool = True
+                  target_ip: str, domains: List[str]
                   ) -> Optional[RawWebSocket]:
         key = (dc, is_media)
         now = time.monotonic()
@@ -53,13 +51,11 @@ class _WsPool:
             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)
-            if allow_refill:
-                self._schedule_refill(key, target_ip, domains)
+            self._schedule_refill(key, target_ip, domains)
             return ws
 
         stats.pool_misses += 1
-        if allow_refill:
-            self._schedule_refill(key, target_ip, domains)
+        self._schedule_refill(key, target_ip, domains)
         return None
 
     def _schedule_refill(self, key, target_ip, domains):
@@ -101,7 +97,7 @@ class _WsPool:
                 self._refill_failures[key] = failures
                 delay = min(
                     self.REFILL_BACKOFF_INITIAL
-                    * (2 ** min(failures - 1, 6)),
+                    * (2 ** min(failures - 1, 12)),
                     self.REFILL_BACKOFF_MAX,
                 )
                 self._refill_after[key] = time.monotonic() + delay

proxy/tg_ws_proxy.py

@@ -302,6 +302,8 @@ async def _handle_client(reader, writer, secret: bytes):
         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
@@ -315,38 +317,40 @@ async def _handle_client(reader, writer, secret: bytes):
                 log.info("[%s] DC%d%s WS blacklisted -> fallback",
                          label, dc, media_tag)
             else:
-                log.info("[%s] DC%d%s WS connect to %s was timed out -> fallback",
-                         label, dc, media_tag, target)
-            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)
-            if not ok:
-                log.warning("[%s] DC%d%s no fallback available",
-                            label, dc, media_tag)
-            return
+                # 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)
+                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
-
-        domains = ws_domains(dc, is_media)
-        ws = None
         ws_failed_redirect = False
         ws_timed_out = False
         all_redirects = True
 
-        # TODO: Sometimes connection is established, and sometimes its timed out
-        # So we don't need to fully restrict pool refill, let it try connecting
-        # in the background. May be remove direct connection there completely and
-        # only use pool for them? But need to fix domains handling first
-        allow_pool_refill = now >= ip_fail_until.get(target, 0)
-        ws = await ws_pool.get(
-            dc, is_media, target, domains,
-            allow_refill=allow_pool_refill,
+        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",
@@ -402,7 +406,7 @@ async def _handle_client(reader, writer, secret: bytes):
                 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 cooldown for %ds",
+                log.info("[%s] DC%d%s WS failed for %ds",
                          label, dc, media_tag, int(DC_FAIL_COOLDOWN))
 
             splitter_fb = None
@@ -419,7 +423,6 @@ async def _handle_client(reader, writer, secret: bytes):
                          label, dc, media_tag)
             return
 
-        dc_fail_until.pop(dc_key, None)
         ip_fail_until.pop(target, None)
         ws_pool.report_success(dc, is_media)
         stats.connections_ws += 1

proxy/utils.py

@@ -59,9 +59,9 @@ WS_PATH_TEST = WS_PATH + '_test'
 def ws_domains(dc: int, is_media) -> List[str]:
     if dc == 203:
         dc = 2
-    if is_media is None or is_media:
-        return [f'kws{dc}-1.web.telegram.org', f'kws{dc}.web.telegram.org']
-    return [f'kws{dc}.web.telegram.org', f'kws{dc}-1.web.telegram.org']
+    if not is_media:
+        return [f'kws{dc}.web.telegram.org', f'kws{dc}-1.web.telegram.org']
+    return [f'kws{dc}-1.web.telegram.org', f'kws{dc}.web.telegram.org']
 
 
 def human_bytes(n: int) -> str: