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/protocol.py
+84 -5
View File
@@ -1,7 +1,11 @@
import gzip
import json
import logging
import logging.handlers
import os
import queue
import re
import threading
from typing import Any, Optional
from .message_content import (
@@ -21,7 +25,6 @@ from .message_content import (
logger = logging.getLogger("douyin_im.protocol")
import os
from datetime import datetime
@@ -42,6 +45,73 @@ def _is_control_payload(content_json: Any, msg_type: int = 0) -> bool:
return False
_WS_DEBUG_PATH = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "ws_media_debug.log")
_WS_DEBUG_WRITER_LOCK = threading.Lock()
_WS_DEBUG_LOGGER: Optional[logging.Logger] = None
def _bounded_env_int(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))
class _DroppingQueueHandler(logging.handlers.QueueHandler):
"""Never let optional diagnostics block the IM event loop."""
def enqueue(self, record) -> None:
try:
self.queue.put_nowait(record)
except queue.Full:
# Debug output is intentionally lossy under pressure. Receiving
# and replying to messages must always take precedence.
return
def _get_ws_debug_logger() -> logging.Logger:
global _WS_DEBUG_LOGGER
if _WS_DEBUG_LOGGER is not None:
return _WS_DEBUG_LOGGER
with _WS_DEBUG_WRITER_LOCK:
if _WS_DEBUG_LOGGER is not None:
return _WS_DEBUG_LOGGER
max_bytes = _bounded_env_int(
"KEFU_WS_DEBUG_MAX_BYTES", 10 * 1024 * 1024, 1024 * 1024, 100 * 1024 * 1024
)
backup_count = _bounded_env_int(
"KEFU_WS_DEBUG_BACKUP_COUNT", 2, 1, 10
)
queue_size = _bounded_env_int(
"KEFU_WS_DEBUG_QUEUE_SIZE", 1000, 100, 10000
)
records: queue.Queue = queue.Queue(maxsize=queue_size)
rotating = logging.handlers.RotatingFileHandler(
_WS_DEBUG_PATH,
maxBytes=max_bytes,
backupCount=backup_count,
encoding="utf-8",
delay=True,
)
rotating.setFormatter(logging.Formatter("%(message)s"))
listener = logging.handlers.QueueListener(
records,
rotating,
respect_handler_level=True,
)
listener.start()
debug_logger = logging.getLogger("douyin_im.ws_raw_debug")
debug_logger.handlers.clear()
debug_logger.addHandler(_DroppingQueueHandler(records))
debug_logger.setLevel(logging.INFO)
debug_logger.propagate = False
# Keep strong references for the lifetime of the logger/listener.
debug_logger._kefu_queue_listener = listener # type: ignore[attr-defined]
debug_logger._kefu_rotating_handler = rotating # type: ignore[attr-defined]
_WS_DEBUG_LOGGER = debug_logger
return debug_logger
def _should_emit_ws_message(
@@ -84,8 +154,13 @@ def _dump_ws_message(msg_type: int, conversation_id: str, content_str: str, msg:
f"{datetime.now().isoformat()} type={msg_type} "
f"conv={conversation_id} content={content_str}{extra}\n"
)
with open(_WS_DEBUG_PATH, "a", encoding="utf-8") as fh:
fh.write(line)
record_limit = _bounded_env_int(
"KEFU_WS_DEBUG_RECORD_MAX_CHARS", 16384, 1024, 262144
)
if len(line) > record_limit:
marker = "...[单条调试记录过长，已截断]\n"
line = line[: max(0, record_limit - len(marker))] + marker
_get_ws_debug_logger().info(line.rstrip("\n"))
except Exception:
pass
@@ -153,8 +228,6 @@ def parse_ws_payload(raw: bytes | str) -> list[dict]:
content_str = msg.content
server_message_id = str(getattr(msg, "server_message_id", "") or "")
_dump_ws_message(msg_type, conversation_id, content_str, msg)
text_content = ""
media_msg: dict = {}
content_json: dict = {}
@@ -175,6 +248,12 @@ def parse_ws_payload(raw: bytes | str) -> list[dict]:
)
return messages
# Raw diagnostics are optional and intentionally run
# only after control/status frames have been filtered.
# The writer itself is queued and rotating, so it can
# never block message parsing or grow without bound.
_dump_ws_message(msg_type, conversation_id, content_str, msg)
if _should_emit_ws_message(conversation_id, msg_type):
sender_uid = str(msg.sender)
if media_msg and (
backend/rpa_engine/douyin_im/service.py
+455 -30
View File
@@ -1,7 +1,11 @@
import asyncio
import inspect
import logging
import os
import time
import weakref
from collections import deque
from dataclasses import dataclass, field
from typing import Awaitable, Callable, Optional
from utils import system_logger
@@ -25,6 +29,14 @@ logger = logging.getLogger("douyin_im.service")
MatchReplyFn = Callable[[str], Awaitable[Optional[list[str]]]]
LogFn = Callable[..., Awaitable[None]]
ReceivedLogFn = Callable[..., Awaitable[None]]
ReadyFn = Callable[[], Optional[Awaitable[None]]]
# These sets live on the FastAPI event-loop thread. They let every account
# derive a polling interval from the current population instead of assuming it
# is the only hosted account on the server.
_active_service_ids: set[int] = set()
_ws_service_ids: set[int] = set()
def _env_poll_seconds(name: str, default: float, minimum: float = 5.0) -> float:
@@ -34,7 +46,11 @@ def _env_poll_seconds(name: str, default: float, minimum: float = 5.0) -> float:
return default
def _conversation_poll_timing(account_id: int, has_ws: bool) -> tuple[float, float]:
def _conversation_poll_timing(
account_id: int,
has_ws: bool,
population: Optional[int] = None,
) -> tuple[float, float]:
"""Return the reconciliation interval and a stable per-account stagger.
WebSocket is the real-time receive path. HTTP polling is only a safety
@@ -42,15 +58,253 @@ def _conversation_poll_timing(account_id: int, has_ws: bool) -> tuple[float, flo
hundreds of accounts wastes bandwidth and eventually starves new starts.
Accounts without WebSocket keep the original fast polling cadence.
"""
interval = _env_poll_seconds(
base_interval = _env_poll_seconds(
"KEFU_WS_RECONCILE_INTERVAL_SECONDS" if has_ws else "KEFU_HTTP_POLL_INTERVAL_SECONDS",
120.0 if has_ws else 15.0,
)
if population is None:
if has_ws:
population = len(_ws_service_ids)
else:
population = len(_active_service_ids - _ws_service_ids)
population = max(1, int(population or 0))
# Reserve a bounded request budget for each class. With 500 connected
# accounts and the default 1 req/s budget, reconciliation automatically
# stretches to 500 seconds instead of permanently saturating the two
# shared background-network slots. Disconnected accounts use a separate
# budget so they cannot create a 15-second retry storm after an outage.
budget_name = "KEFU_WS_POLL_BUDGET_RPS" if has_ws else "KEFU_HTTP_POLL_BUDGET_RPS"
try:
budget_rps = max(0.05, float(os.getenv(budget_name, "1.0")))
except (TypeError, ValueError):
budget_rps = 1.0
interval = max(base_interval, population / budget_rps)
spread_ms = max(1, int(interval * 1000))
stagger = ((int(account_id or 0) * 2654435761) % spread_ms) / 1000.0
return interval, stagger
def _initial_unread_concurrency() -> int:
try:
return max(
1,
min(8, int(os.getenv("KEFU_INITIAL_UNREAD_CONCURRENCY", "2"))),
)
except (TypeError, ValueError):
return 2
@dataclass(eq=False)
class _InitialUnreadJob:
service: "DouyinImService"
messages: deque[dict] = field(default_factory=deque)
cancelled: bool = False
active_task: Optional[asyncio.Task] = None
class _InitialUnreadDispatcher:
"""Process startup unread snapshots with a fixed process-wide worker set.
One queued object represents one account and is requeued after each
message. This keeps account-local FIFO while preventing a single account
from monopolizing both consumers. Only active consumers create handler
tasks, so a 500-account start does not create 500 waiting tasks.
"""
def __init__(self, concurrency: Optional[int] = None) -> None:
self.concurrency = max(
1,
int(concurrency or _initial_unread_concurrency()),
)
self._queue: asyncio.Queue[_InitialUnreadJob] = asyncio.Queue()
self._jobs: dict[int, _InitialUnreadJob] = {}
self._workers: list[asyncio.Task] = []
self._lock = asyncio.Lock()
self._stopping = False
async def _ensure_workers(self) -> None:
async with self._lock:
self._workers = [task for task in self._workers if not task.done()]
if self._workers or self._stopping:
return
for index in range(self.concurrency):
self._workers.append(
asyncio.create_task(
self._worker(index + 1),
name=f"initial-unread-worker-{index + 1}",
)
)
async def submit(
self,
service: "DouyinImService",
messages: list[dict],
) -> None:
pending = [message for message in messages if isinstance(message, dict)]
if not pending or not service._running:
return
await self._ensure_workers()
async with self._lock:
if self._stopping or not service._running:
return
service_key = id(service)
existing = self._jobs.get(service_key)
if existing and not existing.cancelled:
existing.messages.extend(pending)
return
job = _InitialUnreadJob(
service=service,
messages=deque(pending),
)
self._jobs[service_key] = job
self._queue.put_nowait(job)
logger.debug(
"Initial unread queued account=%s count=%s",
service.account_id,
len(pending),
)
async def cancel(self, service: "DouyinImService") -> None:
operation: Optional[asyncio.Task] = None
current_task = asyncio.current_task()
async with self._lock:
job = self._jobs.pop(id(service), None)
if job is None:
return
job.cancelled = True
job.messages.clear()
operation = job.active_task
if (
operation
and operation is not current_task
and not operation.done()
):
operation.cancel()
if (
operation
and operation is not current_task
and not operation.done()
):
await asyncio.gather(operation, return_exceptions=True)
async def _worker(self, worker_number: int) -> None:
while True:
job = await self._queue.get()
operation: Optional[asyncio.Task] = None
try:
async with self._lock:
if (
job.cancelled
or not job.service._running
or not job.messages
):
self._jobs.pop(id(job.service), None)
else:
message = job.messages.popleft()
operation = asyncio.create_task(
job.service._handle_incoming(message),
name=(
"initial-unread-account-"
f"{job.service.account_id}"
),
)
job.active_task = operation
if operation is not None:
try:
await operation
except asyncio.CancelledError:
# Cancelling one account cancels only its bounded child
# operation; cancelling this worker still shuts it down.
if asyncio.current_task().cancelling():
raise
except Exception as exc:
logger.error(
"Initial unread handling failed account=%s "
"worker=%s: %s",
job.service.account_id,
worker_number,
exc,
)
async with self._lock:
if job.active_task is operation:
job.active_task = None
if (
job.cancelled
or not job.service._running
or not job.messages
):
self._jobs.pop(id(job.service), None)
else:
# Round-robin across accounts while retaining FIFO
# within this account's initial unread snapshot.
self._queue.put_nowait(job)
except asyncio.CancelledError:
if operation and not operation.done():
operation.cancel()
await asyncio.gather(operation, return_exceptions=True)
raise
finally:
self._queue.task_done()
async def join(self) -> None:
await self._queue.join()
async def stop(self) -> None:
async with self._lock:
if self._stopping:
workers = list(self._workers)
operations = []
else:
self._stopping = True
workers = list(self._workers)
operations = [
job.active_task
for job in self._jobs.values()
if job.active_task and not job.active_task.done()
]
for job in self._jobs.values():
job.cancelled = True
job.messages.clear()
self._jobs.clear()
for operation in operations:
operation.cancel()
for worker in workers:
worker.cancel()
if operations:
await asyncio.gather(*operations, return_exceptions=True)
if workers:
await asyncio.gather(*workers, return_exceptions=True)
while True:
try:
self._queue.get_nowait()
except asyncio.QueueEmpty:
break
else:
self._queue.task_done()
self._workers.clear()
_INITIAL_UNREAD_DISPATCHERS = weakref.WeakKeyDictionary()
def _get_initial_unread_dispatcher() -> _InitialUnreadDispatcher:
loop = asyncio.get_running_loop()
dispatcher = _INITIAL_UNREAD_DISPATCHERS.get(loop)
if dispatcher is None:
dispatcher = _InitialUnreadDispatcher()
_INITIAL_UNREAD_DISPATCHERS[loop] = dispatcher
return dispatcher
async def _shutdown_initial_unread_dispatcher() -> None:
loop = asyncio.get_running_loop()
dispatcher = _INITIAL_UNREAD_DISPATCHERS.pop(loop, None)
if dispatcher is not None:
await dispatcher.stop()
class DouyinImService:
"""抖音 IM 直连服务：WebSocket 实时监听 + HTTP 轮询 + 自动回复"""
@@ -68,6 +322,7 @@ class DouyinImService:
refresh_credentials: Optional[Callable[[], Awaitable[bool]]] = None,
follow_tick: Optional[Callable[[], Awaitable[None]]] = None,
on_session_invalid: Optional[Callable[[str], Awaitable[None]]] = None,
on_ready: Optional[ReadyFn] = None,
):
self.session = session
self.match_reply = match_reply
@@ -78,6 +333,8 @@ class DouyinImService:
self.follow_tick = follow_tick
# 由 worker 注入：检测到 IM 登录失效（INVALID_REQUEST）时回调，用于自动下线
self.on_session_invalid = on_session_invalid
self._on_ready = on_ready
self._ready_notified = False
self._session_invalid_strikes = 0
self._session_invalid_fired = False
self.reply_delay_seconds = max(0, int(reply_delay_seconds or 0))
@@ -109,6 +366,16 @@ class DouyinImService:
self._ws_client: Optional[DouyinImWsClient] = None
self.last_error: str = ""
async def _notify_ready(self) -> None:
if self._ready_notified:
return
self._ready_notified = True
if not self._on_ready:
return
result = self._on_ready()
if inspect.isawaitable(result):
await result
def _reply_key(self, conversation_key: str, content: str) -> str:
return f"{conversation_key}::{content}"
@@ -725,7 +992,12 @@ class DouyinImService:
f"(sent={sent_any}, count={len(replies)})"
)
async def _index_conversations(self, conversations: list[dict]):
async def _index_conversations(
self,
conversations: list[dict],
*,
enrich_profiles: bool = True,
):
my_uid = int(self.session.my_uid or 0)
for raw in conversations:
conv = enrich_conversation_item(raw, my_uid)
@@ -734,7 +1006,11 @@ class DouyinImService:
avatar = str(conv.get("sender_avatar") or "").strip()
peer_uid = str(conv.get("peer_uid") or "")
if peer_uid and (is_generic_peer_name(name, peer_uid) or not avatar):
if (
enrich_profiles
and peer_uid
and (is_generic_peer_name(name, peer_uid) or not avatar)
):
profile = await fetch_peer_profile(self.session, peer_uid, self.account_id)
if profile.get("nickname"):
name = profile["nickname"]
@@ -755,16 +1031,67 @@ class DouyinImService:
if peer_uid and name:
self._conv_names[peer_uid] = name
async def _poll_conversations(self):
async def _poll_conversations(
self,
*,
initial: bool = False,
defer_handlers: bool = False,
) -> list[dict]:
controller = get_traffic_controller()
async with controller.background_slot(self.account_id, "conversation poll"):
async with DouyinImHttpClient(self.session, account_id=self.account_id) as http:
async with controller.background_slot(
self.account_id,
"conversation poll",
startup=initial,
):
# Do not keep one idle connection pool per hosted account. At 500
# accounts that would retain hundreds of sockets between sparse
# reconciliations. Capacity-aware scheduling makes construction
# infrequent, while the context manager releases the socket as
# soon as this account's poll finishes.
async with DouyinImHttpClient(
self.session,
account_id=self.account_id,
) as http:
conversations = await http.get_conversations(enrich_profiles=False)
# Profile enrichment may involve several slow third-party requests.
# Run it after releasing the conversation-list slot; each individual
# lookup re-enters the shared controller and yields fairly to startup
# validation and other accounts between profiles.
await self._index_conversations(conversations)
# Capture the previous preview before _index_conversations overwrites
# _conv_meta. A conversation-list preview is not inherently a new
# message: after startup we only act when it is unread or has actually
# changed since the last cache snapshot. This prevents the first
# scheduled reconciliation from replying to every historical preview.
previous_previews: list[tuple[bool, str]] = []
previous_by_peer = {
str(meta.get("peer_uid") or ""): str(meta.get("content") or "")
for meta in self._conv_meta.values()
if str(meta.get("peer_uid") or "")
}
for conv in conversations:
conv_id = str(conv.get("conversation_id") or "").strip()
peer_uid = str(
conv.get("peer_uid") or conv.get("sender_id") or ""
).strip()
sender_name = str(conv.get("sender_name") or "").strip()
prior: Optional[str] = None
known = False
if conv_id and conv_id in self._conv_meta:
known = True
prior = str(self._conv_meta[conv_id].get("content") or "")
elif sender_name and sender_name in self._conv_previews:
known = True
prior = str(self._conv_previews.get(sender_name) or "")
elif peer_uid and peer_uid in previous_by_peer:
known = True
prior = previous_by_peer[peer_uid]
previous_previews.append((known, prior or ""))
# Never enrich every row in a reconciliation snapshot. At 500 accounts
# that could turn one list request into tens of thousands of profile
# calls. _handle_incoming resolves the peer lazily only for an unread
# or genuinely changed conversation after this lightweight cache pass.
await self._index_conversations(
conversations,
enrich_profiles=False,
)
unread_total = sum(
max(0, int(item.get("unread_count") or 0))
for item in conversations
@@ -773,10 +1100,27 @@ class DouyinImService:
logger.info(f"IM unread total: {unread_total}")
# Message handling may wait in the global send lane. Do not keep one
# of the scarce background HTTP slots occupied while that happens.
for conv in conversations:
deferred: list[dict] = []
for conv, (preview_known, previous_preview) in zip(
conversations,
previous_previews,
):
unread = int(conv.get("unread_count") or 0)
if unread > 0 or conv.get("content"):
await self._handle_incoming(conv)
current_preview = str(conv.get("content") or "")
preview_changed = bool(
preview_known
and current_preview
and current_preview != previous_preview
)
# Existing previews are useful cache data, but on the authoritative
# startup snapshot they are not proof of a newly received message.
# The same is true on later polls unless the cached preview changed.
if unread > 0 or (not initial and preview_changed):
if defer_handlers:
deferred.append(conv)
else:
await self._handle_incoming(conv)
return deferred
async def _verify_account_uid(self):
"""启动时用 query/user 接口核验账号真实 UID，修正采集端可能取错的 my_uid/device_id。
@@ -790,7 +1134,11 @@ class DouyinImService:
from .auth import DouyinAuth
auth = DouyinAuth.from_im_session(self.session)
controller = get_traffic_controller()
async with controller.background_slot(self.account_id, "account UID verify"):
async with controller.background_slot(
self.account_id,
"account UID verify",
startup=True,
):
async with DouyinImHttpClient(self.session, account_id=self.account_id) as http:
old = int(self.session.my_uid or 0)
resolved = await asyncio.to_thread(http._resolve_authoritative_uid, auth)
@@ -808,11 +1156,16 @@ class DouyinImService:
async def run(self):
"""主循环：WebSocket + HTTP 轮询"""
self._running = True
_active_service_ids.add(self.account_id)
await self._reply_queue.start()
await self._verify_account_uid()
# ensure_frontier_ws 可能触发签名/HTTP（阻塞），放线程池避免多账号启动时卡死事件循环
controller = get_traffic_controller()
async with controller.background_slot(self.account_id, "frontier discovery"):
async with controller.background_slot(
self.account_id,
"frontier discovery",
startup=True,
):
await asyncio.to_thread(ensure_frontier_ws, self.session)
has_ws = bool(self.session.frontier_ws_url())
cred_summary = format_session_credential_summary(self.session)
@@ -833,7 +1186,11 @@ class DouyinImService:
from .emoji_pack import ensure_emoji_map, is_fresh
if not is_fresh():
async with controller.background_slot(self.account_id, "emoji preload"):
async with controller.background_slot(
self.account_id,
"emoji preload",
startup=True,
):
await asyncio.to_thread(ensure_emoji_map, self.session)
except Exception as e:
logger.debug(f"emoji map preload failed: {e}")
@@ -844,8 +1201,12 @@ class DouyinImService:
await self._ws_client.start()
initial_poll_succeeded = False
initial_unread: list[dict] = []
try:
await self._poll_conversations()
initial_unread = await self._poll_conversations(
initial=True,
defer_handlers=True,
)
initial_poll_succeeded = True
except Exception as e:
logger.warning(f"Initial conversation poll failed: {e}")
@@ -857,9 +1218,25 @@ class DouyinImService:
account_id=self.account_id,
)
ws_connected = bool(
self._ws_client and getattr(self._ws_client, "connected", False)
)
# Batch readiness covers transport discovery plus the authoritative
# list/cache snapshot, not potentially slow reply generation, logging,
# or outbound sends for pre-existing unread conversations.
await self._notify_ready()
if initial_unread:
await _get_initial_unread_dispatcher().submit(
self,
initial_unread,
)
# A discovered WS URL is enough to start on the low-frequency
# reconciliation schedule. If the socket does not actually open, the
# first health tick switches the account to fallback mode. This avoids
# a second fast HTTP poll just because the handshake needed a moment.
ws_connected = bool(has_ws)
if ws_connected:
_ws_service_ids.add(self.account_id)
else:
_ws_service_ids.discard(self.account_id)
poll_interval, poll_stagger = _conversation_poll_timing(
self.account_id,
ws_connected,
@@ -884,15 +1261,23 @@ class DouyinImService:
poll_stagger,
"yes" if ws_connected else "no",
)
loop_count = 0
poll_failures = 0
follow_interval = _env_poll_seconds(
"KEFU_FOLLOW_POLL_INTERVAL_SECONDS",
60.0,
minimum=30.0,
)
follow_stagger = (
(int(self.account_id or 0) * 2654435761)
% max(1, int(follow_interval * 1000))
) / 1000.0
next_follow_tick_at = loop.time() + follow_interval + follow_stagger
while self._running:
# The initial poll above is authoritative. Sleep before the next
# recurring tick so startup cannot issue two back-to-back polls.
await asyncio.sleep(5)
if not self._running:
break
loop_count += 1
try:
current_ws_connected = bool(
self._ws_client
@@ -900,6 +1285,10 @@ class DouyinImService:
)
if current_ws_connected != ws_connected:
ws_connected = current_ws_connected
if ws_connected:
_ws_service_ids.add(self.account_id)
else:
_ws_service_ids.discard(self.account_id)
poll_interval, poll_stagger = _conversation_poll_timing(
self.account_id,
ws_connected,
@@ -919,22 +1308,53 @@ class DouyinImService:
poll_interval,
"yes" if ws_connected else "no",
)
else:
# Large batches change the active population while older
# services are already running. Future polls adopt the
# widened capacity interval without moving a poll that is
# already scheduled.
capacity_interval, _ = _conversation_poll_timing(
self.account_id,
ws_connected,
)
if abs(capacity_interval - poll_interval) >= 1.0:
poll_interval = capacity_interval
if loop.time() >= next_conversation_poll_at:
try:
await self._poll_conversations()
max_poll_seconds = _env_poll_seconds(
"KEFU_CONVERSATION_POLL_DEADLINE_SECONDS",
30.0,
)
await asyncio.wait_for(
self._poll_conversations(),
timeout=max_poll_seconds,
)
poll_failures = 0
except asyncio.TimeoutError:
poll_failures = min(4, poll_failures + 1)
logger.debug(
"Dropping stale conversation poll account=%s after %.1fs",
self.account_id,
max_poll_seconds,
)
except Exception:
poll_failures = min(4, poll_failures + 1)
raise
finally:
# Advance on both success and failure. Otherwise a
# past deadline retries every five-second loop tick
# during an outage and amplifies traffic.
next_conversation_poll_at = loop.time() + poll_interval
if loop_count % 6 == 0:
logger.info(f"IM direct tick #{loop_count} account={self.account_id}")
next_conversation_poll_at = (
loop.time() + poll_interval * (2 ** poll_failures)
)
# 关注欢迎语：约每 60s 检测一次新粉丝（独立于私信轮询，失败不影响主循环）
if self.follow_tick and loop_count % 12 == 0:
if self.follow_tick and loop.time() >= next_follow_tick_at:
try:
await self.follow_tick()
except Exception as e:
logger.error(f"follow welcome tick error: {e}")
finally:
next_follow_tick_at = loop.time() + follow_interval
except Exception as e:
logger.error(f"IM poll error: {e}")
system_logger.record(
@@ -947,6 +1367,11 @@ class DouyinImService:
async def stop(self):
self._running = False
_active_service_ids.discard(self.account_id)
_ws_service_ids.discard(self.account_id)
dispatcher = _INITIAL_UNREAD_DISPATCHERS.get(asyncio.get_running_loop())
if dispatcher is not None:
await dispatcher.cancel(self)
await get_traffic_controller().send_queue.cancel_account(self.account_id)
await self._reply_queue.stop()
if self._ws_client:
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()