1pub 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#[derive(Default, Debug, Clone)]
21#[allow(clippy::large_enum_variant)]
22pub enum MessageDisposition {
23 #[default]
25 Ack,
26 Reply(CanonicalMessage),
28 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#[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 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
121pub trait AsyncHandler: Send + Sync + 'static {
126 fn handle<'a>(&'a self, msg: CanonicalMessage) -> BoxFuture<'a, Result<Handled, HandlerError>>;
127}
128
129pub 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
139pub type CommitFunc =
142 Box<dyn FnOnce(MessageDisposition) -> BoxFuture<'static, anyhow::Result<()>> + Send + 'static>;
143
144pub type BatchCommitFunc = Box<
146 dyn FnOnce(Vec<MessageDisposition>) -> BoxFuture<'static, anyhow::Result<()>> + Send + 'static,
147>;
148
149#[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 fn on_connect_hook(&self) -> Option<BoxFuture<'_, anyhow::Result<()>>> {
196 None
197 }
198
199 fn on_disconnect_hook(&self) -> Option<BoxFuture<'_, anyhow::Result<()>>> {
204 None
205 }
206
207 async fn receive_batch(&mut self, _max_messages: usize)
212 -> Result<ReceivedBatch, ConsumerError>;
213
214 async fn receive(&mut self) -> Result<Received, ConsumerError> {
216 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 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?; let batch_commit = Box::new(move |dispositions: Vec<MessageDisposition>| {
245 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 fn on_connect_hook(&self) -> Option<BoxFuture<'_, anyhow::Result<()>>> {
301 None
302 }
303
304 fn on_disconnect_hook(&self) -> Option<BoxFuture<'_, anyhow::Result<()>>> {
309 None
310 }
311
312 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#[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#[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
476pub 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 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 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 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
544pub 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
554pub 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 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 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 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 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}