qyyfw
78ff261b5a
- 前端 React+TypeScript (源码+构建产物) - 后端 Python FastAPI - 数据库迁移 001-015 - Docker部署配置 - 微信ClawBot API频道绑定 - 项目文档 README.md
293 lines
9.4 KiB
Python
293 lines
9.4 KiB
Python
"""
|
|
Redis事件总线 - 利用已有Redis实现消息队列
|
|
参照AstrBot EventBus设计，但使用Redis替代asyncio.Queue
|
|
"""
|
|
import json
|
|
import logging
|
|
import threading
|
|
import time
|
|
from typing import Callable
|
|
import redis
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Redis配置（与项目已有配置一致）
|
|
try:
|
|
from .config import REDIS_HOST, REDIS_PORT, REDIS_PASSWORD
|
|
except ImportError:
|
|
REDIS_HOST = "localhost"
|
|
REDIS_PORT = 6379
|
|
REDIS_PASSWORD = "pwd2026"
|
|
|
|
|
|
def _get_redis():
|
|
"""获取Redis连接"""
|
|
return redis.Redis(
|
|
host=REDIS_HOST,
|
|
port=REDIS_PORT,
|
|
password=REDIS_PASSWORD,
|
|
decode_responses=True,
|
|
)
|
|
|
|
|
|
# Redis键名前缀
|
|
QUEUE_KEY = "aiflow:event_queue"
|
|
PROCESSING_KEY = "aiflow:event_processing"
|
|
|
|
|
|
class MessageEvent:
|
|
"""统一消息事件"""
|
|
|
|
def __init__(self, platform: str, channel_id: str, content: str,
|
|
sender_id: str = "", sender_name: str = "",
|
|
message_id: str = "", group_id: str = None,
|
|
content_type: str = "text", raw: dict = None):
|
|
self.platform = platform
|
|
self.channel_id = channel_id
|
|
self.content = content
|
|
self.sender_id = sender_id
|
|
self.sender_name = sender_name
|
|
self.message_id = message_id
|
|
self.group_id = group_id
|
|
self.content_type = content_type
|
|
self.raw = raw or {}
|
|
self.timestamp = int(time.time())
|
|
|
|
def to_dict(self) -> dict:
|
|
return {
|
|
"platform": self.platform,
|
|
"channel_id": self.channel_id,
|
|
"content": self.content,
|
|
"sender_id": self.sender_id,
|
|
"sender_name": self.sender_name,
|
|
"message_id": self.message_id,
|
|
"group_id": self.group_id,
|
|
"content_type": self.content_type,
|
|
"raw": self.raw,
|
|
"timestamp": self.timestamp,
|
|
}
|
|
|
|
@classmethod
|
|
def from_dict(cls, data: dict) -> "MessageEvent":
|
|
return cls(
|
|
platform=data.get("platform", ""),
|
|
channel_id=data.get("channel_id", ""),
|
|
content=data.get("content", ""),
|
|
sender_id=data.get("sender_id", ""),
|
|
sender_name=data.get("sender_name", ""),
|
|
message_id=data.get("message_id", ""),
|
|
group_id=data.get("group_id"),
|
|
content_type=data.get("content_type", "text"),
|
|
raw=data.get("raw", {}),
|
|
)
|
|
|
|
|
|
class RedisEventBus:
|
|
"""Redis事件总线 - 生产者-消费者模式"""
|
|
|
|
def __init__(self):
|
|
self._handlers: dict[str, list[Callable]] = {}
|
|
self._consumers: list[threading.Thread] = []
|
|
self._running = False
|
|
|
|
def publish(self, event: MessageEvent) -> bool:
|
|
"""发布事件到Redis队列"""
|
|
try:
|
|
r = _get_redis()
|
|
event_data = json.dumps(event.to_dict(), ensure_ascii=False)
|
|
r.lpush(QUEUE_KEY, event_data)
|
|
logger.debug(f"[EventBus] 发布事件: {event.platform}:{event.channel_id}")
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"[EventBus] 发布事件失败: {e}")
|
|
return False
|
|
|
|
def subscribe(self, platform: str, handler: Callable[[MessageEvent], None]):
|
|
"""订阅特定平台的事件"""
|
|
if platform not in self._handlers:
|
|
self._handlers[platform] = []
|
|
self._handlers[platform].append(handler)
|
|
logger.info(f"[EventBus] 注册处理器: {platform}")
|
|
|
|
def _consume_loop(self):
|
|
"""消费者循环：从Redis队列取出事件并分发给处理器"""
|
|
r = _get_redis()
|
|
while self._running:
|
|
try:
|
|
# BRPOP 阻塞等待新事件
|
|
result = r.brpop(QUEUE_KEY, timeout=2)
|
|
if not result:
|
|
continue
|
|
|
|
_, event_data = result
|
|
event_dict = json.loads(event_data)
|
|
event = MessageEvent.from_dict(event_dict)
|
|
|
|
# 分发给对应平台的处理器
|
|
platform = event.platform
|
|
handlers = self._handlers.get(platform, []) + self._handlers.get("*", [])
|
|
|
|
for handler in handlers:
|
|
try:
|
|
handler(event)
|
|
except Exception as e:
|
|
logger.error(f"[EventBus] 处理器异常(platform={platform}): {e}")
|
|
|
|
except Exception as e:
|
|
if self._running:
|
|
logger.error(f"[EventBus] 消费循环异常: {e}")
|
|
time.sleep(1)
|
|
|
|
def start(self, num_consumers: int = 2):
|
|
"""启动消费者线程"""
|
|
if self._running:
|
|
return
|
|
self._running = True
|
|
for i in range(num_consumers):
|
|
t = threading.Thread(target=self._consume_loop, name=f"EventBus-Consumer-{i}", daemon=True)
|
|
t.start()
|
|
self._consumers.append(t)
|
|
logger.info(f"[EventBus] 已启动 {num_consumers} 个消费者")
|
|
|
|
def stop(self):
|
|
"""停止所有消费者"""
|
|
self._running = False
|
|
for t in self._consumers:
|
|
t.join(timeout=5)
|
|
self._consumers.clear()
|
|
logger.info("[EventBus] 已停止所有消费者")
|
|
|
|
def queue_size(self) -> int:
|
|
"""获取队列长度"""
|
|
try:
|
|
r = _get_redis()
|
|
return r.llen(QUEUE_KEY)
|
|
except Exception:
|
|
return -1
|
|
|
|
|
|
# 全局单例
|
|
_event_bus_instance = None
|
|
|
|
|
|
def get_event_bus() -> RedisEventBus:
|
|
"""获取全局事件总线单例"""
|
|
global _event_bus_instance
|
|
if _event_bus_instance is None:
|
|
_event_bus_instance = RedisEventBus()
|
|
return _event_bus_instance
|
|
|
|
|
|
# ============================================================
|
|
# 单元测试
|
|
# ============================================================
|
|
if __name__ == "__main__":
|
|
print("=" * 60)
|
|
print("Redis事件总线 单元测试")
|
|
print("=" * 60)
|
|
|
|
# 测试MessageEvent
|
|
print("\n[测试1] MessageEvent 创建和序列化")
|
|
event = MessageEvent(
|
|
platform="feishu",
|
|
channel_id="ch_001",
|
|
content="你好，帮我查询天气",
|
|
sender_id="user_001",
|
|
sender_name="张三",
|
|
message_id="msg_001",
|
|
group_id="grp_001",
|
|
content_type="text",
|
|
raw={"custom": "data"},
|
|
)
|
|
event_dict = event.to_dict()
|
|
assert event_dict["platform"] == "feishu"
|
|
assert event_dict["channel_id"] == "ch_001"
|
|
assert event_dict["content"] == "你好，帮我查询天气"
|
|
assert event_dict["sender_id"] == "user_001"
|
|
assert event_dict["sender_name"] == "张三"
|
|
assert event_dict["message_id"] == "msg_001"
|
|
assert event_dict["group_id"] == "grp_001"
|
|
assert event_dict["content_type"] == "text"
|
|
assert event_dict["raw"] == {"custom": "data"}
|
|
assert "timestamp" in event_dict
|
|
print(" PASS: MessageEvent序列化正确")
|
|
|
|
# 测试MessageEvent反序列化
|
|
print("\n[测试2] MessageEvent 反序列化")
|
|
event2 = MessageEvent.from_dict(event_dict)
|
|
assert event2.platform == "feishu"
|
|
assert event2.channel_id == "ch_001"
|
|
assert event2.content == "你好，帮我查询天气"
|
|
assert event2.sender_id == "user_001"
|
|
assert event2.sender_name == "张三"
|
|
assert event2.message_id == "msg_001"
|
|
assert event2.group_id == "grp_001"
|
|
assert event2.content_type == "text"
|
|
assert event2.raw == {"custom": "data"}
|
|
print(" PASS: MessageEvent反序列化正确")
|
|
|
|
# 测试RedisEventBus发布事件
|
|
print("\n[测试3] RedisEventBus 发布事件")
|
|
bus = RedisEventBus()
|
|
result = bus.publish(event)
|
|
assert result is True, "事件发布应该成功"
|
|
print(" PASS: 事件发布成功")
|
|
|
|
# 测试队列长度
|
|
print("\n[测试4] 队列长度查询")
|
|
size = bus.queue_size()
|
|
assert size >= 1, f"队列中应该至少有1条消息，实际: {size}"
|
|
print(f" PASS: 队列长度={size}")
|
|
|
|
# 测试事件消费
|
|
print("\n[测试5] 事件消费")
|
|
received_events = []
|
|
|
|
def handler(evt: MessageEvent):
|
|
received_events.append(evt)
|
|
|
|
bus.subscribe("feishu", handler)
|
|
bus.start(num_consumers=1)
|
|
|
|
# 发布一个事件并等待消费
|
|
test_event = MessageEvent(
|
|
platform="feishu",
|
|
channel_id="ch_test",
|
|
content="测试消息",
|
|
sender_id="user_test",
|
|
)
|
|
bus.publish(test_event)
|
|
|
|
# 等待消费者处理
|
|
time.sleep(1)
|
|
|
|
assert len(received_events) >= 1, f"应该接收到事件，实际: {len(received_events)}"
|
|
if received_events:
|
|
assert received_events[-1].platform == "feishu"
|
|
assert received_events[-1].channel_id == "ch_test"
|
|
assert received_events[-1].content == "测试消息"
|
|
print(f" PASS: 事件被正确消费, platform={received_events[-1].platform}, content={received_events[-1].content}")
|
|
|
|
# 清理
|
|
bus.stop()
|
|
print(" PASS: EventBus正常停止")
|
|
|
|
# 测试全局单例
|
|
print("\n[测试6] 全局单例")
|
|
bus1 = get_event_bus()
|
|
bus2 = get_event_bus()
|
|
assert bus1 is bus2, "get_event_bus()应该返回同一个实例"
|
|
print(" PASS: 全局单例正确")
|
|
|
|
# 清理Redis测试队列
|
|
print("\n[测试7] 清理测试数据")
|
|
r = _get_redis()
|
|
r.delete(QUEUE_KEY)
|
|
size = bus.queue_size()
|
|
assert size == 0, f"清理后队列应该为空，实际: {size}"
|
|
print(f" PASS: Redis测试队列已清理, 队列长度={size}")
|
|
|
|
print("\n" + "=" * 60)
|
|
print("所有测试通过!")
|
|
print("=" * 60)
|