import asyncio import logging import time from collections import deque from urllib.parse import urlencode from typing import Dict, List, Optional, Tuple, Set from .raw_websocket import RawWebSocket, WsHandshakeError from .stats import stats from .config import proxy_config from .utils import ws_domains, DC_DEFAULT_IPS 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], *, allow_refill: bool = True ) -> Optional[RawWebSocket]: key = (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 (age > self.WS_POOL_MAX_AGE or ws._closed or ws.writer.transport.is_closing()): asyncio.create_task(self._quiet_close(ws)) continue 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.report_success(dc, is_media) if allow_refill: self._schedule_refill(key, target_ip, domains) return ws stats.pool_misses += 1 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 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: 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)) for _ in range(needed)] for t in tasks: try: 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: self._refilling.discard(key) 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 = key try: while True: bucket = self._idle.get(key) if not bucket: return expires_at = min( created + self.WS_POOL_MAX_AGE for _, created in bucket) await asyncio.sleep(min( self.WS_POOL_CHECK_INTERVAL, max(0, expires_at - time.monotonic()))) now = time.monotonic() expired = [] ready = deque() while bucket: ws, created = bucket.popleft() if (now - created >= self.WS_POOL_MAX_AGE or ws._closed or ws.writer.transport.is_closing()): 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: %d stale, %d ready", dc, 'm' if is_media else '', len(expired), len(bucket)) self._schedule_refill(key, target_ip, domains) finally: if self._rotating.get(key) is asyncio.current_task(): self._rotating.pop(key, None) async def _connect_one(self, target_ip, domains) -> Optional[RawWebSocket]: for domain in domains: if self.try_fronting_first: ws = await self._connect_fronted(target_ip, domain) if ws: return ws try: ws = await RawWebSocket.connect( target_ip, domain, timeout=8) self.try_fronting_first = False return ws except asyncio.TimeoutError: if self.try_fronting_first: return None return await self._connect_fronted(target_ip, domain) except WsHandshakeError as exc: if exc.is_redirect: continue return None except Exception: return None return None async def _connect_fronted(self, target_ip, domain) -> Optional[RawWebSocket]: try: ws = await RawWebSocket.connect( target_ip, domain, timeout=7, sni="sprinthost.ru") except Exception: return None stats.connections_fronting += 1 self.try_fronting_first = True return ws 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) self._schedule_refill((dc, is_media), 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 self._rotating.values(): if not task.done() and task.get_loop() is loop: task.cancel() self._idle.clear() self._refilling.clear() self._rotating.clear() self._refill_failures.clear() self._refill_after.clear() self.try_fronting_first = False class _CfWorkerPool: WS_POOL_MAX_AGE = 100.0 def __init__(self): self._idle: Dict[Tuple[int, str], deque] = {} self._refilling: Set[Tuple[int, str]] = set() async def get(self, dc: int, worker_domain: str, fallback_dst: str) -> Optional[RawWebSocket]: now = time.monotonic() key = (dc, worker_domain) 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 (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 (age=%.1fs, left=%d)", dc, age, len(bucket)) self._schedule_refill(key, fallback_dst) return ws stats.cf_pool_misses += 1 self._schedule_refill(key, fallback_dst) return None def _schedule_refill(self, key, fallback_dst): if key in self._refilling: return self._refilling.add(key) asyncio.create_task(self._refill(key, fallback_dst)) async def _refill(self, key, fallback_dst): dc, worker_domain = key try: bucket = self._idle.setdefault(key, deque()) needed = proxy_config.pool_size - len(bucket) if needed <= 0: return tasks = [asyncio.create_task( self._connect_one(worker_domain, fallback_dst, dc)) for _ in range(needed)] for t in tasks: try: ws = await t if ws: bucket.append((ws, time.monotonic())) except Exception: pass log.debug("CF worker pool refilled DC%d: %d ready", dc, len(bucket)) finally: self._refilling.discard(key) async def _connect_one(self, worker_domain, fallback_dst, dc) -> Optional[RawWebSocket]: query = urlencode({ 'dst': fallback_dst, 'dc': str(dc), }) path = f'/apiws?{query}' try: return await RawWebSocket.connect( worker_domain, worker_domain, timeout=8, path=path) except Exception: return None 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 for worker_domain in proxy_config.cfproxy_worker_domains: for dc, fallback_dst in cf_fallbacks.items(): self._schedule_refill((dc, worker_domain), fallback_dst) log.info("CF worker pool warmup started for %d DC(s)", len(cf_fallbacks)) def reset(self): self._idle.clear() self._refilling.clear() ws_pool = _WsPool() cf_worker_pool = _CfWorkerPool()