562 lines
23 KiB
Python
562 lines
23 KiB
Python
"""渠道出站消息分发器（仅运行在 API 容器）。"""
|
||
|
||
import asyncio
|
||
import json
|
||
import os
|
||
import random
|
||
import socket
|
||
import time
|
||
from datetime import UTC, datetime, timedelta
|
||
|
||
from redis.exceptions import ResponseError
|
||
from sqlalchemy import select
|
||
|
||
from yuxi.channel.config import ChannelConfigManager
|
||
from yuxi.channel.constants import (
|
||
CHANNEL_BACKOFF_BASE_SECONDS,
|
||
CHANNEL_BACKOFF_MAX_SECONDS,
|
||
CHANNEL_CONSUMER_GROUP,
|
||
CHANNEL_DELIVERED_TTL_SECONDS,
|
||
CHANNEL_MAX_DELIVERY_ATTEMPTS,
|
||
CHANNEL_PENDING_CLAIM_MIN_IDLE_MS,
|
||
CHANNEL_PROCESSING_LOCK_RENEW_INTERVAL_SECONDS,
|
||
CHANNEL_PROCESSING_LOCK_TTL_SECONDS,
|
||
CHANNEL_STREAM_KEY,
|
||
DeliveryStatus,
|
||
DispatchResult,
|
||
channel_delivered_key,
|
||
channel_processing_key,
|
||
)
|
||
from yuxi.channel.exceptions import (
|
||
ChannelPermanentError,
|
||
ChannelRateLimitedError,
|
||
ChannelRetryableError,
|
||
)
|
||
from yuxi.channel.metrics import (
|
||
channel_delivery_duration_seconds,
|
||
channel_messages_delivered_total,
|
||
channel_messages_failed_total,
|
||
)
|
||
from yuxi.channel.middlewares.protocols import OutboundContext, OutboundResult
|
||
from yuxi.channel.middlewares.registry import OutboundMiddlewareRegistry
|
||
from yuxi.channel.outbound.lock import RedisLockRenewer
|
||
from yuxi.channel.plugins.registry import ChannelRegistry, get_registry
|
||
from yuxi.channel.ports import OutboundPort
|
||
from yuxi.repositories.conversation_repository import ConversationRepository
|
||
from yuxi.repositories.message_repository import MessageRepository
|
||
from yuxi.services.run_queue_service import get_redis_client
|
||
from yuxi.storage.postgres.manager import pg_manager
|
||
from yuxi.storage.postgres.model_channel import ChannelSession
|
||
from yuxi.utils.logging_config import logger
|
||
|
||
|
||
class OutboundDispatcher:
|
||
"""消费 Redis Stream 中的渠道出站事件，调用插件完成投递。"""
|
||
|
||
MAX_ATTEMPTS = CHANNEL_MAX_DELIVERY_ATTEMPTS
|
||
PROCESSING_LOCK_TTL = CHANNEL_PROCESSING_LOCK_TTL_SECONDS
|
||
LOCK_RENEW_INTERVAL = CHANNEL_PROCESSING_LOCK_RENEW_INTERVAL_SECONDS
|
||
|
||
def __init__(
|
||
self,
|
||
registry: ChannelRegistry | None = None,
|
||
config_manager: ChannelConfigManager | None = None,
|
||
outbound_registry: OutboundMiddlewareRegistry | None = None,
|
||
) -> None:
|
||
self.registry = registry or get_registry()
|
||
self.config = config_manager or ChannelConfigManager()
|
||
self._outbound_registry = outbound_registry or OutboundMiddlewareRegistry()
|
||
self._consumer_name: str | None = None
|
||
self._task: asyncio.Task | None = None
|
||
|
||
@property
|
||
def CONSUMER_NAME(self) -> str:
|
||
"""惰性计算消费者名称，确保在运行时获取正确的 PID。"""
|
||
if self._consumer_name is None:
|
||
self._consumer_name = f"{socket.gethostname()}-{os.getpid()}"
|
||
return self._consumer_name
|
||
|
||
async def start(self) -> None:
|
||
redis = await get_redis_client()
|
||
try:
|
||
await redis.xgroup_create(
|
||
CHANNEL_STREAM_KEY,
|
||
CHANNEL_CONSUMER_GROUP,
|
||
id="0",
|
||
mkstream=True,
|
||
)
|
||
except ResponseError as e:
|
||
if "already exists" not in str(e).lower():
|
||
raise
|
||
await self._outbound_registry.start_all()
|
||
self._task = asyncio.create_task(self._consume_loop())
|
||
logger.info("OutboundDispatcher started (consumer=%s)", self.CONSUMER_NAME)
|
||
|
||
async def stop(self) -> None:
|
||
if self._task:
|
||
self._task.cancel()
|
||
try:
|
||
await self._task
|
||
except asyncio.CancelledError:
|
||
pass
|
||
self._task = None
|
||
await self._outbound_registry.stop_all()
|
||
logger.info("OutboundDispatcher stopped")
|
||
|
||
def invalidate(self, channel_type: str, account_id: str) -> None:
|
||
"""清除指定渠道账户的出站中间件缓存。"""
|
||
self._outbound_registry.invalidate(channel_type, account_id)
|
||
|
||
async def _consume_loop(self) -> None:
|
||
redis = await get_redis_client()
|
||
while True:
|
||
try:
|
||
# 1. 认领死消费者的 pending 消息
|
||
await self._claim_stale_pending(redis)
|
||
# 2. 读取本消费者尚未 ACK 的 pending 消息
|
||
pending = await redis.xreadgroup(
|
||
groupname=CHANNEL_CONSUMER_GROUP,
|
||
consumername=self.CONSUMER_NAME,
|
||
streams={CHANNEL_STREAM_KEY: "0"},
|
||
count=10,
|
||
block=5000,
|
||
)
|
||
for _stream, entries in pending:
|
||
for entry_id, fields in entries:
|
||
await self._handle_event(entry_id, fields)
|
||
# 3. 读取新消息
|
||
new = await redis.xreadgroup(
|
||
groupname=CHANNEL_CONSUMER_GROUP,
|
||
consumername=self.CONSUMER_NAME,
|
||
streams={CHANNEL_STREAM_KEY: ">"},
|
||
count=10,
|
||
block=5000,
|
||
)
|
||
for _stream, entries in new:
|
||
for entry_id, fields in entries:
|
||
await self._handle_event(entry_id, fields)
|
||
except asyncio.CancelledError:
|
||
break
|
||
except Exception:
|
||
logger.exception("OutboundDispatcher consume loop error")
|
||
await asyncio.sleep(1)
|
||
|
||
async def _claim_stale_pending(self, redis) -> None:
|
||
"""认领超过空闲阈值、原消费者可能已死亡的 pending 消息。"""
|
||
try:
|
||
cursor = "0-0"
|
||
while True:
|
||
result = await redis.xautoclaim(
|
||
CHANNEL_STREAM_KEY,
|
||
CHANNEL_CONSUMER_GROUP,
|
||
self.CONSUMER_NAME,
|
||
min_idle_time=CHANNEL_PENDING_CLAIM_MIN_IDLE_MS,
|
||
start_id=cursor,
|
||
count=10,
|
||
justid=False,
|
||
)
|
||
# redis-py 返回 [next_cursor, [(id, fields), ...], [deleted_ids...]]
|
||
next_cursor = result[0] if result else "0-0"
|
||
claimed = result[1] if len(result) > 1 else []
|
||
for entry_id, fields in claimed:
|
||
await self._handle_event(entry_id, fields)
|
||
if not next_cursor or next_cursor == "0-0":
|
||
break
|
||
cursor = next_cursor
|
||
except ResponseError as e:
|
||
logger.warning("XAUTOCLAIM failed: %s", e)
|
||
except Exception:
|
||
logger.exception("Failed to claim stale pending messages")
|
||
|
||
async def _handle_event(self, entry_id: str, fields: dict) -> None:
|
||
payload_raw = fields.get("payload", "{}")
|
||
try:
|
||
event = json.loads(payload_raw)
|
||
except json.JSONDecodeError:
|
||
logger.warning("Skip malformed channel outbound event: %s", payload_raw)
|
||
# 格式错误无法重试，直接 ACK 避免永久堆积
|
||
await self._ack(entry_id)
|
||
return
|
||
|
||
message_id = event.get("message_id")
|
||
if not message_id:
|
||
logger.warning("Skip channel outbound event without message_id")
|
||
await self._ack(entry_id)
|
||
return
|
||
|
||
redis = await get_redis_client()
|
||
delivered_key = channel_delivered_key(message_id)
|
||
processing_key = channel_processing_key(message_id)
|
||
|
||
# 已投递则跳过
|
||
if await redis.exists(delivered_key):
|
||
await self._ack(entry_id)
|
||
return
|
||
|
||
# 检查是否未到重试时间
|
||
next_retry_at = await self._get_next_retry_at(message_id)
|
||
if next_retry_at and datetime.now(UTC) < next_retry_at:
|
||
logger.debug(
|
||
"Message %s retry deferred until %s",
|
||
message_id,
|
||
next_retry_at.isoformat(),
|
||
)
|
||
return
|
||
|
||
# 获取处理锁，owner 信息用于崩溃后识别与排查
|
||
lock_value = json.dumps(
|
||
{
|
||
"owner": self.CONSUMER_NAME,
|
||
"started_at": datetime.now(UTC).isoformat(),
|
||
"entry_id": entry_id,
|
||
},
|
||
ensure_ascii=False,
|
||
)
|
||
lock_acquired = await redis.set(processing_key, lock_value, nx=True, ex=self.PROCESSING_LOCK_TTL)
|
||
if not lock_acquired:
|
||
lock_acquired = await self._try_reclaim_lock(redis, processing_key, lock_value)
|
||
if not lock_acquired:
|
||
logger.debug(
|
||
"Message %s is being processed by another consumer, skip",
|
||
message_id,
|
||
)
|
||
return
|
||
|
||
async with RedisLockRenewer(
|
||
redis, processing_key, ttl=self.PROCESSING_LOCK_TTL, interval=self.LOCK_RENEW_INTERVAL
|
||
):
|
||
await self._process_event(entry_id, event, message_id, redis, delivered_key, processing_key)
|
||
|
||
async def _try_reclaim_lock(self, redis, lock_key: str, lock_value: str) -> bool:
|
||
"""若当前锁属于本消费者（崩溃重启场景），则重新占有。"""
|
||
try:
|
||
current = await redis.get(lock_key)
|
||
if not current:
|
||
return bool(await redis.set(lock_key, lock_value, nx=True, ex=self.PROCESSING_LOCK_TTL))
|
||
data = json.loads(current)
|
||
if data.get("owner") == self.CONSUMER_NAME:
|
||
await redis.set(lock_key, lock_value, xx=True, ex=self.PROCESSING_LOCK_TTL)
|
||
return True
|
||
except Exception:
|
||
logger.warning("Failed to reclaim processing lock %s", lock_key)
|
||
return False
|
||
|
||
async def _process_event(
|
||
self,
|
||
entry_id: str,
|
||
event: dict,
|
||
message_id: int,
|
||
redis,
|
||
delivered_key: str,
|
||
processing_key: str,
|
||
) -> None:
|
||
dispatch_status = DispatchResult.RETRYABLE
|
||
retry_after: int | None = None
|
||
try:
|
||
async with pg_manager.get_async_session_context() as session:
|
||
msg_repo = MessageRepository(session)
|
||
message = await msg_repo.get(message_id)
|
||
if message is None:
|
||
logger.warning("Message %s not found, acking event", message_id)
|
||
await self._ack(entry_id)
|
||
await redis.delete(processing_key)
|
||
return
|
||
|
||
# 增加尝试次数
|
||
attempts = (message.channel_metadata or {}).get("attempts", 0) + 1
|
||
message.update_channel_metadata({"attempts": attempts})
|
||
await session.commit()
|
||
|
||
if attempts > self.MAX_ATTEMPTS:
|
||
logger.warning(
|
||
"Message %s exceeded max delivery attempts (%s), marking failed",
|
||
message_id,
|
||
self.MAX_ATTEMPTS,
|
||
)
|
||
await self._finalize_failure(entry_id, message_id, redis, delivered_key, processing_key)
|
||
return
|
||
|
||
dispatch_status, retry_after = await self._do_dispatch(event, message, session)
|
||
except ChannelPermanentError:
|
||
logger.exception("Permanent error dispatching channel message %s", message_id)
|
||
dispatch_status = DispatchResult.PERMANENT_FAILURE
|
||
except ChannelRateLimitedError as e:
|
||
logger.exception("Rate limited dispatching channel message %s", message_id)
|
||
dispatch_status = DispatchResult.RETRYABLE
|
||
retry_after = e.retry_after
|
||
except ChannelRetryableError:
|
||
logger.exception("Retryable error dispatching channel message %s", message_id)
|
||
dispatch_status = DispatchResult.RETRYABLE
|
||
except Exception:
|
||
logger.exception("Failed to dispatch channel message %s", message_id)
|
||
dispatch_status = DispatchResult.RETRYABLE
|
||
|
||
if dispatch_status == DispatchResult.SUCCESS:
|
||
await redis.setex(delivered_key, CHANNEL_DELIVERED_TTL_SECONDS, "1")
|
||
await self._ack(entry_id)
|
||
await redis.delete(processing_key)
|
||
elif dispatch_status == DispatchResult.PERMANENT_FAILURE:
|
||
await self._finalize_failure(
|
||
entry_id,
|
||
message_id,
|
||
redis,
|
||
delivered_key,
|
||
processing_key,
|
||
)
|
||
else:
|
||
# 可重试失败：保留 pending，更新 next_retry_at，释放 processing 锁
|
||
await self._defer_retry(message_id, retry_after_seconds=retry_after)
|
||
await redis.delete(processing_key)
|
||
|
||
async def _finalize_failure(
|
||
self,
|
||
entry_id: str,
|
||
message_id: int,
|
||
redis,
|
||
delivered_key: str,
|
||
processing_key: str,
|
||
) -> None:
|
||
"""永久失败或超限：ACK + 设置 delivered + 删锁。
|
||
|
||
对 partial_failed 设置 delivered 是为了避免已发出 chunk 被重复投递；
|
||
对其他永久失败设置 delivered 是为了避免无意义重试。
|
||
"""
|
||
try:
|
||
async with pg_manager.get_async_session_context() as session:
|
||
msg_repo = MessageRepository(session)
|
||
message = await msg_repo.get(message_id)
|
||
if message is not None:
|
||
chunks = (message.channel_metadata or {}).get("chunks") or []
|
||
has_success = any(isinstance(cs, dict) and cs.get("status") == "success" for cs in chunks)
|
||
has_failure = any(isinstance(cs, dict) and cs.get("status") == "failed" for cs in chunks)
|
||
if has_success and has_failure:
|
||
message.delivery_status = DeliveryStatus.PARTIAL_FAILED
|
||
else:
|
||
message.delivery_status = DeliveryStatus.FAILED
|
||
await session.commit()
|
||
except Exception:
|
||
logger.exception("Failed to update delivery_status for message %s", message_id)
|
||
|
||
await redis.setex(delivered_key, CHANNEL_DELIVERED_TTL_SECONDS, "1")
|
||
await self._ack(entry_id)
|
||
await redis.delete(processing_key)
|
||
|
||
async def _defer_retry(self, message_id: int, retry_after_seconds: int | None = None) -> None:
|
||
"""可重试失败时，在 channel_metadata 写入下次重试时间。"""
|
||
try:
|
||
async with pg_manager.get_async_session_context() as session:
|
||
msg_repo = MessageRepository(session)
|
||
message = await msg_repo.get(message_id)
|
||
if message is None:
|
||
return
|
||
attempts = (message.channel_metadata or {}).get("attempts", 1)
|
||
if retry_after_seconds is not None:
|
||
delay = retry_after_seconds
|
||
else:
|
||
delay = min(
|
||
CHANNEL_BACKOFF_BASE_SECONDS * (2**attempts),
|
||
CHANNEL_BACKOFF_MAX_SECONDS,
|
||
)
|
||
delay = delay + random.randint(0, delay // 2)
|
||
next_retry_at = datetime.now(UTC) + timedelta(seconds=delay)
|
||
message.update_channel_metadata({"next_retry_at": next_retry_at.isoformat()})
|
||
await session.commit()
|
||
logger.info(
|
||
"Channel message %s deferred to %s (attempt=%s, delay=%ss)",
|
||
message_id,
|
||
next_retry_at.isoformat(),
|
||
attempts,
|
||
delay,
|
||
)
|
||
except Exception:
|
||
logger.exception("Failed to record retry deferral for message %s", message_id)
|
||
|
||
async def _get_next_retry_at(self, message_id: int) -> datetime | None:
|
||
try:
|
||
async with pg_manager.get_async_session_context() as session:
|
||
msg_repo = MessageRepository(session)
|
||
message = await msg_repo.get(message_id)
|
||
if message is None:
|
||
return None
|
||
channel_metadata = message.channel_metadata or {}
|
||
raw = channel_metadata.get("next_retry_at")
|
||
if not raw:
|
||
return None
|
||
return datetime.fromisoformat(raw)
|
||
except Exception:
|
||
logger.warning("Failed to parse next_retry_at for message %s", message_id)
|
||
return None
|
||
|
||
async def _ack(self, entry_id: str) -> None:
|
||
try:
|
||
redis = await get_redis_client()
|
||
await redis.xack(CHANNEL_STREAM_KEY, CHANNEL_CONSUMER_GROUP, entry_id)
|
||
except Exception:
|
||
logger.warning("Failed to ack stream entry %s", entry_id)
|
||
|
||
async def _do_dispatch(
|
||
self,
|
||
event: dict,
|
||
message,
|
||
session,
|
||
) -> tuple[str, int | None]:
|
||
"""执行实际投递，返回 (success/permanent_failure/retryable, retry_after_seconds)。"""
|
||
labels = {
|
||
"channel_type": event["channel_type"],
|
||
"account_id": event["account_id"],
|
||
}
|
||
start_time = time.perf_counter()
|
||
|
||
conv_repo = ConversationRepository(session)
|
||
conversation = await conv_repo.get_conversation_by_id(event["conversation_id"])
|
||
if not conversation or not message:
|
||
logger.warning(
|
||
"Conversation or message not found for message %s (conversation_id=%s)",
|
||
message.id,
|
||
event.get("conversation_id"),
|
||
)
|
||
return DispatchResult.PERMANENT_FAILURE, None
|
||
|
||
result = await session.execute(select(ChannelSession).where(ChannelSession.session_key == event["session_key"]))
|
||
channel_session = result.scalar_one_or_none()
|
||
if channel_session is None:
|
||
logger.warning("ChannelSession not found for session_key=%s", event["session_key"])
|
||
|
||
plugin = self.registry.get_plugin(event["channel_type"])
|
||
if plugin is None:
|
||
logger.warning("No plugin registered for channel_type=%s", event["channel_type"])
|
||
await self._update_message_status(
|
||
session,
|
||
message,
|
||
status=DeliveryStatus.FAILED,
|
||
sent_ids=[],
|
||
chunk_statuses=[{"index": 0, "status": "failed", "error": "plugin not found"}],
|
||
)
|
||
return DispatchResult.PERMANENT_FAILURE, None
|
||
|
||
if not isinstance(plugin, OutboundPort):
|
||
logger.warning(
|
||
"Plugin for %s does not implement OutboundPort",
|
||
event["channel_type"],
|
||
)
|
||
await self._update_message_status(
|
||
session,
|
||
message,
|
||
status=DeliveryStatus.FAILED,
|
||
sent_ids=[],
|
||
chunk_statuses=[{"index": 0, "status": "failed", "error": "plugin does not implement OutboundPort"}],
|
||
)
|
||
return DispatchResult.PERMANENT_FAILURE, None
|
||
|
||
try:
|
||
config = await self.config.get_config(event["channel_type"], event["account_id"])
|
||
except Exception as e:
|
||
logger.exception(
|
||
"Failed to load channel config for %s/%s",
|
||
event["channel_type"],
|
||
event["account_id"],
|
||
)
|
||
await self._update_message_status(
|
||
session,
|
||
message,
|
||
status=DeliveryStatus.FAILED,
|
||
sent_ids=[],
|
||
chunk_statuses=[{"index": 0, "status": "failed", "error": str(e)}],
|
||
)
|
||
duration = time.perf_counter() - start_time
|
||
channel_delivery_duration_seconds.observe(labels, duration)
|
||
channel_messages_failed_total.inc(labels)
|
||
logger.bind(
|
||
event="channel_delivery_failed",
|
||
channel_type=event["channel_type"],
|
||
account_id=event["account_id"],
|
||
session_key=event["session_key"],
|
||
message_id=message.id,
|
||
duration_ms=int(duration * 1000),
|
||
status="failed",
|
||
error=str(e),
|
||
).error("Channel outbound delivery failed: config load error")
|
||
return DispatchResult.PERMANENT_FAILURE, None
|
||
|
||
capabilities = plugin.get_delivery_capabilities()
|
||
ctx = OutboundContext(
|
||
event=event,
|
||
message=None,
|
||
db_message=message,
|
||
config=config,
|
||
config_mw={},
|
||
plugin=plugin,
|
||
channel_session=channel_session,
|
||
conversation=conversation,
|
||
db=session,
|
||
capabilities=capabilities,
|
||
labels=labels,
|
||
start_time=start_time,
|
||
)
|
||
ctx.update_status = self._finalize_dispatch
|
||
|
||
chain = self._outbound_registry.resolve_chain(config)
|
||
try:
|
||
result = await self._execute_chain(ctx, chain)
|
||
return result.status, result.retry_after
|
||
except Exception:
|
||
logger.exception("Outbound middleware chain failed for message %s", message.id)
|
||
return DispatchResult.RETRYABLE, None
|
||
|
||
async def _execute_chain(
|
||
self,
|
||
ctx: OutboundContext,
|
||
chain: list,
|
||
) -> OutboundResult:
|
||
"""按顺序执行出站中间件链。"""
|
||
|
||
async def _next(idx: int) -> OutboundResult:
|
||
if idx >= len(chain):
|
||
return OutboundResult(status=ctx.dispatch_result, retry_after=ctx.retry_after)
|
||
mw = chain[idx]
|
||
ctx.config_mw = self._outbound_registry.get_middleware_config(ctx.config, mw.name)
|
||
return await mw.process(ctx, lambda: _next(idx + 1))
|
||
|
||
return await _next(0)
|
||
|
||
async def _finalize_dispatch(self, ctx: OutboundContext, status: str) -> None:
|
||
"""由 StatusUpdateMiddleware 回调，用于提交状态、指标与日志。"""
|
||
await self._update_message_status(
|
||
ctx.db,
|
||
ctx.db_message,
|
||
status,
|
||
ctx.sent_ids,
|
||
ctx.chunk_statuses,
|
||
)
|
||
|
||
duration = time.perf_counter() - ctx.start_time
|
||
channel_delivery_duration_seconds.observe(ctx.labels, duration)
|
||
if status == DeliveryStatus.COMPLETE:
|
||
channel_messages_delivered_total.inc(ctx.labels)
|
||
else:
|
||
channel_messages_failed_total.inc(ctx.labels)
|
||
|
||
log_event = "channel_message_delivered" if status == DeliveryStatus.COMPLETE else "channel_delivery_failed"
|
||
logger.bind(
|
||
event=log_event,
|
||
channel_type=ctx.event.get("channel_type"),
|
||
account_id=ctx.event.get("account_id"),
|
||
session_key=ctx.event.get("session_key"),
|
||
message_id=ctx.db_message.id,
|
||
duration_ms=int(duration * 1000),
|
||
status=status,
|
||
).info("Channel outbound message %s", status)
|
||
|
||
async def _update_message_status(
|
||
self,
|
||
session,
|
||
message,
|
||
status: DeliveryStatus,
|
||
sent_ids: list[str | None],
|
||
chunk_statuses: list[dict],
|
||
) -> None:
|
||
message.delivery_status = status
|
||
message.channel_message_id = next((sid for sid in sent_ids if sid), None)
|
||
message.update_channel_metadata({"sent_ids": sent_ids, "chunks": chunk_statuses})
|
||
await session.commit()
|