233 lines
7.5 KiB
Python
233 lines
7.5 KiB
Python
"""
|
||
进程内 inbox / outbox 队列（骨架过渡实现）。
|
||
|
||
本文件职责：回调快写入队、按条认领、message_id 去重；后续可替换为 PostgreSQL 表。
|
||
禁止：Redis 全局 Worker 单例锁；禁止在回调线程里等 LLM。
|
||
为何先用内存：先打通「收消息→回复」联调；落 PG 迁移不挡本迭代验收。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import logging
|
||
import threading
|
||
import time
|
||
import uuid
|
||
from collections import OrderedDict
|
||
from dataclasses import dataclass, field
|
||
from enum import Enum
|
||
from typing import Any, Optional
|
||
|
||
from agent.channel.wecom.models import InboundMessage
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class TaskState(str, Enum):
|
||
PENDING = "pending"
|
||
CLAIMED = "claimed"
|
||
DONE = "done"
|
||
FAILED = "failed"
|
||
|
||
|
||
@dataclass
|
||
class InboxItem:
|
||
"""inbox 一条入站任务。"""
|
||
|
||
id: str
|
||
message: InboundMessage
|
||
state: TaskState = TaskState.PENDING
|
||
claimed_at: float = 0.0
|
||
attempts: int = 0
|
||
last_error: str = ""
|
||
|
||
|
||
@dataclass
|
||
class OutboxItem:
|
||
"""
|
||
outbox 一条出站任务。
|
||
|
||
touser：必须是入站 sender_id，禁止另换 ID。
|
||
发送成功 ≠ 业务成功；此处仅通道送达。
|
||
"""
|
||
|
||
id: str
|
||
touser: str
|
||
content: str
|
||
state: TaskState = TaskState.PENDING
|
||
claimed_at: float = 0.0
|
||
attempts: int = 0
|
||
last_error: str = ""
|
||
created_at: float = field(default_factory=time.time)
|
||
dedupe_key: str = ""
|
||
|
||
|
||
class MemoryMessageStore:
|
||
"""
|
||
线程安全的内存 inbox/outbox。
|
||
|
||
认领：同一条任务同一时刻只被一个消费者持有；超时可回收（默认 120s）。
|
||
去重：相同 message_id / dedupe_key 不重复入队。
|
||
"""
|
||
|
||
def __init__(self, *, claim_timeout_sec: float = 120.0, dedupe_maxlen: int = 5000) -> None:
|
||
self._lock = threading.RLock()
|
||
self._inbox: OrderedDict[str, InboxItem] = OrderedDict()
|
||
self._outbox: OrderedDict[str, OutboxItem] = OrderedDict()
|
||
self._seen_msg: OrderedDict[str, float] = OrderedDict()
|
||
self._seen_out: OrderedDict[str, float] = OrderedDict()
|
||
self._claim_timeout = claim_timeout_sec
|
||
self._dedupe_maxlen = dedupe_maxlen
|
||
self._inbox_event = threading.Event()
|
||
self._outbox_event = threading.Event()
|
||
|
||
def enqueue_inbound(self, message: InboundMessage) -> tuple[bool, str]:
|
||
"""
|
||
写入 inbox。
|
||
|
||
返回：(是否新入队, 说明)。无 sender_id 拒绝；message_id 重复则跳过。
|
||
副作用：唤醒 inbox 等待线程。
|
||
"""
|
||
if not message.has_sender():
|
||
return False, "missing_sender_id"
|
||
with self._lock:
|
||
if message.message_id in self._seen_msg:
|
||
return False, "duplicate_message_id"
|
||
self._remember(self._seen_msg, message.message_id)
|
||
item = InboxItem(id=str(uuid.uuid4()), message=message)
|
||
self._inbox[item.id] = item
|
||
self._inbox_event.set()
|
||
logger.info(
|
||
"inbox 入队 id=%s msg_id=%s sender=%s type=%s",
|
||
item.id,
|
||
message.message_id,
|
||
message.sender_id,
|
||
message.msg_type,
|
||
)
|
||
return True, item.id
|
||
|
||
def claim_inbound(self) -> Optional[InboxItem]:
|
||
"""认领一条 pending/超时 inbox；无任务返回 None。"""
|
||
now = time.time()
|
||
with self._lock:
|
||
for item in self._inbox.values():
|
||
if item.state == TaskState.PENDING or (
|
||
item.state == TaskState.CLAIMED
|
||
and item.claimed_at
|
||
and now - item.claimed_at > self._claim_timeout
|
||
):
|
||
item.state = TaskState.CLAIMED
|
||
item.claimed_at = now
|
||
item.attempts += 1
|
||
return item
|
||
self._inbox_event.clear()
|
||
return None
|
||
|
||
def complete_inbound(self, item_id: str, *, error: str = "") -> None:
|
||
"""标记 inbox 完成或失败。"""
|
||
with self._lock:
|
||
item = self._inbox.get(item_id)
|
||
if not item:
|
||
return
|
||
if error:
|
||
item.state = TaskState.FAILED
|
||
item.last_error = error
|
||
else:
|
||
item.state = TaskState.DONE
|
||
|
||
def enqueue_outbound(
|
||
self,
|
||
*,
|
||
touser: str,
|
||
content: str,
|
||
dedupe_key: str = "",
|
||
) -> tuple[bool, str]:
|
||
"""写入 outbox；dedupe_key 重复则跳过。"""
|
||
if not touser or not content:
|
||
return False, "invalid_outbox"
|
||
with self._lock:
|
||
if dedupe_key and dedupe_key in self._seen_out:
|
||
return False, "duplicate_outbox"
|
||
if dedupe_key:
|
||
self._remember(self._seen_out, dedupe_key)
|
||
item = OutboxItem(
|
||
id=str(uuid.uuid4()),
|
||
touser=touser,
|
||
content=content,
|
||
dedupe_key=dedupe_key,
|
||
)
|
||
self._outbox[item.id] = item
|
||
self._outbox_event.set()
|
||
logger.info("outbox 入队 id=%s touser=%s", item.id, touser)
|
||
return True, item.id
|
||
|
||
def claim_outbound(self) -> Optional[OutboxItem]:
|
||
"""认领一条 pending/超时 outbox。"""
|
||
now = time.time()
|
||
with self._lock:
|
||
for item in self._outbox.values():
|
||
if item.state == TaskState.PENDING or (
|
||
item.state == TaskState.CLAIMED
|
||
and item.claimed_at
|
||
and now - item.claimed_at > self._claim_timeout
|
||
):
|
||
item.state = TaskState.CLAIMED
|
||
item.claimed_at = now
|
||
item.attempts += 1
|
||
return item
|
||
self._outbox_event.clear()
|
||
return None
|
||
|
||
def complete_outbound(self, item_id: str, *, error: str = "", permanent: bool = False) -> None:
|
||
"""
|
||
标记 outbox 完成或失败。
|
||
|
||
permanent=True：不再重试（如 IP 未加白 60020）。
|
||
"""
|
||
with self._lock:
|
||
item = self._outbox.get(item_id)
|
||
if not item:
|
||
return
|
||
if error:
|
||
if permanent or item.attempts >= 5:
|
||
item.state = TaskState.FAILED
|
||
else:
|
||
item.state = TaskState.PENDING
|
||
item.claimed_at = 0.0
|
||
self._outbox_event.set()
|
||
item.last_error = error
|
||
else:
|
||
item.state = TaskState.DONE
|
||
|
||
def wait_inbox(self, timeout: float = 1.0) -> None:
|
||
self._inbox_event.wait(timeout=timeout)
|
||
|
||
def wait_outbox(self, timeout: float = 1.0) -> None:
|
||
self._outbox_event.wait(timeout=timeout)
|
||
|
||
def stats(self) -> dict[str, Any]:
|
||
with self._lock:
|
||
return {
|
||
"inbox": len(self._inbox),
|
||
"outbox": len(self._outbox),
|
||
"seen_msg": len(self._seen_msg),
|
||
}
|
||
|
||
def _remember(self, store: OrderedDict[str, float], key: str) -> None:
|
||
store[key] = time.time()
|
||
while len(store) > self._dedupe_maxlen:
|
||
store.popitem(last=False)
|
||
|
||
|
||
# 进程内单例：HTTP 回调与消费线程共享（同进程）。多进程部署时必须换 PG。
|
||
_STORE: Optional[MemoryMessageStore] = None
|
||
_STORE_LOCK = threading.Lock()
|
||
|
||
|
||
def get_message_store() -> MemoryMessageStore:
|
||
"""获取进程内消息账本单例。"""
|
||
global _STORE
|
||
with _STORE_LOCK:
|
||
if _STORE is None:
|
||
_STORE = MemoryMessageStore()
|
||
return _STORE
|