Added pool refill backoff; don't refill pool on .get if ip is blocked

proxy/pool.py

@@ -16,15 +16,20 @@ 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_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 = False
 
     async def get(self, dc: int, is_media: bool,
-                  target_ip: str, domains: List[str]
+                  target_ip: str, domains: List[str],
+                  *, allow_refill: bool = True
                   ) -> Optional[RawWebSocket]:
         key = (dc, is_media)
         now = time.monotonic()
@@ -43,19 +48,28 @@ class _WsPool:
             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._schedule_refill(key, target_ip, domains)
+            self.report_success(dc, is_media)
+            if allow_refill:
+                self._schedule_refill(key, target_ip, domains)
             return ws
 
         stats.pool_misses += 1
-        self._schedule_refill(key, target_ip, domains)
+        if allow_refill:
+            self._schedule_refill(key, target_ip, domains)
         return None
 
     def _schedule_refill(self, key, target_ip, domains):
-        if key in self._refilling:
+        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)
+
     async def _refill(self, key, target_ip, domains):
         dc, is_media = key
         try:
@@ -63,6 +77,7 @@ class _WsPool:
             needed = proxy_config.pool_size - len(bucket)
             if needed <= 0:
                 return
+            connected = 0
             tasks = [asyncio.create_task(
                 self._connect_one(target_ip, domains))
                 for _ in range(needed)]
@@ -71,9 +86,24 @@ class _WsPool:
                     ws = await t
                     if ws:
                         bucket.append((ws, time.monotonic()))
+                        connected += 1
                         self._schedule_rotation(key, target_ip, domains)
                 except Exception:
                     pass
+            if connected:
+                self.report_success(dc, is_media)
+            else:
+                failures = self._refill_failures.get(key, 0) + 1
+                self._refill_failures[key] = failures
+                delay = min(
+                    self.REFILL_BACKOFF_INITIAL
+                    * (2 ** min(failures - 1, 6)),
+                    self.REFILL_BACKOFF_MAX,
+                )
+                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))
         finally:
@@ -181,6 +211,8 @@ class _WsPool:
         self._idle.clear()
         self._refilling.clear()
         self._rotating.clear()
+        self._refill_failures.clear()
+        self._refill_after.clear()
         self.try_fronting_first = False
 
 

proxy/tg_ws_proxy.py

@@ -339,7 +339,11 @@ async def _handle_client(reader, writer, secret: bytes):
         ws_timed_out = False
         all_redirects = True
 
-        ws = await ws_pool.get(dc, is_media, target, domains) if not is_test_dc else None
+        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,
+        ) if not is_test_dc else None
         if ws:
             log.info("[%s] DC%d%s -> pool hit via %s",
                      label, dc, media_tag, target)
@@ -413,6 +417,7 @@ async def _handle_client(reader, writer, secret: bytes):
 
         dc_fail_until.pop(dc_key, None)
         ip_fail_until.pop(target, None)
+        ws_pool.report_success(dc, is_media)
         stats.connections_ws += 1
 
         splitter = None