ZYT/dy
Template
2
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

更新

Browse Source
This commit is contained in:
Your Name
2026-07-28 15:04:17 +08:00
parent ac406a5f99
commit 8f68af1c2c
27 changed files with 3442 additions and 296 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/rpa_engine/douyin_im/ws_client.py
+359 -143
View File
@@ -1,9 +1,10 @@
import asyncio
import logging
import threading
import os
import weakref
from typing import Awaitable, Callable, Optional
from websocket import WebSocketApp
from websockets.legacy.client import WebSocketClientProtocol, connect as websocket_connect
from utils import system_logger
from .protocol import parse_ws_payload
@@ -13,9 +14,102 @@ 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:
"""直连 frontier-im WebSocket（websocket-client，与 DouYin_Spider 一致）"""
"""Async frontier-im WebSocket client with bounded message backpressure."""
def __init__(
self,
@@ -29,11 +123,14 @@ class DouyinImWsClient:
self._running = False
self.connected = False
self._task: Optional[asyncio.Task] = None
self._loop: Optional[asyncio.AbstractEventLoop] = None
self._ws_app: Optional[WebSocketApp] = None
self._ws_lock = threading.Lock()
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")
@@ -46,190 +143,309 @@ class DouyinImWsClient:
)
return
self._running = True
self._loop = asyncio.get_running_loop()
self._task = asyncio.create_task(self._run_loop(url))
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
with self._ws_lock:
if self._ws_app:
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:
self._ws_app.close()
connection.fail_connection()
except Exception:
pass
self._ws_app = None
if self._task:
self._task.cancel()
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 self._task
await task
except asyncio.CancelledError:
pass
if self._task is task:
self._task = None
self._connection = None
await self._stop_dispatcher()
async def _run_ws_thread(self, url: str):
"""在独立守护线程中跑 run_forever，直到连接断开/关闭。
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."""
不能用共享默认线程池（run_in_executor(None)/asyncio.to_thread）：
WS 长连接会永久占用一个池线程，账号数超过池大小（默认 64）后，
所有账号的签名/轮询任务被饿死，表现为“启动几十个账号后全部卡死超时”。
"""
loop = asyncio.get_running_loop()
done = asyncio.Event()
error: list[BaseException] = []
def _runner():
try:
self._connect_sync(url)
except BaseException as e:
error.append(e)
finally:
try:
loop.call_soon_threadsafe(done.set)
except RuntimeError:
pass # 事件循环已关闭
thread = threading.Thread(
target=_runner,
name=f"im-ws-{self.account_id or 'na'}",
daemon=True,
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,
)
thread.start()
try:
await done.wait()
except asyncio.CancelledError:
# stop() 会 close ws_app 使 run_forever 退出，线程随之结束
raise
if error:
raise error[0]
return True
async def _run_loop(self, url: str):
retry = 0
while self._running:
from .frontier import ensure_frontier_ws
from .traffic_control import get_traffic_controller
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)
# ensure_frontier_ws 可能触发签名/HTTP（阻塞），放线程池避免卡事件循环
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)
connect_url = self.session.frontier_ws_url() or url
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:
logger.info(f"Connecting IM WebSocket: {connect_url[:100]}...")
await self._run_ws_thread(connect_url)
retry = 0
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 e:
logger.warning(f"IM WebSocket error: {e}")
system_logger.record(
"实时接收连接异常",
detail=f"建立 frontier WebSocket 失败：{e}",
level="error",
category="ws",
account_id=self.account_id,
)
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
# Stable per-account jitter prevents every hosted account from
# reconnecting in the same second after a shared network outage.
jitter = ((int(self.account_id or 0) * 2654435761) % 5000) / 1000.0
wait = min(30.0, 2.0 * retry) + jitter
logger.info(f"IM WebSocket reconnect in {wait:.1f}s...")
system_logger.record(
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",
category="ws",
account_id=self.account_id,
)
await asyncio.sleep(wait)
try:
await asyncio.sleep(wait)
except asyncio.CancelledError:
break
def _connect_sync(self, url: str):
if not self._loop:
return
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
def on_open(_ws):
self.connected = True
logger.info("IM WebSocket connected")
system_logger.record(
"实时接收通道已连接",
detail="frontier WebSocket 已建立，可实时接收私信。",
level="success",
category="ws",
account_id=self.account_id,
)
async def _run_connection(self, url: str) -> None:
"""Open one connection and dispatch messages sequentially.
def on_message(_ws, message):
asyncio.run_coroutine_threadsafe(self._dispatch(message), self._loop)
``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.
"""
def on_error(_ws, error):
if self._running:
logger.warning(f"IM WebSocket error: {error}")
system_logger.record(
"实时接收通道报错",
detail=f"{error}",
level="error",
category="ws",
account_id=self.account_id,
)
def on_close(_ws, code, msg):
self.connected = False
logger.info(f"IM WebSocket closed: code={code}, msg={msg}")
if self._running:
system_logger.record(
"实时接收通道关闭",
detail=f"code={code}, msg={msg}",
level="warning",
category="ws",
account_id=self.account_id,
)
headers = {
"User-Agent": self.session.user_agent,
"Pragma": "no-cache",
"Cache-Control": "no-cache",
"Accept-Language": "zh-CN,zh;q=0.9,en;q=0.8",
"Sec-WebSocket-Protocol": "binary, base64, pbbp2",
"Sec-WebSocket-Extensions": "permessage-deflate; client_max_window_bits",
}
ws_app = WebSocketApp(
url,
header=headers,
cookie=self.session.cookie_header(),
on_open=on_open,
on_message=on_message,
on_error=on_error,
on_close=on_close,
)
with self._ws_lock:
self._ws_app = ws_app
loop = asyncio.get_running_loop()
connected_at: float | None = None
connection: Optional[WebSocketClientProtocol] = None
try:
ws_app.run_forever(origin="https://www.douyin.com", ping_interval=20, ping_timeout=10)
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
with self._ws_lock:
if self._ws_app is ws_app:
self._ws_app = None
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:
await self.on_message(item)
except Exception as e:
logger.debug(f"WS message handler error: {e}")
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"处理收到的私信时出错：{e}",
detail=f"处理收到的私信时出错：{exc}",
level="error",
category="recv",
account_id=self.account_id,
)
finally:
queue.task_done()