flare_core/client/
router.rs1use 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#[async_trait]
15pub trait MessageHandler: Send + Sync {
16 async fn handle(&self, frame: &Frame) -> Result<Option<Frame>>;
24}
25
26pub struct MessageRouter {
30 handlers: HashMap<String, Vec<Arc<dyn MessageHandler>>>,
32 default_handler: Option<Arc<dyn MessageHandler>>,
34}
35
36impl MessageRouter {
37 pub fn new() -> Self {
39 Self {
40 handlers: HashMap::new(),
41 default_handler: None,
42 }
43 }
44
45 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 pub fn set_default_handler(&mut self, handler: Arc<dyn MessageHandler>) {
57 self.default_handler = Some(handler);
58 }
59
60 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 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 }
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 match default_handler.handle(frame).await {
94 Ok(Some(reply)) => {
95 replies.push(reply);
96 }
97 Ok(None) => {
98 }
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 fn extract_route_key(frame: &Frame) -> String {
113 if let Some(ref command) = frame.command {
115 match &command.r#type {
116 Some(crate::common::protocol::command::Type::System(sys_cmd)) => {
117 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 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 pub fn remove_route(&mut self, route: &str) {
169 self.handlers.remove(route);
170 }
171
172 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
185pub struct SimpleHandler {
187 #[allow(clippy::type_complexity)]
188 handler: Box<dyn Fn(&Frame) -> Result<Option<Frame>> + Send + Sync>,
189}
190
191impl SimpleHandler {
192 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
210pub 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 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}