import asyncio import logging import os import weakref from typing import Awaitable, Callable, Optional from websockets.legacy.client import WebSocketClientProtocol, connect as websocket_connect from utils import system_logger from .protocol import parse_ws_payload from .session import DouyinImSession logger = logging.getLogger("douyin_im.ws") MessageHandler = Callable[[dict], Awaitable[None]] # Both stages are finite. The transport queue gives the receive coroutine a # small amount of breathing room, while the application queue decouples Pong / # frame reads from potentially slow database and reply work. Once both fill, # backpressure intentionally reaches TCP instead of allocating more tasks. _TRANSPORT_MAX_QUEUE = 4 _APPLICATION_QUEUE_SIZE = 8 _INCOMING_MAX_SIZE = 2**20 _STABLE_CONNECTION_SECONDS = 60.0 _MAX_RECONNECT_BASE_SECONDS = 60.0 _CLOSE_GRACE_SECONDS = 2.0 _PING_TIMEOUT_SECONDS = 120.0 _HANDLER_CONCURRENCY_ENV = "KEFU_WS_HANDLER_CONCURRENCY" _SYSTEM_LOG_THROTTLE_ENV = "KEFU_WS_SYSTEM_LOG_THROTTLE_SECONDS" def _env_int_clamped(name: str, default: int, minimum: int, maximum: int) -> int: try: value = int(os.getenv(name, str(default)) or default) except (TypeError, ValueError): value = default return max(minimum, min(maximum, value)) def _env_float_clamped( name: str, default: float, minimum: float, maximum: float, ) -> float: try: value = float(os.getenv(name, str(default)) or default) except (TypeError, ValueError): value = default return max(minimum, min(maximum, value)) def _handler_concurrency_limit() -> int: # SQLite serializes writes. Eight allows unrelated parsing / reads to # progress without letting an accidental value such as 500 recreate the # original event-loop and database stampede. return _env_int_clamped(_HANDLER_CONCURRENCY_ENV, 8, 1, 32) def _system_log_throttle_seconds() -> float: return _env_float_clamped(_SYSTEM_LOG_THROTTLE_ENV, 300.0, 10.0, 3600.0) class _LoopWsState: """Shared limits for all WS clients owned by one asyncio event loop.""" def __init__(self) -> None: self.handler_slots = asyncio.Semaphore(_handler_concurrency_limit()) self.system_log_last_at: dict[tuple[int, str], float] = {} # asyncio synchronization primitives belong to their creating event loop. # Keeping one weakly-keyed state per loop gives production a process-wide # limit while keeping isolated test loops and uncommon threaded loops safe. _LOOP_STATES: "weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, _LoopWsState]" = ( weakref.WeakKeyDictionary() ) def _get_loop_state() -> _LoopWsState: loop = asyncio.get_running_loop() state = _LOOP_STATES.get(loop) if state is None: state = _LoopWsState() _LOOP_STATES[loop] = state return state def _reconnect_delay(account_id: int | None, retry: int) -> float: """Return exponential backoff with stable, account-specific full jitter. A connection that is accepted and immediately closed is still a failed attempt. The old client reset its retry counter whenever ``run_forever`` returned normally, which kept those accounts reconnecting every 2-7s. This delay reaches a 60-90s range after repeated short-lived connections. """ attempt = max(1, int(retry or 1)) base = min(_MAX_RECONNECT_BASE_SECONDS, float(2 ** min(attempt, 6))) # Spread later retries across half of the base interval. Keep at least a # five-second spread on early retries so a shared outage doesn't reconnect # every account in the same instant. spread = max(5.0, base / 2.0) # Keep one account's fraction stable across attempts. This preserves the # exponential ordering while different accounts remain spread apart. seed = (int(account_id or 0) * 2654435761) & 0xFFFFFFFF fraction = (seed % 10000) / 10000.0 return base + (spread * fraction) class DouyinImWsClient: """Async frontier-im WebSocket client with bounded message backpressure.""" def __init__( self, session: DouyinImSession, on_message: MessageHandler, account_id: int | None = None, ): self.session = session self.on_message = on_message self.account_id = account_id self._running = False self.connected = False self._task: Optional[asyncio.Task] = None self._connection: Optional[WebSocketClientProtocol] = None self._last_connection_lifetime = 0.0 self._message_queue: Optional[asyncio.Queue[dict]] = None self._dispatcher_task: Optional[asyncio.Task] = None async def start(self): if self._task and not self._task.done(): return url = self.session.frontier_ws_url() if not url: logger.warning("No frontier WebSocket URL captured; WS listener disabled") system_logger.record( "实时接收未启用：未获取到 frontier WebSocket 地址", detail="缺少有效的 device_id 或 sessionid，无法建立实时私信通道，将仅依赖 HTTP 轮询。", level="warning", category="ws", account_id=self.account_id, ) return self._running = True self._ensure_dispatcher() self._task = asyncio.create_task( self._run_loop(url), name=f"im-ws-{self.account_id or 'na'}", ) async def stop(self): self._running = False self.connected = False connection = self._connection if connection is not None: try: await asyncio.wait_for( connection.close(code=1000, reason="client stopping"), timeout=_CLOSE_GRACE_SECONDS, ) except asyncio.TimeoutError: # Shutdown iterates over every hosted account. One broken # peer must not consume close_timeout repeatedly and turn a # 500-account shutdown into a many-minute operation. logger.debug("Timed out closing IM WebSocket; aborting transport") try: connection.fail_connection() except Exception: pass except Exception: logger.debug("Failed to close IM WebSocket cleanly", exc_info=True) task = self._task if task and task is not asyncio.current_task() and not task.done(): task.cancel() try: await task except asyncio.CancelledError: pass if self._task is task: self._task = None self._connection = None await self._stop_dispatcher() def _record_connection_system_event( self, event_key: str, message: str, *, detail: str, level: str, ) -> bool: """Persist at most one repeated lifecycle event per account/window.""" loop = asyncio.get_running_loop() state = _get_loop_state() key = (int(self.account_id or 0), event_key) now = loop.time() last_at = state.system_log_last_at.get(key) if last_at is not None and now - last_at < _system_log_throttle_seconds(): return False state.system_log_last_at[key] = now system_logger.record( message, detail=detail, level=level, category="ws", account_id=self.account_id, ) return True def _reset_connection_system_log_throttle(self) -> None: loop = asyncio.get_running_loop() state = _LOOP_STATES.get(loop) if state is None: return account_key = int(self.account_id or 0) state.system_log_last_at.pop((account_key, "connected"), None) state.system_log_last_at.pop((account_key, "retry"), None) def _ensure_dispatcher(self) -> None: if self._dispatcher_task and not self._dispatcher_task.done(): return if self._message_queue is None: self._message_queue = asyncio.Queue(maxsize=_APPLICATION_QUEUE_SIZE) self._dispatcher_task = asyncio.create_task( self._dispatch_loop(), name=f"im-ws-dispatch-{self.account_id or 'na'}", ) async def _stop_dispatcher(self) -> None: task = self._dispatcher_task self._dispatcher_task = None if task and task is not asyncio.current_task() and not task.done(): task.cancel() try: await task except asyncio.CancelledError: pass queue = self._message_queue self._message_queue = None if queue is not None: # Dropped messages must decrement the unfinished counter so tests, # diagnostics, and a later restart can never hang on queue.join(). while True: try: queue.get_nowait() except asyncio.QueueEmpty: break else: queue.task_done() async def _prepare_url(self, fallback_url: str) -> str: from .frontier import ensure_frontier_ws from .traffic_control import get_traffic_controller # Frontier discovery can perform synchronous signing / HTTP work. It # remains in the shared background lane and off the FastAPI event loop. controller = get_traffic_controller() async with controller.background_slot( self.account_id or 0, "websocket prepare", ): await asyncio.to_thread(ensure_frontier_ws, self.session) return self.session.frontier_ws_url() or fallback_url async def _run_loop(self, initial_url: str): retry = 0 first_attempt = True while self._running: self._last_connection_lifetime = 0.0 try: # Startup validation already prepared the captured URL. Avoid # repeating signing / frontier discovery for all 500 accounts # on their first connect; refresh only after a disconnect. if first_attempt and initial_url: connect_url = initial_url else: connect_url = await self._prepare_url(initial_url) first_attempt = False if not connect_url: raise RuntimeError("frontier WebSocket URL is unavailable") logger.info("Connecting IM WebSocket: %s...", connect_url[:100]) await self._run_connection(connect_url) except asyncio.CancelledError: break except Exception as exc: if self._running: logger.warning("IM WebSocket error: %s", exc) self._record_connection_system_event( "retry", "实时接收连接异常", detail=f"建立 frontier WebSocket 失败：{exc}", level="error", ) # Only a genuinely stable connection earns a retry reset. A # successful handshake followed by an immediate normal close must # continue exponential backoff rather than reconnect forever at # the first delay. if self._last_connection_lifetime >= _STABLE_CONNECTION_SECONDS: retry = 0 if not self._running: break retry += 1 wait = _reconnect_delay(self.account_id, retry) logger.info("IM WebSocket reconnect in %.1fs...", wait) self._record_connection_system_event( "retry", f"实时接收断开，{wait:.1f}s 后重连", detail="frontier WebSocket 连接已断开，正在自动重连。", level="warning", ) try: await asyncio.sleep(wait) except asyncio.CancelledError: break def _connection_headers(self) -> list[tuple[str, str]]: headers = [ ("Pragma", "no-cache"), ("Cache-Control", "no-cache"), ("Accept-Language", "zh-CN,zh;q=0.9,en;q=0.8"), ] cookie = self.session.cookie_header() if cookie: headers.append(("Cookie", cookie)) return headers async def _run_connection(self, url: str) -> None: """Open one connection and dispatch messages sequentially. ``max_queue`` bounds the library's receive buffer and ``_dispatch`` feeds one lifecycle-owned, bounded application queue. This receive loop therefore remains responsive to control frames during ordinary database stalls without creating one task per incoming frame. """ loop = asyncio.get_running_loop() connected_at: float | None = None connection: Optional[WebSocketClientProtocol] = None try: async with websocket_connect( url, origin="https://www.douyin.com", subprotocols=["binary", "base64", "pbbp2"], extra_headers=self._connection_headers(), user_agent_header=self.session.user_agent, compression="deflate", open_timeout=10, ping_interval=20, # A handler may legitimately wait up to SQLite's 30s busy # timeout. Leave enough headroom for queued work so a healthy # socket isn't mistaken for a dead peer during that stall. ping_timeout=_PING_TIMEOUT_SECONDS, close_timeout=3, # Frontier frames contain metadata and media URLs rather than # media bytes. A finite frame limit plus a finite queue makes # receive memory genuinely bounded across hundreds of peers. max_size=_INCOMING_MAX_SIZE, max_queue=_TRANSPORT_MAX_QUEUE, ) as websocket: connection = websocket self._connection = websocket connected_at = loop.time() self.connected = True logger.info("IM WebSocket connected") self._record_connection_system_event( "connected", "实时接收通道已连接", detail="frontier WebSocket 已建立，可实时接收私信。", level="success", ) async for raw in websocket: if not self._running: break await self._dispatch(raw) finally: if connected_at is not None: self._last_connection_lifetime = max(0.0, loop.time() - connected_at) if self._last_connection_lifetime >= _STABLE_CONNECTION_SECONDS: # A genuinely healthy session starts a new lifecycle. Its # next outage should be visible immediately rather than # hidden by an old retry window. self._reset_connection_system_log_throttle() self.connected = False if self._connection is connection: self._connection = None if connection is not None: code = connection.close_code reason = connection.close_reason logger.info("IM WebSocket closed: code=%s, msg=%s", code, reason) if self._running: self._record_connection_system_event( "retry", "实时接收通道关闭", detail=f"code={code}, msg={reason}", level="warning", ) async def _dispatch(self, raw): self._ensure_dispatcher() queue = self._message_queue if queue is None: return if isinstance(raw, str): payload = raw.encode("utf-8", errors="ignore") else: payload = raw items = parse_ws_payload(payload) for item in items: if not self._running: return await queue.put(item) async def _dispatch_loop(self) -> None: queue = self._message_queue if queue is None: return handler_slots = _get_loop_state().handler_slots while True: item = await queue.get() try: if self._running: # Every account owns one dispatcher, preserving its FIFO. # The shared semaphore prevents 500 dispatchers from # entering SQLite / reply work at the same instant. A # dispatcher waiting here is directly cancellable by # stop(); no detached per-message task is created. async with handler_slots: if self._running: await self.on_message(item) except asyncio.CancelledError: raise except Exception as exc: logger.debug("WS message handler error: %s", exc) system_logger.record( "实时消息处理失败", detail=f"处理收到的私信时出错：{exc}", level="error", category="recv", account_id=self.account_id, ) finally: queue.task_done()