refactor: worker进程 → tokio task 架构 + 多渠道消息队列 + JSON工具系统
核心变更: - daemon: 废弃 worker 子进程模型,替换为 LLM/Tool/Send 三消费者 tokio task - queue: 新增公平轮转 MessageQueue + QueueRunner 路由,严格 ChannelId 渠道隔离 - channel: ChannelId(platform + user_id) 统一渠道标识,替代散列 user_id+platform - tools: 全部工具统一经由 execute_tool() → Tool Consumer 单一执行路径 - llm: chat_with_tools 返回 ChatResult,支持异步工具调用 → 保留/恢复 session - tools: 新增 JSON 工具定义系统 (ToolDef) 及 SubprocessRunner (HTTP 模板引擎) - bin: ias-auth 独立 JWT 生成器,ias-fetch headless Chrome 网页抓取 - 所有出站消息通过 Send Consumer 队列转发,无直接 send_text 调用 - Session TTL 清理 (5分钟),QueueRunner 通道容错 - cargo check 零警告,cargo test 23/23 通过
This commit is contained in:
@@ -0,0 +1,552 @@
|
||||
//! ## 多渠道消息队列 — MessageQueue
|
||||
//!
|
||||
//! # 设计目标
|
||||
//!
|
||||
//! 多用户场景下,保证公平性:
|
||||
//! - A 发了 3 条消息、B 发了 1 条 → 处理顺序 A1, B1, A2, A3
|
||||
//! - 不会因为 A 消息多而饿死 B
|
||||
//!
|
||||
//! 每个消息携带 user_id + correlation_id,保证回复能追溯到正确的用户。
|
||||
//!
|
||||
//! # 架构位置
|
||||
//!
|
||||
//! ```text
|
||||
//! receive_messages → enqueue PipelineMessage
|
||||
//! ┌─────────────────────────────┐
|
||||
//! │ MessageQueue │
|
||||
//! │ [User A] [User B] │ ← 每用户独立 VecDeque
|
||||
//! │ pending: [A, B, A] │ ← 公平轮转索引
|
||||
//! └─────────────────────────────┘
|
||||
//! QueueRunner::run() → dequeue → 按 kind 路由到 mpsc 通道
|
||||
//! ├─ llm_tx (UserMessage, ToolResult)
|
||||
//! ├─ tool_tx (ToolCall, ApprovalRequest)
|
||||
//! └─ send_tx (LLMReply)
|
||||
//! ```
|
||||
|
||||
use crate::channel::ChannelId;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::{HashMap, VecDeque};
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
use tokio::sync::Notify;
|
||||
use uuid::Uuid;
|
||||
|
||||
// ─── 管线消息类型 ───
|
||||
|
||||
/// 消息种类 —— 区分不同消费路由
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub enum MessageKind {
|
||||
/// 从外部收到的用户消息 → 交给 LLM Consumer
|
||||
UserMessage {
|
||||
text: String,
|
||||
message_id: String,
|
||||
context_token: Option<String>,
|
||||
/// 审批通过后携带的已批准工具名
|
||||
approved_tool: Option<String>,
|
||||
},
|
||||
/// LLM 生成的回复文本 → 发给 Send Consumer
|
||||
LLMReply {
|
||||
text: String,
|
||||
context_token: Option<String>,
|
||||
},
|
||||
/// LLM 请求的工具调用 → 发给 Tool Consumer
|
||||
ToolCall {
|
||||
session_id: Uuid,
|
||||
tool_call_id: String,
|
||||
tool_name: String,
|
||||
arguments: String,
|
||||
context_token: Option<String>,
|
||||
},
|
||||
/// 工具执行结果 → 回喂 LLM Consumer
|
||||
ToolResult {
|
||||
session_id: Uuid,
|
||||
tool_call_id: String,
|
||||
tool_name: String,
|
||||
result: String,
|
||||
context_token: Option<String>,
|
||||
},
|
||||
/// 高风险工具 → 需用户审批
|
||||
ApprovalRequest {
|
||||
session_id: Uuid,
|
||||
tool_name: String,
|
||||
tool_args: String,
|
||||
/// 审批后要追加到用户消息中的文本
|
||||
approval_prompt: String,
|
||||
},
|
||||
/// 定时任务到期 → 交给 LLM Consumer
|
||||
ScheduledTask {
|
||||
text: String,
|
||||
task_name: String,
|
||||
},
|
||||
}
|
||||
|
||||
/// 管线消息 —— 携带路由元数据的统一消息类型
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct PipelineMessage {
|
||||
/// 来源/目标渠道(平台 + 用户),全程不变
|
||||
pub channel: ChannelId,
|
||||
/// 关联 ID:追踪一次用户请求的完整链路(用户消息→LLM→工具→回喂→发送)
|
||||
pub correlation_id: Uuid,
|
||||
/// 消息种类
|
||||
pub kind: MessageKind,
|
||||
}
|
||||
|
||||
impl PipelineMessage {
|
||||
/// 便捷:渠道中的 user_id
|
||||
pub fn user_id(&self) -> &str { &self.channel.user_id }
|
||||
/// 便捷:渠道中的 platform
|
||||
#[allow(dead_code)]
|
||||
pub fn platform(&self) -> &str { &self.channel.platform }
|
||||
/// 快速构造一条用户消息
|
||||
#[allow(dead_code)]
|
||||
pub fn user_message(
|
||||
user_id: impl Into<String>,
|
||||
text: impl Into<String>,
|
||||
message_id: impl Into<String>,
|
||||
context_token: Option<String>,
|
||||
) -> Self {
|
||||
Self {
|
||||
channel: ChannelId::wechat(user_id),
|
||||
correlation_id: Uuid::new_v4(),
|
||||
kind: MessageKind::UserMessage {
|
||||
text: text.into(),
|
||||
message_id: message_id.into(),
|
||||
context_token,
|
||||
approved_tool: None,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
/// 快速构造一条 LLM 回复
|
||||
pub fn llm_reply(
|
||||
channel: ChannelId,
|
||||
correlation_id: Uuid,
|
||||
text: impl Into<String>,
|
||||
context_token: Option<String>,
|
||||
) -> Self {
|
||||
Self {
|
||||
channel,
|
||||
correlation_id,
|
||||
kind: MessageKind::LLMReply {
|
||||
text: text.into(),
|
||||
context_token,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
/// 快速构造一条工具调用
|
||||
pub fn tool_call(
|
||||
channel: ChannelId,
|
||||
correlation_id: Uuid,
|
||||
session_id: Uuid,
|
||||
tool_call_id: impl Into<String>,
|
||||
tool_name: impl Into<String>,
|
||||
arguments: impl Into<String>,
|
||||
context_token: Option<String>,
|
||||
) -> Self {
|
||||
Self {
|
||||
channel,
|
||||
correlation_id,
|
||||
kind: MessageKind::ToolCall {
|
||||
session_id,
|
||||
tool_call_id: tool_call_id.into(),
|
||||
tool_name: tool_name.into(),
|
||||
arguments: arguments.into(),
|
||||
context_token,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
/// 快速构造一条工具结果
|
||||
pub fn tool_result(
|
||||
channel: ChannelId,
|
||||
correlation_id: Uuid,
|
||||
session_id: Uuid,
|
||||
tool_call_id: impl Into<String>,
|
||||
tool_name: impl Into<String>,
|
||||
result: impl Into<String>,
|
||||
context_token: Option<String>,
|
||||
) -> Self {
|
||||
Self {
|
||||
channel,
|
||||
correlation_id,
|
||||
kind: MessageKind::ToolResult {
|
||||
session_id,
|
||||
tool_call_id: tool_call_id.into(),
|
||||
tool_name: tool_name.into(),
|
||||
result: result.into(),
|
||||
context_token,
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ─── 公平轮转队列 ───
|
||||
|
||||
/// 多渠道公平轮转消息队列
|
||||
///
|
||||
/// 按 ChannelId(platform + user_id)分队列,轮转出队保证公平。
|
||||
/// 出队后由 QueueRunner 按 MessageKind 路由到对应 mpsc 通道。
|
||||
pub struct MessageQueue {
|
||||
/// 每渠道独立队列
|
||||
queues: HashMap<ChannelId, VecDeque<PipelineMessage>>,
|
||||
/// 有待处理消息的渠道(轮转顺序)
|
||||
pending: VecDeque<ChannelId>,
|
||||
/// 空→非空时触发 Notify 唤醒消费者
|
||||
signal: Arc<Notify>,
|
||||
/// 定时唤醒间隔(兜底用)
|
||||
poll_interval: Duration,
|
||||
}
|
||||
|
||||
impl MessageQueue {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
queues: HashMap::new(),
|
||||
pending: VecDeque::new(),
|
||||
signal: Arc::new(Notify::new()),
|
||||
poll_interval: Duration::from_millis(200),
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
pub fn with_poll_interval(mut self, d: Duration) -> Self {
|
||||
self.poll_interval = d;
|
||||
self
|
||||
}
|
||||
|
||||
/// 入队。如果该渠道之前为空 → 加入 pending 尾部 + 通知
|
||||
pub fn enqueue(&mut self, item: PipelineMessage) {
|
||||
let ch = item.channel.clone();
|
||||
let q = self.queues.entry(ch.clone()).or_default();
|
||||
let was_empty = q.is_empty();
|
||||
q.push_back(item);
|
||||
if was_empty {
|
||||
self.pending.push_back(ch);
|
||||
self.signal.notify_one();
|
||||
}
|
||||
}
|
||||
|
||||
/// 公平轮转出队
|
||||
///
|
||||
/// 每次调用切换到下一个有消息的渠道:
|
||||
/// - 从 pending 头部取出一个渠道
|
||||
/// - pop 其一条消息
|
||||
/// - 该渠道还有消息 → 放回 pending 尾部
|
||||
pub fn dequeue(&mut self) -> Option<PipelineMessage> {
|
||||
while let Some(ch) = self.pending.pop_front() {
|
||||
if let Some(q) = self.queues.get_mut(&ch) {
|
||||
if let Some(item) = q.pop_front() {
|
||||
if !q.is_empty() {
|
||||
self.pending.push_back(ch);
|
||||
} else {
|
||||
self.queues.remove(&ch);
|
||||
}
|
||||
return Some(item);
|
||||
}
|
||||
}
|
||||
self.queues.remove(&ch);
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
pub fn signal(&self) -> Arc<Notify> {
|
||||
Arc::clone(&self.signal)
|
||||
}
|
||||
#[allow(dead_code)]
|
||||
pub fn poll_interval(&self) -> Duration {
|
||||
self.poll_interval
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
pub fn has_pending(&self, ch: &ChannelId) -> bool {
|
||||
self.queues.get(ch).map_or(false, |q| !q.is_empty())
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
pub fn pending_count(&self) -> usize {
|
||||
self.pending.len()
|
||||
}
|
||||
|
||||
/// 清空指定渠道的全部待处理消息
|
||||
#[allow(dead_code)]
|
||||
pub fn drain_channel(&mut self, ch: &ChannelId) -> Vec<PipelineMessage> {
|
||||
let drained = self
|
||||
.queues
|
||||
.remove(ch)
|
||||
.map(|mut q| q.drain(..).collect())
|
||||
.unwrap_or_default();
|
||||
self.pending.retain(|c| c != ch);
|
||||
drained
|
||||
}
|
||||
}
|
||||
|
||||
// ─── 消费者通道标识 ───
|
||||
|
||||
/// 消息路由目标
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum ConsumerTarget {
|
||||
/// LLM 消费者(处理 UserMessage, ToolResult)
|
||||
LLM,
|
||||
/// 工具消费者(处理 ToolCall, ApprovalRequest)
|
||||
Tool,
|
||||
/// 发送消费者(处理 LLMReply)
|
||||
Send,
|
||||
}
|
||||
|
||||
impl PipelineMessage {
|
||||
/// 判断这条消息应该路由到哪个消费者
|
||||
pub fn target(&self) -> ConsumerTarget {
|
||||
match self.kind {
|
||||
MessageKind::UserMessage { .. }
|
||||
| MessageKind::ToolResult { .. }
|
||||
| MessageKind::ScheduledTask { .. } => ConsumerTarget::LLM,
|
||||
MessageKind::ToolCall { .. } | MessageKind::ApprovalRequest { .. } => {
|
||||
ConsumerTarget::Tool
|
||||
}
|
||||
MessageKind::LLMReply { .. } => ConsumerTarget::Send,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ─── 测试 ───
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::queue::QueueRunner;
|
||||
|
||||
fn make_user_msg(user: &str, text: &str) -> PipelineMessage {
|
||||
PipelineMessage::user_message(user, text, "", None)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_target_routing() {
|
||||
let msg = PipelineMessage::user_message("u1", "hello", "m1", None);
|
||||
assert_eq!(msg.target(), ConsumerTarget::LLM);
|
||||
|
||||
let msg = PipelineMessage::llm_reply(ChannelId::wechat("u1"), Uuid::new_v4(), "hi", None);
|
||||
assert_eq!(msg.target(), ConsumerTarget::Send);
|
||||
|
||||
let msg = PipelineMessage::tool_call(
|
||||
ChannelId::wechat("u1"),
|
||||
Uuid::new_v4(),
|
||||
Uuid::new_v4(),
|
||||
"tc1",
|
||||
"query_weather",
|
||||
r#"{"location":"北京"}"#,
|
||||
None,
|
||||
);
|
||||
assert_eq!(msg.target(), ConsumerTarget::Tool);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_fair_round_robin() {
|
||||
let mut q = MessageQueue::new();
|
||||
q.enqueue(make_user_msg("A", "a1"));
|
||||
q.enqueue(make_user_msg("A", "a2"));
|
||||
q.enqueue(make_user_msg("A", "a3"));
|
||||
q.enqueue(make_user_msg("B", "b1"));
|
||||
assert_eq!(q.dequeue().unwrap().channel.user_id, "A");
|
||||
assert_eq!(q.dequeue().unwrap().channel.user_id, "B");
|
||||
assert_eq!(q.dequeue().unwrap().channel.user_id, "A");
|
||||
assert_eq!(q.dequeue().unwrap().channel.user_id, "A");
|
||||
assert!(q.dequeue().is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_single_user() {
|
||||
let mut q = MessageQueue::new();
|
||||
q.enqueue(make_user_msg("X", "hi"));
|
||||
q.enqueue(make_user_msg("X", "hello"));
|
||||
assert_eq!(q.dequeue().unwrap().channel.user_id, "X");
|
||||
assert_eq!(q.dequeue().unwrap().channel.user_id, "X");
|
||||
assert!(q.dequeue().is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_empty_queue() {
|
||||
let mut q = MessageQueue::new();
|
||||
assert!(q.dequeue().is_none());
|
||||
assert_eq!(q.pending_count(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_user_id_preserved() {
|
||||
let msg = PipelineMessage::user_message("wx_user_1", "天气?", "m1", None);
|
||||
assert_eq!(msg.channel.user_id, "wx_user_1");
|
||||
|
||||
let reply = PipelineMessage::llm_reply(ChannelId::wechat("wx_user_1"), msg.correlation_id, "晴天", None);
|
||||
assert_eq!(reply.channel.user_id, "wx_user_1");
|
||||
assert_eq!(reply.correlation_id, msg.correlation_id);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_mixed_item_types() {
|
||||
let mut q = MessageQueue::new();
|
||||
q.enqueue(PipelineMessage::user_message("U", "user msg", "1", None));
|
||||
q.enqueue(PipelineMessage {
|
||||
channel: ChannelId::wechat("U"),
|
||||
correlation_id: Uuid::new_v4(),
|
||||
kind: MessageKind::ScheduledTask {
|
||||
text: "task result".into(),
|
||||
task_name: "reminder".into(),
|
||||
},
|
||||
});
|
||||
assert!(matches!(
|
||||
q.dequeue().unwrap().kind,
|
||||
MessageKind::UserMessage { .. }
|
||||
));
|
||||
assert!(matches!(
|
||||
q.dequeue().unwrap().kind,
|
||||
MessageKind::ScheduledTask { .. }
|
||||
));
|
||||
}
|
||||
|
||||
// ─── 端到端关联测试 ───
|
||||
|
||||
/// 验证 user_id 从用户消息到 LLM 回复全程不变
|
||||
#[test]
|
||||
fn test_user_id_pipeline_preserved() {
|
||||
let user_id = "wx_user_1";
|
||||
|
||||
// 用户消息入队
|
||||
let user_msg = PipelineMessage::user_message(user_id, "天气怎么样?", "msg1", None);
|
||||
assert_eq!(user_msg.channel.user_id, user_id);
|
||||
let correlation_id = user_msg.correlation_id;
|
||||
|
||||
// LLM 回复(模拟 LLM Consumer 生成)
|
||||
let reply = PipelineMessage::llm_reply(ChannelId::wechat(user_id), correlation_id, "今天晴天", None);
|
||||
assert_eq!(reply.channel.user_id, user_id);
|
||||
assert_eq!(reply.correlation_id, correlation_id);
|
||||
|
||||
// 工具调用(模拟 LLM Consumer 生成)
|
||||
let tool_call = PipelineMessage::tool_call(
|
||||
ChannelId::wechat(user_id),
|
||||
correlation_id,
|
||||
Uuid::new_v4(),
|
||||
"tc1",
|
||||
"query_weather",
|
||||
r#"{"location":"北京"}"#,
|
||||
None,
|
||||
);
|
||||
assert_eq!(tool_call.channel.user_id, user_id);
|
||||
assert_eq!(tool_call.correlation_id, correlation_id);
|
||||
|
||||
// 工具结果回喂(模拟 Tool Consumer 生成)
|
||||
let tool_result = PipelineMessage::tool_result(
|
||||
ChannelId::wechat(user_id),
|
||||
correlation_id,
|
||||
Uuid::new_v4(),
|
||||
"tc1",
|
||||
"query_weather",
|
||||
"25°C",
|
||||
None,
|
||||
);
|
||||
assert_eq!(tool_result.channel.user_id, user_id);
|
||||
assert_eq!(tool_result.correlation_id, correlation_id);
|
||||
}
|
||||
|
||||
/// 验证多个用户的消息不会串
|
||||
#[test]
|
||||
fn test_multi_user_isolation() {
|
||||
let mut q = MessageQueue::new();
|
||||
|
||||
q.enqueue(PipelineMessage::user_message("u1", "你好", "1", None));
|
||||
q.enqueue(PipelineMessage::user_message("u2", "hello", "2", None));
|
||||
q.enqueue(PipelineMessage::user_message("u1", "天气?", "3", None));
|
||||
|
||||
// 按公平轮转出队,验证 user_id 正确
|
||||
let msg1 = q.dequeue().unwrap();
|
||||
assert_eq!(msg1.channel.user_id, "u1");
|
||||
let msg2 = q.dequeue().unwrap();
|
||||
assert_eq!(msg2.channel.user_id, "u2");
|
||||
let msg3 = q.dequeue().unwrap();
|
||||
assert_eq!(msg3.channel.user_id, "u1");
|
||||
}
|
||||
|
||||
/// 验证路由目标正确
|
||||
#[test]
|
||||
fn test_consumer_target_routing() {
|
||||
// UserMessage → LLM
|
||||
let m = PipelineMessage::user_message("u1", "hi", "1", None);
|
||||
assert_eq!(m.target(), ConsumerTarget::LLM);
|
||||
|
||||
// ToolResult → LLM
|
||||
let m = PipelineMessage::tool_result(ChannelId::wechat("u1"), Uuid::new_v4(), Uuid::new_v4(), "tc1", "test", "ok", None);
|
||||
assert_eq!(m.target(), ConsumerTarget::LLM);
|
||||
|
||||
// ScheduledTask → LLM
|
||||
let m = PipelineMessage {
|
||||
channel: ChannelId::wechat("u1"),
|
||||
correlation_id: Uuid::new_v4(),
|
||||
kind: MessageKind::ScheduledTask { text: "task".into(), task_name: "t1".into() },
|
||||
};
|
||||
assert_eq!(m.target(), ConsumerTarget::LLM);
|
||||
|
||||
// ToolCall → Tool
|
||||
let m = PipelineMessage::tool_call(ChannelId::wechat("u1"), Uuid::new_v4(), Uuid::new_v4(), "tc1", "weather", "{}", None);
|
||||
assert_eq!(m.target(), ConsumerTarget::Tool);
|
||||
|
||||
// ApprovalRequest → Tool
|
||||
let m = PipelineMessage {
|
||||
channel: ChannelId::wechat("u1"),
|
||||
correlation_id: Uuid::new_v4(),
|
||||
kind: MessageKind::ApprovalRequest {
|
||||
session_id: Uuid::new_v4(),
|
||||
tool_name: "high_risk_tool".into(),
|
||||
tool_args: "{}".into(),
|
||||
approval_prompt: "".into(),
|
||||
},
|
||||
};
|
||||
assert_eq!(m.target(), ConsumerTarget::Tool);
|
||||
|
||||
// LLMReply → Send
|
||||
let m = PipelineMessage::llm_reply(ChannelId::wechat("u1"), Uuid::new_v4(), "reply", None);
|
||||
assert_eq!(m.target(), ConsumerTarget::Send);
|
||||
}
|
||||
|
||||
/// 验证 QueueRunner 的路由逻辑:入队 → 路由到正确通道
|
||||
#[tokio::test]
|
||||
async fn test_queue_runner_routing() {
|
||||
let (runner, mut channels) = QueueRunner::new(16);
|
||||
let enqueue = runner.enqueue_handle();
|
||||
|
||||
// 启动 runner(在后台运行)
|
||||
let handle = tokio::spawn(async move {
|
||||
runner.run().await;
|
||||
});
|
||||
|
||||
// 入队一个 UserMessage → 应到 llm_rx
|
||||
let user_msg = PipelineMessage::user_message("u1", "hello", "m1", None);
|
||||
enqueue.enqueue(user_msg).await;
|
||||
|
||||
// 入队一个 LLMReply → 应到 send_rx
|
||||
let reply_msg = PipelineMessage::llm_reply(ChannelId::wechat("u1"), Uuid::new_v4(), "reply", None);
|
||||
enqueue.enqueue(reply_msg).await;
|
||||
|
||||
// 验证 llm_rx 收到 UserMessage
|
||||
let llm_msg = tokio::time::timeout(
|
||||
std::time::Duration::from_secs(2),
|
||||
channels.llm_rx.recv(),
|
||||
)
|
||||
.await
|
||||
.expect("llm_rx 超时")
|
||||
.expect("llm_rx 通道关闭");
|
||||
assert_eq!(llm_msg.channel.user_id, "u1");
|
||||
assert!(matches!(llm_msg.kind, MessageKind::UserMessage { .. }));
|
||||
|
||||
// 验证 send_rx 收到 LLMReply
|
||||
let send_msg = tokio::time::timeout(
|
||||
std::time::Duration::from_secs(2),
|
||||
channels.send_rx.recv(),
|
||||
)
|
||||
.await
|
||||
.expect("send_rx 超时")
|
||||
.expect("send_rx 通道关闭");
|
||||
assert_eq!(send_msg.channel.user_id, "u1");
|
||||
assert!(matches!(send_msg.kind, MessageKind::LLMReply { .. }));
|
||||
|
||||
// 关闭 runner
|
||||
handle.abort();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,5 @@
|
||||
pub mod message_queue;
|
||||
pub mod runner;
|
||||
|
||||
pub use message_queue::{ConsumerTarget, MessageKind, MessageQueue, PipelineMessage};
|
||||
pub use runner::{ConsumerChannels, EnqueueHandle, QueueRunner};
|
||||
@@ -0,0 +1,212 @@
|
||||
//! QueueRunner — 消息队列运行时
|
||||
//!
|
||||
//! 从公平轮转 MessageQueue 中出队消息,按 MessageKind 路由到对应的
|
||||
//! tokio mpsc 通道,由注册的消费者 task 处理。
|
||||
//!
|
||||
//! # 架构
|
||||
//!
|
||||
//! ```text
|
||||
//! ┌──────────────┐
|
||||
//! enqueue ──────→ │ MessageQueue │ ← 公平轮转
|
||||
//! └──────┬───────┘
|
||||
//! │ dequeue
|
||||
//! ▼
|
||||
//! ┌──────────────┐
|
||||
//! │ QueueRunner │ ← 路由循环
|
||||
//! └──┬───┬───┬──┘
|
||||
//! ┌─┘ │ └─┐
|
||||
//! ▼ ▼ ▼
|
||||
//! llm_tx tool_tx send_tx
|
||||
//! │ │ │
|
||||
//! ┌────┘ ┌─┘ ┌──┘
|
||||
//! ▼ ▼ ▼
|
||||
//! LLMCons. ToolCons. SendCons.
|
||||
//! ```
|
||||
|
||||
use crate::queue::{ConsumerTarget, MessageQueue, PipelineMessage};
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::{mpsc, Mutex, Notify};
|
||||
use tracing::{error, info, warn};
|
||||
|
||||
/// 消费者通道集 —— 创建 QueueRunner 时获得接收端
|
||||
pub struct ConsumerChannels {
|
||||
pub llm_rx: mpsc::Receiver<PipelineMessage>,
|
||||
pub tool_rx: mpsc::Receiver<PipelineMessage>,
|
||||
pub send_rx: mpsc::Receiver<PipelineMessage>,
|
||||
}
|
||||
|
||||
/// 队列运行时 —— 负责从公平轮转队列出队并按类型路由
|
||||
pub struct QueueRunner {
|
||||
/// 公平轮转队列
|
||||
queue: Arc<Mutex<MessageQueue>>,
|
||||
/// LLM 消费者通道
|
||||
llm_tx: mpsc::Sender<PipelineMessage>,
|
||||
/// 工具消费者通道
|
||||
tool_tx: mpsc::Sender<PipelineMessage>,
|
||||
/// 发送消费者通道
|
||||
send_tx: mpsc::Sender<PipelineMessage>,
|
||||
/// 关闭信号
|
||||
shutdown: Arc<AtomicBool>,
|
||||
/// 队列信号(空→非空通知)
|
||||
signal: Arc<Notify>,
|
||||
/// 轮询间隔
|
||||
poll_interval: std::time::Duration,
|
||||
}
|
||||
|
||||
impl QueueRunner {
|
||||
/// 创建 QueueRunner 和对应的消费者通道
|
||||
///
|
||||
/// `channel_size`: 每个 mpsc 通道的缓冲区大小
|
||||
pub fn new(channel_size: usize) -> (Self, ConsumerChannels) {
|
||||
let (llm_tx, llm_rx) = mpsc::channel(channel_size);
|
||||
let (tool_tx, tool_rx) = mpsc::channel(channel_size);
|
||||
let (send_tx, send_rx) = mpsc::channel(channel_size);
|
||||
|
||||
let queue = MessageQueue::new();
|
||||
let signal = queue.signal();
|
||||
let queue = Arc::new(Mutex::new(queue));
|
||||
|
||||
let runner = Self {
|
||||
queue,
|
||||
llm_tx,
|
||||
tool_tx,
|
||||
send_tx,
|
||||
shutdown: Arc::new(AtomicBool::new(false)),
|
||||
signal,
|
||||
poll_interval: std::time::Duration::from_millis(200),
|
||||
};
|
||||
|
||||
let channels = ConsumerChannels {
|
||||
llm_rx,
|
||||
tool_rx,
|
||||
send_rx,
|
||||
};
|
||||
|
||||
(runner, channels)
|
||||
}
|
||||
|
||||
/// 获取队列的写句柄(用于外部入队)
|
||||
pub fn enqueue_handle(&self) -> EnqueueHandle {
|
||||
EnqueueHandle {
|
||||
queue: self.queue.clone(),
|
||||
shutdown: self.shutdown.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
/// 获取关闭信号句柄
|
||||
pub fn shutdown_handle(&self) -> ShutdownHandle {
|
||||
ShutdownHandle {
|
||||
shutdown: self.shutdown.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
/// 启动路由主循环(阻塞,应在 tokio::spawn 中运行)
|
||||
pub async fn run(&self) {
|
||||
info!("QueueRunner 已启动");
|
||||
loop {
|
||||
// 检查关闭信号
|
||||
if self.shutdown.load(Ordering::Relaxed) {
|
||||
info!("QueueRunner 收到关闭信号,停止路由");
|
||||
break;
|
||||
}
|
||||
|
||||
// 出队
|
||||
let item = {
|
||||
let mut q = self.queue.lock().await;
|
||||
q.dequeue()
|
||||
};
|
||||
|
||||
if let Some(msg) = item {
|
||||
// 预提取字段(避免 borrow + move 冲突)
|
||||
let target = msg.target();
|
||||
let user_id = msg.channel.user_id.clone();
|
||||
let correlation_id = msg.correlation_id;
|
||||
|
||||
let result = match target {
|
||||
ConsumerTarget::LLM => self.llm_tx.send(msg).await,
|
||||
ConsumerTarget::Tool => self.tool_tx.send(msg).await,
|
||||
ConsumerTarget::Send => self.send_tx.send(msg).await,
|
||||
};
|
||||
|
||||
if let Err(_e) = result {
|
||||
warn!(
|
||||
"路由消息失败 (target={:?}, user={}, corr={}): 通道已关闭,跳过",
|
||||
target, user_id, correlation_id
|
||||
);
|
||||
// 不中断整个路由循环,跳过该消息继续处理
|
||||
}
|
||||
} else {
|
||||
// 队列为空,等待通知或超时
|
||||
tokio::select! {
|
||||
_ = self.signal.notified() => {
|
||||
// 有新的消息入队,继续循环
|
||||
}
|
||||
_ = tokio::time::sleep(self.poll_interval) => {
|
||||
// 超时兜底,继续检查
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
info!("QueueRunner 已停止");
|
||||
}
|
||||
}
|
||||
|
||||
/// 入队句柄 —— 外部模块通过它向队列投放消息
|
||||
#[derive(Clone)]
|
||||
pub struct EnqueueHandle {
|
||||
queue: Arc<Mutex<MessageQueue>>,
|
||||
shutdown: Arc<AtomicBool>,
|
||||
}
|
||||
|
||||
impl EnqueueHandle {
|
||||
/// 向队列投放一条消息
|
||||
pub async fn enqueue(&self, msg: PipelineMessage) {
|
||||
if self.shutdown.load(Ordering::Relaxed) {
|
||||
warn!(
|
||||
"队列已关闭,丢弃消息: user={}",
|
||||
msg.user_id()
|
||||
);
|
||||
return;
|
||||
}
|
||||
self.queue.lock().await.enqueue(msg);
|
||||
}
|
||||
}
|
||||
|
||||
/// 关闭句柄
|
||||
#[derive(Clone)]
|
||||
pub struct ShutdownHandle {
|
||||
shutdown: Arc<AtomicBool>,
|
||||
}
|
||||
|
||||
impl ShutdownHandle {
|
||||
/// 发送关闭信号
|
||||
pub fn shutdown(&self) {
|
||||
self.shutdown.store(true, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
|
||||
/// 等所有消费者处理完并关闭后,等待队列排空
|
||||
#[allow(dead_code)]
|
||||
pub async fn drain_and_shutdown(
|
||||
_enqueue: &EnqueueHandle,
|
||||
runner_shutdown: &ShutdownHandle,
|
||||
consumer_handles: Vec<tokio::task::JoinHandle<()>>,
|
||||
) {
|
||||
info!("开始关闭队列系统...");
|
||||
|
||||
// 1. 阻止新消息入队(设置关闭标志)
|
||||
runner_shutdown.shutdown();
|
||||
|
||||
// 2. 等待排空(给一个合理的时间)
|
||||
tokio::time::sleep(std::time::Duration::from_millis(500)).await;
|
||||
|
||||
// 3. 等待消费者 task 完成
|
||||
for handle in consumer_handles {
|
||||
if let Err(e) = handle.await {
|
||||
error!("消费者 task 结束异常: {:?}", e);
|
||||
}
|
||||
}
|
||||
|
||||
info!("队列系统已完全关闭");
|
||||
}
|
||||
Reference in New Issue
Block a user