import asyncio import logging import random import time from collections import deque from urllib.parse import urlencode from typing import Dict, List, Optional, Tuple from .raw_websocket import RawWebSocket, WsHandshakeError from .stats import stats from .config import proxy_config from .utils import ws_domains, DC_DEFAULT_IPS, WS_PATH, WS_PATH_TEST log = logging.getLogger('tg-mtproto-proxy') class _WsPool: WS_POOL_MAX_AGE = 120.0 WS_POOL_CHECK_INTERVAL = 5.0 REFILL_BACKOFF_INITIAL = 1.0 REFILL_BACKOFF_MAX = 3600.0 def __init__(self): 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, *, is_test_dc: bool = False ) -> Optional[RawWebSocket]: 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) if bucket is None: bucket = deque() self._idle[key] = bucket while bucket: ws, created = bucket.popleft() age = now - created 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%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 stats.pool_misses += 1 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[key] = asyncio.create_task( self._refill(key, target_ip, domains)) async def _refill(self, key, target_ip, domains): dc, is_media, is_test_dc = key tasks = [] adopted = set() try: bucket = self._idle.setdefault(key, deque()) needed = proxy_config.pool_size - len(bucket) if needed <= 0: return connected = 0 tasks = [asyncio.create_task( self._connect_one(target_ip, domains, WS_PATH_TEST if is_test_dc else WS_PATH)) for _ in range(needed)] 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 except Exception as exc: log.debug("WS pool connect failed: %r", exc) if connected: 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 delay = min( self.REFILL_BACKOFF_INITIAL * (2 ** min(failures - 1, 12)), self.REFILL_BACKOFF_MAX, ) self._refill_after[key] = time.monotonic() + delay log.info( "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: 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: return self._rotating[key] = asyncio.create_task( self._rotate(key, target_ip, domains)) async def _rotate(self, key, target_ip, domains): dc, is_media, is_test_dc = key try: 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 self._is_stale(ws, created, now): expired.append(ws) else: ready.append((ws, created)) bucket.extend(ready) if expired: for ws in expired: asyncio.create_task(self._quiet_close(ws)) log.debug( "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, path=WS_PATH) -> Optional[RawWebSocket]: for domain in domains: 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 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 _quiet_close(self, ws): try: await ws.close() except Exception: pass async def warmup(self): for dc, target_ip in proxy_config.dc_redirects.items(): if target_ip is None: continue for is_media in (False, True): domains = ws_domains(dc, is_media) 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 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() self._refill_failures.clear() 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 PER_DC_LIMIT = 1 def __init__(self): self._idle: Dict[int, deque] = {} self._refilling: Dict[int, asyncio.Task] = {} self._exhausted_until: Dict[str, float] = {} async def get(self, dc: int, fallback_dst: str, worker_domains: List[str] ) -> Optional[Tuple[RawWebSocket, str]]: now = time.monotonic() bucket = self._idle.get(dc) if bucket is None: bucket = deque() self._idle[dc] = bucket while bucket: ws, created, worker_domain = bucket.popleft() age = now - created if (age > self.WS_POOL_MAX_AGE or ws._closed or ws.writer.transport.is_closing()): asyncio.create_task(self._quiet_close(ws)) continue stats.cf_pool_hits += 1 log.debug( "CF worker pool hit DC%d via %s (age=%.1fs, left=%d)", dc, worker_domain, age, len(bucket)) self._schedule_refill(dc, fallback_dst, worker_domains) return ws, worker_domain stats.cf_pool_misses += 1 return None def _schedule_refill(self, dc, fallback_dst, worker_domains): if dc in self._refilling: return self._refilling[dc] = asyncio.create_task(self._refill( dc, fallback_dst, list(worker_domains))) async def _refill(self, dc, fallback_dst, worker_domains): try: bucket = self._idle.setdefault(dc, deque()) target_size = min(proxy_config.pool_size, self.PER_DC_LIMIT) needed = target_size - len(bucket) if needed <= 0: return for _ in range(needed): connected = await self._connect_one( worker_domains, fallback_dst, dc) if connected is None: break ws, worker_domain = connected if self._refilling.get(dc) is not asyncio.current_task(): await self._quiet_close(ws) return bucket.append((ws, time.monotonic(), worker_domain)) log.debug("CF worker pool refilled DC%d: %d ready", dc, len(bucket)) finally: if self._refilling.get(dc) is asyncio.current_task(): self._refilling.pop(dc, None) async def _connect_one(self, worker_domains, fallback_dst, dc): query = urlencode({ 'dst': fallback_dst, 'dc': str(dc), }) path = f'/apiws?{query}' for worker_domain in self.available_domains(worker_domains): try: ws = await RawWebSocket.connect( worker_domain, worker_domain, timeout=8, path=path, secure=not proxy_config.disable_secure) return ws, worker_domain except Exception as exc: self.report_failure(worker_domain, exc) return None def available_domains(self, worker_domains: List[str]) -> List[str]: now = time.time() domains = [] for domain in worker_domains: if domain in domains: continue exhausted_until = self._exhausted_until.get(domain, 0) if exhausted_until > now: continue if exhausted_until: self._exhausted_until.pop(domain, None) domains.append(domain) random.shuffle(domains) return domains def report_failure(self, worker_domain: str, exc: Exception) -> None: return # TODO: check status code after daily limit reached if not isinstance(exc, WsHandshakeError) or exc.status_code != 429: return now = time.time() if self._exhausted_until.get(worker_domain, 0) > now: return exhausted_until = now + (86400 - (now % 86400)) self._exhausted_until[worker_domain] = exhausted_until log.warning( "CF worker %s reached its request limit, disabled for %d seconds", worker_domain, int(exhausted_until - now)) async def _quiet_close(self, ws): try: await ws.close() except Exception: pass async def warmup(self): cf_fallbacks = { dc: ip for dc, ip in DC_DEFAULT_IPS.items() if dc not in proxy_config.dc_redirects } if not cf_fallbacks or not proxy_config.cfproxy_worker_domains: return worker_domains = list(proxy_config.cfproxy_worker_domains) for dc, fallback_dst in cf_fallbacks.items(): self._schedule_refill(dc, fallback_dst, worker_domains) log.info("CF worker pool warmup started for %d DC(s)", len(cf_fallbacks)) def reset(self): loop = asyncio.get_running_loop() for task in 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("CF worker pool close failed: %r", exc) self._idle.clear() self._refilling.clear() self._exhausted_until.clear() async def close(self): tasks = list(self._refilling.values()) self.reset() await asyncio.gather(*tasks, return_exceptions=True) ws_pool = _WsPool() cf_worker_pool = _CfWorkerPool()