Skip to main content

mq_bridge/
traits.rs

1//  mq-bridge
2//  © Copyright 2025, by Marco Mengelkoch
3//  Licensed under MIT License, see License file for more details
4//  git clone https://github.com/marcomq/mq-bridge
5
6pub use crate::errors::{ConsumerError, HandlerError, PublisherError};
7pub use crate::outcomes::{Handled, Received, ReceivedBatch, Sent, SentBatch};
8use crate::CanonicalMessage;
9use anyhow::anyhow;
10use async_trait::async_trait;
11pub use futures::future::BoxFuture;
12use std::any::Any;
13use std::sync::Arc;
14use tracing::warn;
15
16/// The disposition of a processed message.
17///
18/// Implements `From<Option<CanonicalMessage>>` for compatibility:
19/// `None` maps to `Ack`, `Some(msg)` maps to `Reply(msg)`.
20#[derive(Default, Debug, Clone)]
21#[allow(clippy::large_enum_variant)]
22pub enum MessageDisposition {
23    /// Acknowledge processing (success).
24    #[default]
25    Ack,
26    /// Acknowledge processing and send a reply.
27    Reply(CanonicalMessage),
28    /// Negative acknowledgement (failure).
29    Nack,
30}
31
32impl From<Option<CanonicalMessage>> for MessageDisposition {
33    fn from(opt: Option<CanonicalMessage>) -> Self {
34        match opt {
35            Some(msg) => MessageDisposition::Reply(msg),
36            None => MessageDisposition::Ack,
37        }
38    }
39}
40
41impl From<Handled> for MessageDisposition {
42    fn from(handled: Handled) -> Self {
43        match handled {
44            Handled::Ack => MessageDisposition::Ack,
45            Handled::Publish(msg) => MessageDisposition::Reply(msg),
46        }
47    }
48}
49
50/// A generic trait for handling messages (commands or events).
51///
52/// Handlers process an incoming message and can optionally return a new
53/// message (e.g. a reply) via `Handled::Publish`, or acknowledge processing via `Handled::Ack`.
54#[async_trait]
55pub trait Handler: Send + Sync + 'static {
56    async fn handle(&self, msg: CanonicalMessage) -> Result<Handled, HandlerError>;
57
58    async fn handle_many(&self, msgs: Vec<CanonicalMessage>) -> Vec<Result<Handled, HandlerError>> {
59        let mut results = Vec::with_capacity(msgs.len());
60        let mut remaining = msgs.len();
61        for msg in msgs {
62            remaining -= 1;
63            let result = self.handle(msg).await;
64            let aborted = match &result {
65                Err(HandlerError::Retryable(_)) => Some("retryable"),
66                Err(HandlerError::Connection(_)) => Some("connection"),
67                Err(HandlerError::NonRetryable(_)) => Some("non-retryable"),
68                Ok(_) => None,
69            };
70            results.push(result);
71            if let Some(kind) = aborted {
72                for _ in 0..remaining {
73                    results.push(Err(match kind {
74                        "retryable" => HandlerError::Retryable(anyhow!(
75                            "batch aborted after earlier retryable handler failure"
76                        )),
77                        "connection" => HandlerError::Connection(anyhow!(
78                            "batch aborted after earlier handler connection failure"
79                        )),
80                        _ => HandlerError::NonRetryable(anyhow!(
81                            "batch aborted after earlier non-retryable handler failure"
82                        )),
83                    }));
84                }
85                break;
86            }
87        }
88        results
89    }
90
91    /// Tries to register a handler for a specific type.
92    /// Returns `None` if this handler does not support registration (e.g. it's not a TypeHandler).
93    fn register_handler(
94        &self,
95        _type_name: &str,
96        _handler: Arc<dyn Handler>,
97    ) -> Option<Arc<dyn Handler>> {
98        None
99    }
100}
101
102#[async_trait]
103impl<T: Handler + ?Sized> Handler for Arc<T> {
104    async fn handle(&self, msg: CanonicalMessage) -> Result<Handled, HandlerError> {
105        (**self).handle(msg).await
106    }
107
108    async fn handle_many(&self, msgs: Vec<CanonicalMessage>) -> Vec<Result<Handled, HandlerError>> {
109        (**self).handle_many(msgs).await
110    }
111
112    fn register_handler(
113        &self,
114        type_name: &str,
115        handler: Arc<dyn Handler>,
116    ) -> Option<Arc<dyn Handler>> {
117        (**self).register_handler(type_name, handler)
118    }
119}
120
121/// A helper trait that allows implementing handlers using native `async fn` syntax
122/// without the `#[async_trait]` macro.
123///
124/// Implementations of this trait can be adapted to `Handler` using `SimpleHandler`.
125pub trait AsyncHandler: Send + Sync + 'static {
126    fn handle<'a>(&'a self, msg: CanonicalMessage) -> BoxFuture<'a, Result<Handled, HandlerError>>;
127}
128
129/// A wrapper struct that adapts an `AsyncHandler` to the `Handler` trait.
130pub struct SimpleHandler<T>(pub T);
131
132#[async_trait]
133impl<T: AsyncHandler> Handler for SimpleHandler<T> {
134    async fn handle(&self, msg: CanonicalMessage) -> Result<Handled, HandlerError> {
135        self.0.handle(msg).await
136    }
137}
138
139/// A closure that can be called to commit the message.
140/// It returns a `BoxFuture` to allow for async commit operations.
141pub type CommitFunc =
142    Box<dyn FnOnce(MessageDisposition) -> BoxFuture<'static, anyhow::Result<()>> + Send + 'static>;
143
144/// A closure for committing a batch of messages.
145pub type BatchCommitFunc = Box<
146    dyn FnOnce(Vec<MessageDisposition>) -> BoxFuture<'static, anyhow::Result<()>> + Send + 'static,
147>;
148
149/// Status information about an endpoint (Consumer or Publisher).
150#[derive(Debug, Clone, serde::Serialize)]
151pub struct EndpointStatus {
152    pub healthy: bool,
153    pub target: String,
154    #[serde(skip_serializing_if = "Option::is_none")]
155    pub pending: Option<usize>,
156    #[serde(skip_serializing_if = "Option::is_none")]
157    pub capacity: Option<usize>,
158    #[serde(skip_serializing_if = "Option::is_none")]
159    pub error: Option<String>,
160    pub details: serde_json::Value,
161}
162impl Default for EndpointStatus {
163    fn default() -> Self {
164        Self {
165            healthy: true,
166            target: String::new(),
167            pending: None,
168            capacity: None,
169            error: None,
170            details: serde_json::Value::Null,
171        }
172    }
173}
174
175#[async_trait]
176pub trait MessageConsumer: Send + Sync {
177    /// Returns an optional lifecycle hook that runs once after the consumer connection is created.
178    ///
179    /// The route awaits this hook before it reports itself as ready. Returning an error fails
180    /// route startup and lets the outer route runner reconnect or surface the startup failure.
181    ///
182    /// Use this for per-connection setup that should be shared by all messages read through this
183    /// consumer, such as warming a connection pool, creating SQLite tables or indexes, setting up
184    /// a Kafka consumer group, or authenticating a RabbitMQ channel.
185    ///
186    /// ```ignore
187    /// fn on_connect_hook(&self) -> Option<BoxFuture<'_, anyhow::Result<()>>> {
188    ///     Some(Box::pin(async move {
189    ///         self.pool.get().await?;
190    ///         self.db.execute("CREATE TABLE IF NOT EXISTS embeddings (...)").await?;
191    ///         Ok(())
192    ///     }))
193    /// }
194    /// ```
195    fn on_connect_hook(&self) -> Option<BoxFuture<'_, anyhow::Result<()>>> {
196        None
197    }
198
199    /// Returns an optional lifecycle hook that runs before the consumer is dropped.
200    ///
201    /// The route awaits this hook during shutdown or reconnect cleanup. Errors are logged as
202    /// warnings and do not replace the route's original result.
203    fn on_disconnect_hook(&self) -> Option<BoxFuture<'_, anyhow::Result<()>>> {
204        None
205    }
206
207    /// Receives a batch of messages.
208    ///
209    /// This method must be implemented by all consumers.
210    /// If in doubt, implement `receive_batch` to return a single message as a vector.
211    async fn receive_batch(&mut self, _max_messages: usize)
212        -> Result<ReceivedBatch, ConsumerError>;
213
214    /// Receives a single message.
215    async fn receive(&mut self) -> Result<Received, ConsumerError> {
216        // This default implementation ensures we get exactly one message,
217        // looping if the underlying batch consumer returns an empty batch.
218        loop {
219            let mut batch = self.receive_batch(1).await?;
220            if let Some(msg) = batch.messages.pop() {
221                debug_assert!(batch.messages.is_empty());
222                if !batch.messages.is_empty() {
223                    tracing::error!(
224                        "receive_batch(1) returned {} extra messages; dropping them (implementation bug)",
225                        batch.messages.len()
226                    );
227                }
228                return Ok(Received {
229                    message: msg,
230                    commit: into_commit_func(batch.commit),
231                });
232            }
233            // Batch was success but empty, which is unexpected for receive(1). Loop.
234            tokio::time::sleep(std::time::Duration::from_millis(1)).await;
235            tokio::task::yield_now().await;
236        }
237    }
238
239    async fn receive_batch_helper(
240        &mut self,
241        _max_messages: usize,
242    ) -> Result<ReceivedBatch, ConsumerError> {
243        let received = self.receive().await?; // The `?` now correctly handles ConsumerError
244        let batch_commit = Box::new(move |dispositions: Vec<MessageDisposition>| {
245            // The default implementation only handles one message, so we take the first disposition.
246            let single_disposition = dispositions
247                .into_iter()
248                .next()
249                .unwrap_or(MessageDisposition::Ack);
250            (received.commit)(single_disposition)
251        }) as BatchCommitFunc;
252        Ok(ReceivedBatch {
253            messages: vec![received.message],
254            commit: batch_commit,
255        })
256    }
257
258    async fn status(&self) -> EndpointStatus {
259        EndpointStatus {
260            healthy: true,
261            ..Default::default()
262        }
263    }
264    fn as_any(&self) -> &dyn Any;
265}
266
267#[async_trait]
268pub trait MessagePublisher: Send + Sync + 'static {
269    /// Returns an optional lifecycle hook that runs once after the publisher connection is created.
270    ///
271    /// The route awaits this hook before it reports itself as ready. Returning an error fails
272    /// route startup and lets the outer route runner reconnect or surface the startup failure.
273    ///
274    /// Use this for per-connection setup that should be shared by all messages published through
275    /// this publisher, such as loading an embedding model, warming a connection pool, creating
276    /// SQLite tables or indexes, setting up a Kafka producer transaction context, or
277    /// authenticating a RabbitMQ channel.
278    ///
279    /// ```ignore
280    /// struct SqliteEmbeddingPublisher {
281    ///     model: Arc<tokio::sync::Mutex<Option<EmbeddingModel>>>,
282    ///     db: sqlx::SqlitePool,
283    /// }
284    ///
285    /// impl MessagePublisher for SqliteEmbeddingPublisher {
286    ///     fn on_connect_hook(&self) -> Option<BoxFuture<'_, anyhow::Result<()>>> {
287    ///         Some(Box::pin(async move {
288    ///             let mut model = self.model.lock().await;
289    ///             if model.is_none() {
290    ///                 *model = Some(EmbeddingModel::load("all-MiniLM-L6-v2").await?);
291    ///             }
292    ///             sqlx::query("CREATE INDEX IF NOT EXISTS idx_embeddings_id ON embeddings(id)")
293    ///                 .execute(&self.db)
294    ///                 .await?;
295    ///             Ok(())
296    ///         }))
297    ///     }
298    /// }
299    /// ```
300    fn on_connect_hook(&self) -> Option<BoxFuture<'_, anyhow::Result<()>>> {
301        None
302    }
303
304    /// Returns an optional lifecycle hook that runs before the publisher is dropped.
305    ///
306    /// The route awaits this hook during shutdown or reconnect cleanup. Errors are logged as
307    /// warnings and do not replace the route's original result.
308    fn on_disconnect_hook(&self) -> Option<BoxFuture<'_, anyhow::Result<()>>> {
309        None
310    }
311
312    /// Sends a batch of messages.
313    ///
314    /// This method must be implemented by all publishers.
315    /// If in doubt, implement `send_batch` to send messages one at a time.
316    async fn send_batch(
317        &self,
318        messages: Vec<CanonicalMessage>,
319    ) -> Result<SentBatch, PublisherError>;
320
321    async fn send(&self, message: CanonicalMessage) -> Result<Sent, PublisherError> {
322        let message_id = message.message_id;
323        let expects_reply = message.metadata.contains_key("reply_to");
324        match self.send_batch(vec![message]).await {
325            Ok(SentBatch::Ack) => {
326                if expects_reply {
327                    warn!("Message {:032x} expected a reply (reply_to set), but publisher returned Ack. Response loop might be broken.", message_id);
328                }
329                Ok(Sent::Ack)
330            }
331            Ok(SentBatch::Partial {
332                mut responses,
333                mut failed,
334            }) => {
335                if let Some((_, err)) = failed.pop() {
336                    Err(err)
337                } else if let Some(res) = responses.as_mut().and_then(|r| r.pop()) {
338                    Ok(Sent::Response(res))
339                } else {
340                    if expects_reply {
341                        warn!("Message {:032x} expected a reply (reply_to set), but publisher returned Ack. Response loop might be broken.", message_id);
342                    }
343                    Ok(Sent::Ack)
344                }
345            }
346            Err(e) => Err(e),
347        }
348    }
349
350    async fn flush(&self) -> anyhow::Result<()> {
351        Ok(())
352    }
353
354    async fn status(&self) -> EndpointStatus {
355        EndpointStatus {
356            healthy: true,
357            ..Default::default()
358        }
359    }
360    fn as_any(&self) -> &dyn Any;
361}
362
363#[async_trait]
364impl<T: MessagePublisher + ?Sized> MessagePublisher for Arc<T> {
365    fn on_connect_hook(&self) -> Option<BoxFuture<'_, anyhow::Result<()>>> {
366        (**self).on_connect_hook()
367    }
368
369    fn on_disconnect_hook(&self) -> Option<BoxFuture<'_, anyhow::Result<()>>> {
370        (**self).on_disconnect_hook()
371    }
372
373    async fn send(&self, message: CanonicalMessage) -> Result<Sent, PublisherError> {
374        (**self).send(message).await
375    }
376
377    async fn send_batch(
378        &self,
379        messages: Vec<CanonicalMessage>,
380    ) -> Result<SentBatch, PublisherError> {
381        (**self).send_batch(messages).await
382    }
383
384    async fn flush(&self) -> anyhow::Result<()> {
385        (**self).flush().await
386    }
387
388    async fn status(&self) -> EndpointStatus {
389        (**self).status().await
390    }
391
392    fn as_any(&self) -> &dyn Any {
393        (**self).as_any()
394    }
395}
396
397#[async_trait]
398impl<T: MessagePublisher + ?Sized> MessagePublisher for Box<T> {
399    fn on_connect_hook(&self) -> Option<BoxFuture<'_, anyhow::Result<()>>> {
400        (**self).on_connect_hook()
401    }
402
403    fn on_disconnect_hook(&self) -> Option<BoxFuture<'_, anyhow::Result<()>>> {
404        (**self).on_disconnect_hook()
405    }
406
407    async fn send(&self, message: CanonicalMessage) -> Result<Sent, PublisherError> {
408        (**self).send(message).await
409    }
410
411    async fn send_batch(
412        &self,
413        messages: Vec<CanonicalMessage>,
414    ) -> Result<SentBatch, PublisherError> {
415        (**self).send_batch(messages).await
416    }
417
418    async fn flush(&self) -> anyhow::Result<()> {
419        (**self).flush().await
420    }
421
422    async fn status(&self) -> EndpointStatus {
423        (**self).status().await
424    }
425
426    fn as_any(&self) -> &dyn Any {
427        (**self).as_any()
428    }
429}
430
431/// Factory for creating custom endpoints (consumers and publishers).
432#[async_trait]
433pub trait CustomEndpointFactory: Send + Sync + std::fmt::Debug {
434    async fn create_consumer(
435        &self,
436        _route_name: &str,
437        _config: &serde_json::Value,
438    ) -> anyhow::Result<Box<dyn MessageConsumer>> {
439        Err(anyhow::anyhow!(
440            "This custom endpoint does not support creating consumers"
441        ))
442    }
443    async fn create_publisher(
444        &self,
445        _route_name: &str,
446        _config: &serde_json::Value,
447    ) -> anyhow::Result<Box<dyn MessagePublisher>> {
448        Err(anyhow::anyhow!(
449            "This custom endpoint does not support creating publishers"
450        ))
451    }
452}
453
454/// Factory for creating custom middleware.
455#[async_trait]
456pub trait CustomMiddlewareFactory: Send + Sync + std::fmt::Debug {
457    async fn apply_consumer(
458        &self,
459        consumer: Box<dyn MessageConsumer>,
460        _route_name: &str,
461        _config: &serde_json::Value,
462    ) -> anyhow::Result<Box<dyn MessageConsumer>> {
463        Ok(consumer)
464    }
465
466    async fn apply_publisher(
467        &self,
468        publisher: Box<dyn MessagePublisher>,
469        _route_name: &str,
470        _config: &serde_json::Value,
471    ) -> anyhow::Result<Box<dyn MessagePublisher>> {
472        Ok(publisher)
473    }
474}
475
476/// A helper function to send messages in bulk by calling `send` for each one.
477/// This is useful for `MessagePublisher` implementations that don't have a native bulk sending mechanism.
478/// Requires that "send" is implemented for the publisher. Otherwise causes an infinite loop,
479/// as send is calling "send_batch" by default.
480pub async fn send_batch_helper<P: MessagePublisher + ?Sized>(
481    publisher: &P,
482    messages: Vec<CanonicalMessage>,
483    callback: impl for<'a> Fn(&'a P, CanonicalMessage) -> BoxFuture<'a, Result<Sent, PublisherError>>
484        + Send
485        + Sync,
486) -> Result<SentBatch, PublisherError> {
487    let mut responses = Vec::new();
488    let mut failed_messages = Vec::new();
489
490    let mut iter = messages.into_iter();
491    while let Some(msg) = iter.next() {
492        match callback(publisher, msg.clone()).await {
493            Ok(Sent::Response(resp)) => responses.push(resp),
494            Ok(Sent::Ack) => {}
495            Err(PublisherError::Retryable(e)) => {
496                // A retryable error likely affects the whole connection.
497                // We must return what succeeded so far (responses) and mark the rest as failed.
498                failed_messages.push((msg, PublisherError::Retryable(e)));
499                for m in iter {
500                    failed_messages.push((
501                        m,
502                        PublisherError::Retryable(anyhow::anyhow!(
503                            "Batch aborted due to previous error"
504                        )),
505                    ));
506                }
507                break;
508            }
509            Err(PublisherError::Connection(e)) => {
510                // Treat connection errors as affecting the whole batch, propagate immediately.
511                failed_messages.push((msg, PublisherError::Connection(e)));
512                for m in iter {
513                    failed_messages.push((
514                        m,
515                        PublisherError::Connection(anyhow::anyhow!(
516                            "Batch aborted due to previous connection error"
517                        )),
518                    ));
519                }
520                break;
521            }
522            Err(PublisherError::NonRetryable(e)) => {
523                // A non-retryable error is specific to this message.
524                // Collect it and continue with the rest of the batch.
525                failed_messages.push((msg, PublisherError::NonRetryable(e)));
526            }
527        }
528    }
529
530    if failed_messages.is_empty() && responses.is_empty() {
531        Ok(SentBatch::Ack)
532    } else {
533        Ok(SentBatch::Partial {
534            responses: if responses.is_empty() {
535                None
536            } else {
537                Some(responses)
538            },
539            failed: failed_messages,
540        })
541    }
542}
543
544/// Converts a `BatchCommitFunc` into a `CommitFunc` by wrapping it.
545/// This allows a function that commits a batch of messages to be used where a
546/// function that commits a single message is expected.
547pub fn into_commit_func(batch_commit: BatchCommitFunc) -> CommitFunc {
548    Box::new(move |disposition: MessageDisposition| {
549        let batch_disposition = vec![disposition];
550        batch_commit(batch_disposition)
551    })
552}
553
554/// Converts a `CommitFunc` into a `BatchCommitFunc` by wrapping it.
555/// This allows a function that commits a single message to be used where a
556/// function that commits a batch of messages is expected. It does so by
557/// extracting the first message from the response vector (if any) and passing
558/// it to the underlying single-message commit function.
559pub fn into_batch_commit_func(commit: CommitFunc) -> BatchCommitFunc {
560    Box::new(move |mut dispositions: Vec<MessageDisposition>| {
561        let single_disposition = if dispositions.len() > 1 {
562            warn!(
563                "into_batch_commit_func called with batch of {} messages; dropping all responses to avoid partial commit (incorrect usage)",
564                dispositions.len()
565            );
566            // Default to Ack to avoid hanging if we can't process the batch correctly
567            MessageDisposition::Ack
568        } else {
569            dispositions.pop().unwrap_or(MessageDisposition::Ack)
570        };
571        commit(single_disposition)
572    })
573}
574
575#[cfg(test)]
576mod tests {
577    use super::*;
578    use crate::CanonicalMessage;
579    use anyhow::anyhow;
580
581    struct MockPublisher;
582    #[async_trait]
583    impl MessagePublisher for MockPublisher {
584        async fn send_batch(
585            &self,
586            _msgs: Vec<CanonicalMessage>,
587        ) -> Result<SentBatch, PublisherError> {
588            Ok(SentBatch::Ack)
589        }
590        fn as_any(&self) -> &dyn Any {
591            self
592        }
593    }
594
595    #[tokio::test]
596    async fn test_send_batch_helper_partial_failure() {
597        let publisher = MockPublisher;
598        let msgs = vec![
599            CanonicalMessage::from("1"),
600            CanonicalMessage::from("2"),
601            CanonicalMessage::from("3"),
602        ];
603
604        let result = send_batch_helper(&publisher, msgs.clone(), |_pub, msg| {
605            Box::pin(async move {
606                let payload = msg.get_payload_str();
607                if payload == "1" {
608                    Ok(Sent::Response(CanonicalMessage::from("resp1")))
609                } else if payload == "2" {
610                    Err(PublisherError::Retryable(anyhow!("fail")))
611                } else {
612                    Ok(Sent::Ack)
613                }
614            })
615        })
616        .await;
617
618        match result {
619            Ok(SentBatch::Partial { responses, failed }) => {
620                // 1. Verify response from first message
621                assert!(responses.is_some());
622                let resps = responses.unwrap();
623                assert_eq!(resps.len(), 1);
624                assert_eq!(resps[0].get_payload_str(), "resp1");
625
626                // 2. Verify failures
627                // Message 2 failed explicitly
628                // Message 3 failed implicitly because batch was aborted
629                assert_eq!(failed.len(), 2);
630                assert_eq!(failed[0].0.get_payload_str(), "2");
631                assert!(matches!(failed[0].1, PublisherError::Retryable(_)));
632
633                assert_eq!(failed[1].0.get_payload_str(), "3");
634                assert!(matches!(failed[1].1, PublisherError::Retryable(_)));
635            }
636            _ => panic!("Expected Partial result"),
637        }
638    }
639
640    #[tokio::test]
641    async fn test_send_propagates_single_error() {
642        struct FailPublisher;
643        #[async_trait]
644        impl MessagePublisher for FailPublisher {
645            async fn send_batch(
646                &self,
647                msgs: Vec<CanonicalMessage>,
648            ) -> Result<SentBatch, PublisherError> {
649                // Simulate what send_batch_helper does on single failure
650                Ok(SentBatch::Partial {
651                    responses: None,
652                    failed: vec![(
653                        msgs[0].clone(),
654                        PublisherError::NonRetryable(anyhow!("inner")),
655                    )],
656                })
657            }
658            fn as_any(&self) -> &dyn Any {
659                self
660            }
661        }
662
663        let publ = FailPublisher;
664        let res = publ.send(CanonicalMessage::from("test")).await;
665
666        assert!(res.is_err());
667        match res.unwrap_err() {
668            PublisherError::NonRetryable(e) => assert_eq!(e.to_string(), "inner"),
669            _ => panic!("Expected NonRetryable error"),
670        }
671    }
672
673    #[tokio::test]
674    async fn test_simple_handler_wrapper() {
675        struct MyLogic;
676        impl AsyncHandler for MyLogic {
677            fn handle<'a>(
678                &'a self,
679                _msg: CanonicalMessage,
680            ) -> BoxFuture<'a, Result<Handled, HandlerError>> {
681                Box::pin(async { Ok(Handled::Ack) })
682            }
683        }
684
685        let handler = SimpleHandler(MyLogic);
686        let res = handler.handle(CanonicalMessage::from("test")).await;
687        assert!(matches!(res, Ok(Handled::Ack)));
688    }
689}