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/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: