diff --git a/proxy/pool.py b/proxy/pool.py index b86f675..587aeb1 100644 --- a/proxy/pool.py +++ b/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 diff --git a/proxy/tg_ws_proxy.py b/proxy/tg_ws_proxy.py index 51e43ad..e8b94fc 100644 --- a/proxy/tg_ws_proxy.py +++ b/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 diff --git a/proxy/utils.py b/proxy/utils.py index 12ce777..251f6ab 100644 --- a/proxy/utils.py +++ b/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: