mq_bridge/
type_handler.rs

1use crate::response::ErgonomicResponse;
2use crate::traits::{Handled, Handler, HandlerError};
3use crate::{CanonicalMessage, MessageContext};
4use async_trait::async_trait;
5use futures::future::BoxFuture;
6use serde::de::DeserializeOwned;
7use std::collections::HashMap;
8use std::future::Future;
9use std::sync::Arc;
10
11/// A handler that dispatches messages to other handlers based on a metadata field (e.g., "type").
12///
13/// # Example
14/// ```rust
15/// use mq_bridge::type_handler::TypeHandler;
16/// use mq_bridge::{CanonicalMessage, Handled};
17/// use serde::Deserialize;
18///
19/// #[derive(Deserialize)]
20/// struct MyCommand { id: String }
21///
22/// # fn main() {
23/// let handler = TypeHandler::new()
24///     .add("my_command", |cmd: MyCommand| async move {
25///         println!("Received command: {}", cmd.id);
26///     }); // No return needed! Implicitly returns (), which maps to Handled::Ack.
27/// # }
28/// ```
29#[derive(Clone)]
30pub struct TypeHandler {
31    pub(crate) handlers: HashMap<String, Arc<dyn Handler>>,
32    pub(crate) type_key: String, // will be the key in msg metadata, default is "kind"
33    pub(crate) fallback: Option<Arc<dyn Handler>>,
34}
35
36pub const KIND_KEY: &str = "kind";
37
38/// A helper trait to allow registering handlers with or without context.
39pub trait IntoTypedHandler<T, Args>: Send + Sync + 'static {
40    type Future: Future<Output = Result<Handled, HandlerError>> + Send + 'static;
41    fn call(&self, msg: T, ctx: MessageContext) -> Self::Future;
42}
43
44/// Implementation for standard closures returning Result<Handled, HandlerError>.
45/// This path is unique and allows correct inference for legacy code.
46impl<F, Fut, T> IntoTypedHandler<T, (T, Result<Handled, HandlerError>)> for F
47where
48    T: DeserializeOwned + Send + Sync + 'static,
49    F: Fn(T) -> Fut + Send + Sync + 'static,
50    Fut: Future<Output = Result<Handled, HandlerError>> + Send + 'static,
51{
52    type Future = Fut;
53    fn call(&self, msg: T, _ctx: MessageContext) -> Self::Future {
54        (self)(msg)
55    }
56}
57
58impl<F, Fut, T> IntoTypedHandler<T, (T, MessageContext, Result<Handled, HandlerError>)> for F
59where
60    T: DeserializeOwned + Send + Sync + 'static,
61    F: Fn(T, MessageContext) -> Fut + Send + Sync + 'static,
62    Fut: Future<Output = Result<Handled, HandlerError>> + Send + 'static,
63{
64    type Future = Fut;
65    fn call(&self, msg: T, ctx: MessageContext) -> Self::Future {
66        (self)(msg, ctx)
67    }
68}
69
70/// Ergonomic implementation for closures returning types that implement ErgonomicResponse.
71impl<F, Fut, T, R> IntoTypedHandler<T, (T, R)> for F
72where
73    T: DeserializeOwned + Send + Sync + 'static,
74    R: ErgonomicResponse,
75    F: Fn(T) -> Fut + Send + Sync + 'static,
76    Fut: Future<Output = R> + Send + 'static,
77{
78    type Future = BoxFuture<'static, Result<Handled, HandlerError>>;
79    fn call(&self, msg: T, _ctx: MessageContext) -> Self::Future {
80        let fut = (self)(msg);
81        Box::pin(async move { fut.await.into_handler_result() })
82    }
83}
84
85impl<F, Fut, T, R> IntoTypedHandler<T, (T, MessageContext, R)> for F
86where
87    T: DeserializeOwned + Send + Sync + 'static,
88    R: ErgonomicResponse,
89    F: Fn(T, MessageContext) -> Fut + Send + Sync + 'static,
90    Fut: Future<Output = R> + Send + 'static,
91{
92    type Future = BoxFuture<'static, Result<Handled, HandlerError>>;
93    fn call(&self, msg: T, ctx: MessageContext) -> Self::Future {
94        let fut = (self)(msg, ctx);
95        Box::pin(async move { fut.await.into_handler_result() })
96    }
97}
98
99impl Default for TypeHandler {
100    fn default() -> Self {
101        Self::new()
102    }
103}
104
105impl TypeHandler {
106    /// Creates a new TypeHandler that looks for the specified key in message metadata to determine the message type.
107    pub fn new() -> Self {
108        Self {
109            handlers: HashMap::new(),
110            type_key: KIND_KEY.into(),
111            fallback: None,
112        }
113    }
114
115    /// Registers a generic handler for a specific type name.
116    pub fn add_handler(mut self, type_name: &str, handler: impl Handler + 'static) -> Self {
117        self.handlers
118            .insert(type_name.to_string(), Arc::new(handler));
119        self
120    }
121
122    /// Sets a fallback handler to be used when no type match is found.
123    pub fn with_fallback(mut self, handler: Arc<dyn Handler>) -> Self {
124        self.fallback = Some(handler);
125        self
126    }
127
128    #[doc(hidden)]
129    pub fn add_simple<T, H, Args>(self, type_name: &str, handler: H) -> Self
130    where
131        T: DeserializeOwned + Send + Sync + 'static,
132        H: IntoTypedHandler<T, Args>,
133        Args: Send + Sync + 'static,
134    {
135        self.add(type_name, handler)
136    }
137
138    /// Registers a typed handler function.
139    ///
140    /// The handler can accept either:
141    /// - `fn(T) -> Future<Output = Result<Handled, HandlerError>>`
142    /// - `fn(T, MessageContext) -> Future<Output = Result<Handled, HandlerError>>`
143    pub fn add<T, H, Args>(mut self, type_name: &str, handler: H) -> Self
144    where
145        T: DeserializeOwned + Send + Sync + 'static,
146        H: IntoTypedHandler<T, Args>,
147        Args: Send + Sync + 'static,
148    {
149        let handler = Arc::new(handler);
150        let wrapper = move |msg: CanonicalMessage| {
151            let handler = handler.clone();
152            async move {
153                let data = msg.parse::<T>().map_err(|e| {
154                    HandlerError::NonRetryable(anyhow::anyhow!("Deserialization failed: {}", e))
155                })?;
156                let ctx = MessageContext::from(msg);
157                handler.call(data, ctx).await
158            }
159        };
160        self.handlers
161            .insert(type_name.to_string(), Arc::new(wrapper));
162        self
163    }
164}
165
166#[async_trait]
167impl Handler for TypeHandler {
168    async fn handle(&self, msg: CanonicalMessage) -> Result<Handled, HandlerError> {
169        if let Some(type_val) = msg.metadata.get(&self.type_key) {
170            if let Some(handler) = self.handlers.get(type_val) {
171                return handler.handle(msg).await;
172            }
173        }
174
175        if let Some(fallback) = &self.fallback {
176            return fallback.handle(msg).await;
177        }
178
179        Err(HandlerError::NonRetryable(anyhow::anyhow!(
180            "No handler registered for type: '{:?}' and no fallback provided",
181            msg.metadata.get(&self.type_key)
182        )))
183    }
184
185    async fn handle_many(&self, msgs: Vec<CanonicalMessage>) -> Vec<Result<Handled, HandlerError>> {
186        #[derive(Clone, PartialEq)]
187        enum BatchTarget {
188            Handler(String),
189            Fallback,
190        }
191
192        let mut results: Vec<Option<Result<Handled, HandlerError>>> =
193            std::iter::repeat_with(|| None).take(msgs.len()).collect();
194        let mut runs: Vec<(BatchTarget, Vec<(usize, CanonicalMessage)>)> = Vec::new();
195
196        for (index, msg) in msgs.into_iter().enumerate() {
197            let target = if let Some(type_val) = msg.metadata.get(&self.type_key) {
198                if self.handlers.contains_key(type_val) {
199                    Some(BatchTarget::Handler(type_val.clone()))
200                } else if self.fallback.is_some() {
201                    Some(BatchTarget::Fallback)
202                } else {
203                    None
204                }
205            } else if self.fallback.is_some() {
206                Some(BatchTarget::Fallback)
207            } else {
208                None
209            };
210
211            if let Some(target) = target {
212                if let Some((last_target, group)) = runs.last_mut() {
213                    if *last_target == target {
214                        group.push((index, msg));
215                        continue;
216                    }
217                }
218                runs.push((target, vec![(index, msg)]));
219            } else {
220                results[index] = Some(Err(HandlerError::NonRetryable(anyhow::anyhow!(
221                    "No handler registered for type: '{:?}' and no fallback provided",
222                    msg.metadata.get(&self.type_key)
223                ))));
224            }
225        }
226
227        for (target, group) in runs {
228            let expected = group.len();
229            let (indices, messages): (Vec<_>, Vec<_>) = group.into_iter().unzip();
230            let (handler, label) = match target {
231                BatchTarget::Handler(type_name) => {
232                    let Some(handler) = self.handlers.get(&type_name) else {
233                        for index in indices {
234                            results[index] = Some(Err(HandlerError::NonRetryable(
235                                anyhow::anyhow!("Handler batch dispatch did not produce a result"),
236                            )));
237                        }
238                        continue;
239                    };
240                    (handler, format!("Handler for type '{type_name}'"))
241                }
242                BatchTarget::Fallback => {
243                    let Some(handler) = &self.fallback else {
244                        for index in indices {
245                            results[index] = Some(Err(HandlerError::NonRetryable(
246                                anyhow::anyhow!("Handler batch dispatch did not produce a result"),
247                            )));
248                        }
249                        continue;
250                    };
251                    (handler, "Fallback handler".to_string())
252                }
253            };
254
255            let group_results = handler.handle_many(messages).await;
256
257            if group_results.len() != expected {
258                for index in indices {
259                    results[index] = Some(Err(HandlerError::NonRetryable(anyhow::anyhow!(
260                        "{} returned {} results for {} messages",
261                        label,
262                        group_results.len(),
263                        expected
264                    ))));
265                }
266                continue;
267            }
268
269            for (index, result) in indices.into_iter().zip(group_results) {
270                results[index] = Some(result);
271            }
272        }
273
274        results
275            .into_iter()
276            .map(|result| {
277                result.unwrap_or_else(|| {
278                    Err(HandlerError::NonRetryable(anyhow::anyhow!(
279                        "Handler batch dispatch did not produce a result"
280                    )))
281                })
282            })
283            .collect()
284    }
285
286    fn register_handler(
287        &self,
288        type_name: &str,
289        handler: Arc<dyn Handler>,
290    ) -> Option<Arc<dyn Handler>> {
291        let mut th = self.clone();
292        th.handlers.insert(type_name.to_string(), handler);
293        Some(Arc::new(th))
294    }
295}
296
297#[cfg(test)]
298mod tests {
299    use super::*;
300    use crate::msg;
301    use serde::{Deserialize, Serialize};
302    use std::sync::{Arc, Mutex};
303
304    #[derive(Serialize, Deserialize)]
305    struct TestMsg {
306        val: String,
307    }
308
309    #[tokio::test]
310    async fn test_typed_handler_dispatch() {
311        let handler = TypeHandler::new().add("test_a", |msg: TestMsg| async move {
312            assert_eq!(msg.val, "hello");
313            Ok(Handled::Ack)
314        });
315
316        let msg = msg!(
317            &TestMsg {
318                val: "hello".into(),
319            },
320            "test_a"
321        );
322
323        let res = handler.handle(msg).await;
324        assert!(res.is_ok());
325    }
326
327    #[tokio::test]
328    async fn test_typed_handler_with_context() {
329        let handler =
330            TypeHandler::new().add("test_ctx", |msg: TestMsg, ctx: MessageContext| async move {
331                assert_eq!(msg.val, "hello");
332                assert_eq!(ctx.metadata.get("meta").map(|s| s.as_str()), Some("data"));
333                Ok(Handled::Ack)
334            });
335
336        let msg = CanonicalMessage::from_type(&TestMsg {
337            val: "hello".into(),
338        })
339        .unwrap()
340        .with_metadata(HashMap::from([
341            ("kind".to_string(), "test_ctx".to_string()),
342            ("meta".to_string(), "data".to_string()),
343        ]));
344
345        let res = handler.handle(msg).await;
346        assert!(res.is_ok());
347    }
348
349    #[tokio::test]
350    async fn test_typed_handler_no_match_error() {
351        let handler = TypeHandler::new();
352        let msg = msg!(b"{}".to_vec(), "unknown");
353
354        let res = handler.handle(msg).await;
355        assert!(res.is_err());
356        match res.unwrap_err() {
357            HandlerError::NonRetryable(e) => {
358                assert!(e.to_string().contains("No handler registered"))
359            }
360            _ => panic!("Expected NonRetryable error"),
361        }
362    }
363
364    #[tokio::test]
365    async fn test_typed_handler_fallback_ack() {
366        let fallback = Arc::new(|_: CanonicalMessage| async { Ok(Handled::Ack) });
367        let handler = TypeHandler::new().with_fallback(fallback);
368
369        let msg = msg!(b"{}".to_vec(), "unknown");
370
371        let res = handler.handle(msg).await;
372        assert!(matches!(res, Ok(Handled::Ack)));
373    }
374
375    #[tokio::test]
376    async fn test_typed_handler_failure() {
377        let handler = TypeHandler::new().add("fail", |_: TestMsg| async {
378            Result::<Handled, HandlerError>::Err(HandlerError::Retryable(anyhow::anyhow!(
379                "failure"
380            )))
381        });
382
383        let msg = CanonicalMessage::from_type(&TestMsg { val: "x".into() })
384            .unwrap()
385            .with_type_key("fail");
386
387        let res = handler.handle(msg).await;
388        assert!(matches!(res, Err(HandlerError::Retryable(_))));
389    }
390
391    #[tokio::test]
392    async fn test_typed_handler_missing_type_key() {
393        let handler = TypeHandler::new().add("test", |_: TestMsg| async { Ok(Handled::Ack) });
394
395        // Message without "kind" metadata
396        let msg = CanonicalMessage::new(b"{}".to_vec(), None);
397
398        let res = handler.handle(msg).await;
399        assert!(res.is_err());
400    }
401
402    #[tokio::test]
403    async fn test_typed_handler_deserialization_failure() {
404        let handler = TypeHandler::new().add("test", |_: TestMsg| async { Ok(Handled::Ack) });
405
406        // Invalid JSON for TestMsg (missing required field)
407        let msg = CanonicalMessage::new(b"{}".to_vec(), None)
408            .with_metadata(HashMap::from([("kind".to_string(), "test".to_string())]));
409
410        let res = handler.handle(msg).await;
411        assert!(matches!(res, Err(HandlerError::NonRetryable(_))));
412    }
413
414    #[tokio::test]
415    async fn test_typed_handler_handle_many_preserves_dispatch_order() {
416        struct RecordingHandler {
417            name: &'static str,
418            calls: Arc<Mutex<Vec<String>>>,
419        }
420
421        #[async_trait]
422        impl Handler for RecordingHandler {
423            async fn handle(&self, msg: CanonicalMessage) -> Result<Handled, HandlerError> {
424                self.calls
425                    .lock()
426                    .unwrap()
427                    .push(format!("{}:{}", self.name, msg.get_payload_str()));
428                Ok(Handled::Ack)
429            }
430
431            async fn handle_many(
432                &self,
433                msgs: Vec<CanonicalMessage>,
434            ) -> Vec<Result<Handled, HandlerError>> {
435                let payloads = msgs
436                    .iter()
437                    .map(|msg| msg.get_payload_str())
438                    .collect::<Vec<_>>()
439                    .join(",");
440                self.calls
441                    .lock()
442                    .unwrap()
443                    .push(format!("{}:{}", self.name, payloads));
444                msgs.into_iter().map(|_| Ok(Handled::Ack)).collect()
445            }
446        }
447
448        fn typed_msg(payload: &str, kind: &str) -> CanonicalMessage {
449            CanonicalMessage::from(payload)
450                .with_metadata(HashMap::from([("kind".to_string(), kind.to_string())]))
451        }
452
453        let calls = Arc::new(Mutex::new(Vec::new()));
454        let handler = TypeHandler::new()
455            .add_handler(
456                "a",
457                RecordingHandler {
458                    name: "a",
459                    calls: calls.clone(),
460                },
461            )
462            .add_handler(
463                "b",
464                RecordingHandler {
465                    name: "b",
466                    calls: calls.clone(),
467                },
468            );
469
470        let results = handler
471            .handle_many(vec![
472                typed_msg("a1", "a"),
473                typed_msg("a2", "a"),
474                typed_msg("b1", "b"),
475                typed_msg("a3", "a"),
476            ])
477            .await;
478
479        assert!(results.into_iter().all(|result| result.is_ok()));
480        assert_eq!(
481            *calls.lock().unwrap(),
482            vec![
483                "a:a1,a2".to_string(),
484                "b:b1".to_string(),
485                "a:a3".to_string()
486            ]
487        );
488    }
489
490    #[tokio::test]
491    async fn test_cqrs_pattern_example() {
492        #[derive(Serialize, Deserialize)]
493        struct SubmitOrder {
494            id: u32,
495        }
496
497        #[derive(Serialize, Deserialize)]
498        struct OrderSubmitted {
499            id: u32,
500        }
501
502        // 1. Command Handler (Write Side)
503        let command_bus = TypeHandler::new().add("submit_order", |cmd: SubmitOrder| async move {
504            // Execute business logic...
505            // Emit event
506            let evt = OrderSubmitted { id: cmd.id };
507            Ok(Handled::Publish(msg!(&evt, "order_submitted")))
508        });
509
510        // 2. Event Handler (Read Side / Projection)
511        let projection_handler =
512            TypeHandler::new().add("order_submitted", |evt: OrderSubmitted| async move {
513                // Update read database / cache
514                assert_eq!(evt.id, 101);
515                Ok(Handled::Ack)
516            });
517
518        // Simulate incoming command
519        let cmd = SubmitOrder { id: 101 };
520        let cmd_msg = msg!(&cmd, "submit_order");
521
522        // Process command
523        let result = command_bus.handle(cmd_msg).await.unwrap();
524
525        if let Handled::Publish(event_msg) = result {
526            // Verify event type
527            assert_eq!(
528                event_msg.metadata.get("kind").map(|s| s.as_str()),
529                Some("order_submitted")
530            );
531
532            // Process event (Projection)
533            let proj_result = projection_handler.handle(event_msg).await.unwrap();
534            assert!(matches!(proj_result, Handled::Ack));
535        } else {
536            panic!("Expected Handled::Publish");
537        }
538    }
539
540    #[tokio::test]
541    async fn test_cqrs_integration_with_routes() {
542        use crate::models::{Endpoint, Route};
543        use std::sync::atomic::{AtomicU32, Ordering};
544
545        #[derive(Serialize, Deserialize)]
546        struct SubmitOrder {
547            id: u32,
548        }
549
550        #[derive(Serialize, Deserialize)]
551        struct OrderSubmitted {
552            id: u32,
553        }
554
555        // Shared state to verify projection update
556        let read_model_state = Arc::new(AtomicU32::new(0));
557        let read_model_clone = read_model_state.clone();
558
559        // 1. Command Handler (Write Side)
560        let command_handler =
561            TypeHandler::new().add("submit_order", |cmd: SubmitOrder| async move {
562                let evt = OrderSubmitted { id: cmd.id };
563                Ok(Handled::Publish(msg!(&evt, "order_submitted")))
564            });
565
566        // 2. Event Handler (Read Side)
567        let event_handler =
568            TypeHandler::new().add("order_submitted", move |evt: OrderSubmitted| {
569                let state = read_model_clone.clone();
570                async move {
571                    state.store(evt.id, Ordering::SeqCst);
572                    Ok(Handled::Ack)
573                }
574            });
575
576        // 3. Define Endpoints & Routes
577        let cmd_in_ep = Endpoint::new_memory("cmd_in", 10);
578        let event_bus_ep = Endpoint::new_memory("event_bus", 10);
579        let proj_out_ep = Endpoint::new_memory("proj_out", 10);
580
581        let command_route =
582            Route::new(cmd_in_ep.clone(), event_bus_ep.clone()).with_handler(command_handler);
583
584        let event_route =
585            Route::new(event_bus_ep.clone(), proj_out_ep.clone()).with_handler(event_handler);
586
587        // 4. Run Routes
588        let h1 = tokio::spawn(async move {
589            command_route
590                .run_until_err("command_route", None, None)
591                .await
592        });
593        let h2 =
594            tokio::spawn(async move { event_route.run_until_err("event_route", None, None).await });
595
596        // 5. Send Command
597        let cmd_channel = cmd_in_ep.channel().unwrap();
598        let cmd = SubmitOrder { id: 777 };
599        let msg = CanonicalMessage::from_type(&cmd)
600            .unwrap()
601            .with_type_key("submit_order");
602        cmd_channel.send_message(msg).await.unwrap();
603
604        // 6. Wait for consistency
605        let mut attempts = 0;
606        while read_model_state.load(Ordering::SeqCst) != 777 && attempts < 50 {
607            tokio::time::sleep(std::time::Duration::from_millis(20)).await;
608            attempts += 1;
609        }
610
611        assert_eq!(read_model_state.load(Ordering::SeqCst), 777);
612
613        // Cleanup
614        cmd_channel.close();
615        event_bus_ep.channel().unwrap().close();
616
617        let _ = h1.await;
618        let _ = h2.await;
619    }
620}
mq_bridge/type_handler.rs

mq_bridge/
type_handler.rs
