Skip to main content

flare_core/client/
router.rs

1//! 客户端消息路由
2//!
3//! 根据消息类型将消息路由到不同的处理器
4//! 支持自定义路由规则和处理器
5
6use crate::common::error::Result;
7use crate::common::protocol::Frame;
8use async_trait::async_trait;
9use std::collections::HashMap;
10use std::sync::Arc;
11use tracing::{debug, warn};
12
13/// 消息处理器
14#[async_trait]
15pub trait MessageHandler: Send + Sync {
16    /// 处理消息
17    ///
18    /// # 参数
19    /// - `frame`: 接收到的消息帧
20    ///
21    /// # 返回
22    /// 如果需要回复,返回 `Some(Frame)`,否则返回 `None`
23    async fn handle(&self, frame: &Frame) -> Result<Option<Frame>>;
24}
25
26/// 消息路由
27///
28/// 根据消息类型将消息路由到不同的处理器
29pub struct MessageRouter {
30    /// 路由规则:消息类型 -> 处理器
31    handlers: HashMap<String, Vec<Arc<dyn MessageHandler>>>,
32    /// 默认处理器(当没有匹配的路由时使用)
33    default_handler: Option<Arc<dyn MessageHandler>>,
34}
35
36impl MessageRouter {
37    /// 创建新的消息路由
38    pub fn new() -> Self {
39        Self {
40            handlers: HashMap::new(),
41            default_handler: None,
42        }
43    }
44
45    /// 注册处理器
46    ///
47    /// # 参数
48    /// - `route`: 路由键(例如 "system.ping", "message.chat" 等)
49    /// - `handler`: 消息处理器
50    pub fn register(&mut self, route: impl Into<String>, handler: Arc<dyn MessageHandler>) {
51        let route = route.into();
52        self.handlers.entry(route).or_default().push(handler);
53    }
54
55    /// 设置默认处理器
56    pub fn set_default_handler(&mut self, handler: Arc<dyn MessageHandler>) {
57        self.default_handler = Some(handler);
58    }
59
60    /// 路由消息
61    ///
62    /// # 参数
63    /// - `frame`: 要路由的消息帧
64    ///
65    /// # 返回
66    /// 所有处理器的回复(如果有)
67    pub async fn route(&self, frame: &Frame) -> Result<Vec<Frame>> {
68        let route_key = Self::extract_route_key(frame);
69        debug!(
70            "Routing message: route={}, frame_id={}",
71            route_key, frame.message_id
72        );
73
74        let mut replies = Vec::new();
75
76        // 查找匹配的处理器
77        if let Some(handlers) = self.handlers.get(&route_key) {
78            for handler in handlers {
79                match handler.handle(frame).await {
80                    Ok(Some(reply)) => {
81                        replies.push(reply);
82                    }
83                    Ok(None) => {
84                        // 处理器不需要回复
85                    }
86                    Err(e) => {
87                        warn!("Handler error for route {}: {}", route_key, e);
88                    }
89                }
90            }
91        } else if let Some(ref default_handler) = self.default_handler {
92            // 使用默认处理器
93            match default_handler.handle(frame).await {
94                Ok(Some(reply)) => {
95                    replies.push(reply);
96                }
97                Ok(None) => {
98                    // 默认处理器不需要回复
99                }
100                Err(e) => {
101                    warn!("Default handler error: {}", e);
102                }
103            }
104        } else {
105            debug!("No handler found for route: {}", route_key);
106        }
107
108        Ok(replies)
109    }
110
111    /// 从 Frame 中提取路由键
112    fn extract_route_key(frame: &Frame) -> String {
113        // 根据 Command 类型提取路由键
114        if let Some(ref command) = frame.command {
115            match &command.r#type {
116                Some(crate::common::protocol::command::Type::System(sys_cmd)) => {
117                    // SystemCommand 的 type 是 i32,使用 TryFrom 替代已弃用的 from_i32
118                    use crate::common::protocol::system_command::Type as SysType;
119                    use std::convert::TryFrom;
120                    match SysType::try_from(sys_cmd.r#type) {
121                        Ok(SysType::Connect) => "system.connect".to_string(),
122                        Ok(SysType::ConnectAck) => "system.connect_ack".to_string(),
123                        Ok(SysType::Close) => "system.close".to_string(),
124                        Ok(SysType::Ping) => "system.ping".to_string(),
125                        Ok(SysType::Pong) => "system.pong".to_string(),
126                        Ok(SysType::Error) => "system.error".to_string(),
127                        Ok(SysType::Event) => "system.event".to_string(),
128                        Ok(SysType::Auth) => "system.auth".to_string(),
129                        Ok(SysType::AuthAck) => "system.auth_ack".to_string(),
130                        _ => "system.unknown".to_string(),
131                    }
132                }
133                Some(crate::common::protocol::command::Type::Payload(msg_cmd)) => {
134                    // 使用消息类型作为路由键(使用 TryFrom 替代已弃用的 from_i32)
135                    use crate::common::protocol::payload_command::Type as MsgType;
136                    use std::convert::TryFrom;
137                    match MsgType::try_from(msg_cmd.r#type) {
138                        Ok(MsgType::Message) => "payload.message".to_string(),
139                        Ok(MsgType::Event) => "payload.event".to_string(),
140                        Ok(MsgType::Ack) => "payload.ack".to_string(),
141                        Ok(MsgType::Data) => "payload.data".to_string(),
142                        _ => format!("payload.{}", msg_cmd.r#type),
143                    }
144                }
145                Some(crate::common::protocol::command::Type::Notification(notif_cmd)) => {
146                    use crate::common::protocol::notification_command::Type as NotifType;
147                    use std::convert::TryFrom;
148                    match NotifType::try_from(notif_cmd.r#type) {
149                        Ok(NotifType::System) => "notification.system".to_string(),
150                        Ok(NotifType::Broadcast) => "notification.broadcast".to_string(),
151                        Ok(NotifType::Alert) => "notification.alert".to_string(),
152                        Ok(NotifType::User) => "notification.user".to_string(),
153                        Ok(NotifType::Connection) => "notification.connection".to_string(),
154                        _ => format!("notification.{}", notif_cmd.r#type),
155                    }
156                }
157                Some(crate::common::protocol::command::Type::Custom(custom_cmd)) => {
158                    format!("custom.{}", custom_cmd.name)
159                }
160                None => "unknown".to_string(),
161            }
162        } else {
163            "unknown".to_string()
164        }
165    }
166
167    /// 移除指定路由的所有处理器
168    pub fn remove_route(&mut self, route: &str) {
169        self.handlers.remove(route);
170    }
171
172    /// 清除所有路由
173    pub fn clear(&mut self) {
174        self.handlers.clear();
175        self.default_handler = None;
176    }
177}
178
179impl Default for MessageRouter {
180    fn default() -> Self {
181        Self::new()
182    }
183}
184
185/// 简单的消息处理器实现
186pub struct SimpleHandler {
187    #[allow(clippy::type_complexity)]
188    handler: Box<dyn Fn(&Frame) -> Result<Option<Frame>> + Send + Sync>,
189}
190
191impl SimpleHandler {
192    /// 创建新的简单处理器
193    pub fn new<F>(handler: F) -> Self
194    where
195        F: Fn(&Frame) -> Result<Option<Frame>> + Send + Sync + 'static,
196    {
197        Self {
198            handler: Box::new(handler),
199        }
200    }
201}
202
203#[async_trait]
204impl MessageHandler for SimpleHandler {
205    async fn handle(&self, frame: &Frame) -> Result<Option<Frame>> {
206        (self.handler)(frame)
207    }
208}
209
210/// 异步消息处理器实现
211pub struct AsyncHandler {
212    #[allow(clippy::type_complexity)]
213    handler: Arc<
214        dyn Fn(
215                &Frame,
216            ) -> std::pin::Pin<
217                Box<dyn std::future::Future<Output = Result<Option<Frame>>> + Send + '_>,
218            > + Send
219            + Sync,
220    >,
221}
222
223impl AsyncHandler {
224    /// 创建新的异步处理器
225    pub fn new<F, Fut>(handler: F) -> Self
226    where
227        F: Fn(&Frame) -> Fut + Send + Sync + 'static,
228        Fut: std::future::Future<Output = Result<Option<Frame>>> + Send + 'static,
229    {
230        Self {
231            handler: Arc::new(move |frame| Box::pin(handler(frame))),
232        }
233    }
234}
235
236#[async_trait]
237impl MessageHandler for AsyncHandler {
238    async fn handle(&self, frame: &Frame) -> Result<Option<Frame>> {
239        (self.handler)(frame).await
240    }
241}