1use crate::error::MqError;
2use async_trait::async_trait;
3use serde::{Deserialize, Serialize};
4use std::collections::{HashMap, VecDeque};
5use std::sync::Arc;
6use std::time::{Duration, Instant};
7use tokio::sync::{Notify, RwLock};
8
9#[async_trait]
18pub trait MessageQueue: Send + Sync {
19 async fn publish(&self, topic: &str, message: &[u8]) -> Result<(), MqError>;
21
22 async fn consume(&self, topic: &str) -> Result<Option<Message>, MqError>;
24
25 async fn ack(&self, message_id: &str) -> Result<(), MqError>;
27
28 async fn subscribe(&self, topic: &str) -> Result<(), MqError>;
30
31 async fn nack(&self, _message_id: &str) -> Result<(), MqError> {
38 Err(MqError::NotSupported("nack not supported".to_string()))
39 }
40
41 async fn reject(&self, _message_id: &str) -> Result<(), MqError> {
45 Err(MqError::NotSupported("reject not supported".to_string()))
46 }
47}
48
49#[derive(Debug, Clone, Serialize, Deserialize)]
57pub struct Message {
58 pub topic: String,
59 pub payload: Vec<u8>,
60 pub key: Option<String>,
61 pub timestamp: i64,
62 pub headers: HashMap<String, String>,
63 #[serde(default)]
64 pub id: String,
65 #[serde(default)]
67 pub retry_count: u32,
68}
69
70impl Message {
71 pub fn new(topic: impl Into<String>, payload: Vec<u8>) -> Self {
72 Self {
73 topic: topic.into(),
74 payload,
75 key: None,
76 timestamp: current_timestamp(),
77 headers: HashMap::new(),
78 id: String::new(),
79 retry_count: 0,
80 }
81 }
82
83 pub fn with_key(mut self, key: impl Into<String>) -> Self {
84 self.key = Some(key.into());
85 self
86 }
87
88 pub fn text(&self) -> Option<&str> {
89 std::str::from_utf8(&self.payload).ok()
90 }
91
92 pub fn json<T: serde::de::DeserializeOwned>(&self) -> Option<T> {
93 serde_json::from_slice(&self.payload).ok()
94 }
95
96 pub fn text_message(topic: impl Into<String>, text: impl Into<String>) -> Self {
97 Self::new(topic, text.into().into_bytes())
98 }
99
100 pub fn json_message<T: serde::Serialize>(
101 topic: impl Into<String>,
102 data: &T,
103 ) -> Result<Self, MqError> {
104 let payload = serde_json::to_vec(data)?;
105 Ok(Self::new(topic, payload))
106 }
107}
108
109fn current_timestamp() -> i64 {
115 use std::time::{SystemTime, UNIX_EPOCH};
116 match SystemTime::now().duration_since(UNIX_EPOCH) {
117 Ok(d) => d.as_millis() as i64,
118 Err(e) => {
119 eprintln!(
122 "WARN: current_timestamp: system time before UNIX_EPOCH: {} (duration_secs={})",
123 e,
124 e.duration().as_secs()
125 );
126 0
127 }
128 }
129}
130
131pub struct QueueConfig {
136 pub provider: MqProvider,
137 pub brokers: Vec<String>,
138 pub group_id: Option<String>,
139 pub username: Option<String>,
140 pub password: Option<String>,
141}
142
143impl Default for QueueConfig {
144 fn default() -> Self {
145 Self {
146 provider: MqProvider::Kafka(KafkaConfig::default()),
147 brokers: vec!["localhost:9092".to_string()],
148 group_id: None,
149 username: None,
150 password: None,
151 }
152 }
153}
154
155impl QueueConfig {
156 pub fn new() -> Self {
157 Self::default()
158 }
159
160 pub fn with_provider(mut self, provider: MqProvider) -> Self {
161 self.provider = provider;
162 self
163 }
164
165 pub fn with_brokers(mut self, brokers: Vec<String>) -> Self {
166 self.brokers = brokers;
167 self
168 }
169
170 pub fn with_group(mut self, group: impl Into<String>) -> Self {
171 self.group_id = Some(group.into());
172 self
173 }
174
175 pub fn with_auth(mut self, username: impl Into<String>, password: impl Into<String>) -> Self {
176 self.username = Some(username.into());
177 self.password = Some(password.into());
178 self
179 }
180}
181
182#[derive(Debug, Clone)]
183pub enum MqProvider {
184 Kafka(KafkaConfig),
185 RabbitMQ(RabbitConfig),
186 RocketMQ(RocketConfig),
187 ActiveMQ(ActiveConfig),
188 Nats(NatsConfig),
189 Pulsar(PulsarConfig),
190}
191
192#[derive(Debug, Clone, Default)]
193pub struct KafkaConfig {
194 pub client_id: Option<String>,
195 pub acks: Option<String>,
196 pub retries: Option<u32>,
197}
198
199#[derive(Debug, Clone, Default)]
200pub struct RabbitConfig {
201 pub virtual_host: Option<String>,
202}
203
204#[derive(Debug, Clone, Default)]
205pub struct RocketConfig {
206 pub namespace: Option<String>,
207}
208
209#[derive(Debug, Clone, Default)]
210pub struct ActiveConfig {
211 pub broker_url: Option<String>,
212}
213
214#[derive(Debug, Clone, Default)]
215pub struct NatsConfig {
216 pub name: Option<String>,
217}
218
219#[derive(Debug, Clone, Default)]
220pub struct PulsarConfig {
221 pub service_url: Option<String>,
222}
223
224#[derive(Debug, Clone)]
237pub struct ReconnectPolicy {
238 pub max_retries: u32,
240 pub initial_delay_ms: u64,
242 pub max_delay_ms: u64,
244 pub multiplier: f64,
246}
247
248impl Default for ReconnectPolicy {
249 fn default() -> Self {
250 Self {
251 max_retries: 5,
252 initial_delay_ms: 100,
253 max_delay_ms: 10_000,
254 multiplier: 2.0,
255 }
256 }
257}
258
259impl ReconnectPolicy {
260 pub fn new() -> Self {
261 Self::default()
262 }
263
264 pub fn next_delay(&self, attempt: u32) -> Duration {
272 let delay_ms = (self.initial_delay_ms as f64) * self.multiplier.powi(attempt as i32);
273 let delay_ms = delay_ms.max(0.0).min(self.max_delay_ms as f64);
275 Duration::from_millis(delay_ms as u64)
276 }
277}
278
279#[derive(Debug, Clone, Default)]
283pub struct ReconnectState {
284 pub attempts: u32,
286 pub last_reconnect: Option<Instant>,
288}
289
290#[derive(Debug, Clone)]
301pub struct BackpressurePolicy {
302 pub max_queue_size: usize,
304 pub on_overflow: OverflowStrategy,
306}
307
308impl Default for BackpressurePolicy {
309 fn default() -> Self {
310 Self {
311 max_queue_size: 10_000,
312 on_overflow: OverflowStrategy::Reject,
313 }
314 }
315}
316
317#[derive(Debug, Clone, Copy, PartialEq, Eq)]
324pub enum OverflowStrategy {
325 Block,
327 DropOldest,
329 DropNewest,
331 Reject,
333}
334
335pub struct InMemoryQueue {
340 inner: Arc<RwLock<InMemoryQueueInner>>,
341}
342
343impl Clone for InMemoryQueue {
344 fn clone(&self) -> Self {
345 Self {
346 inner: Arc::clone(&self.inner),
347 }
348 }
349}
350
351struct InMemoryQueueInner {
352 queues: HashMap<String, VecDeque<Message>>,
354 in_flight: HashMap<String, Message>,
356 subscribers: HashMap<String, usize>,
358 next_id: u64,
360 max_messages_per_topic: usize,
363 dead_letters: HashMap<String, VecDeque<Message>>,
365 max_retries: u32,
367 backpressure: Option<BackpressurePolicy>,
369 notify: HashMap<String, Arc<Notify>>,
371}
372
373const DEFAULT_MAX_MESSAGES_PER_TOPIC: usize = 100_000;
375
376const DEFAULT_MAX_RETRIES: u32 = 3;
378
379impl InMemoryQueue {
380 pub fn new() -> Self {
381 Self::with_max_messages_per_topic(DEFAULT_MAX_MESSAGES_PER_TOPIC)
382 }
383
384 pub fn with_max_messages_per_topic(max: usize) -> Self {
389 Self {
390 inner: Arc::new(RwLock::new(InMemoryQueueInner {
391 queues: HashMap::new(),
392 in_flight: HashMap::new(),
393 subscribers: HashMap::new(),
394 next_id: 1,
395 max_messages_per_topic: max,
396 dead_letters: HashMap::new(),
397 max_retries: DEFAULT_MAX_RETRIES,
398 backpressure: None,
399 notify: HashMap::new(),
400 })),
401 }
402 }
403
404 pub fn with_max_retries(max_retries: u32) -> Self {
409 Self {
410 inner: Arc::new(RwLock::new(InMemoryQueueInner {
411 queues: HashMap::new(),
412 in_flight: HashMap::new(),
413 subscribers: HashMap::new(),
414 next_id: 1,
415 max_messages_per_topic: DEFAULT_MAX_MESSAGES_PER_TOPIC,
416 dead_letters: HashMap::new(),
417 max_retries,
418 backpressure: None,
419 notify: HashMap::new(),
420 })),
421 }
422 }
423
424 pub fn with_backpressure(policy: BackpressurePolicy) -> Self {
429 let max = policy.max_queue_size;
430 Self {
431 inner: Arc::new(RwLock::new(InMemoryQueueInner {
432 queues: HashMap::new(),
433 in_flight: HashMap::new(),
434 subscribers: HashMap::new(),
435 next_id: 1,
436 max_messages_per_topic: max,
437 dead_letters: HashMap::new(),
438 max_retries: DEFAULT_MAX_RETRIES,
439 backpressure: Some(policy),
440 notify: HashMap::new(),
441 })),
442 }
443 }
444
445 pub async fn message_count(&self, topic: &str) -> usize {
446 let inner = self.inner.read().await;
447 inner.queues.get(topic).map(|q| q.len()).unwrap_or(0)
448 }
449
450 pub async fn subscriber_count(&self, topic: &str) -> usize {
451 let inner = self.inner.read().await;
452 *inner.subscribers.get(topic).unwrap_or(&0)
453 }
454
455 pub async fn in_flight_count(&self) -> usize {
456 let inner = self.inner.read().await;
457 inner.in_flight.len()
458 }
459
460 pub async fn dead_letter_count(&self, topic: &str) -> usize {
464 let inner = self.inner.read().await;
465 inner.dead_letters.get(topic).map(|q| q.len()).unwrap_or(0)
466 }
467
468 pub async fn consume_dead_letter(&self, topic: &str) -> Option<Message> {
473 let mut inner = self.inner.write().await;
474 if let Some(dq) = inner.dead_letters.get_mut(topic) {
475 return dq.pop_front();
476 }
477 None
478 }
479
480 pub async fn requeue_dead_letter(&self, message_id: &str) -> Result<(), MqError> {
485 let mut inner = self.inner.write().await;
486 for dq in inner.dead_letters.values_mut() {
488 let mut found_idx = None;
489 for (idx, m) in dq.iter().enumerate() {
490 if m.id == message_id {
491 found_idx = Some(idx);
492 break;
493 }
494 }
495 if let Some(idx) = found_idx {
496 let mut msg = dq.remove(idx).expect("checked: idx exists");
497 msg.retry_count = 0;
499 inner
501 .queues
502 .entry(msg.topic.clone())
503 .or_insert_with(VecDeque::new)
504 .push_back(msg);
505 return Ok(());
506 }
507 }
508 Err(MqError::NotSupported(format!(
509 "Dead letter not found: {}",
510 message_id
511 )))
512 }
513}
514
515impl Default for InMemoryQueue {
516 fn default() -> Self {
517 Self::new()
518 }
519}
520
521#[async_trait]
522impl MessageQueue for InMemoryQueue {
523 async fn publish(&self, topic: &str, message: &[u8]) -> Result<(), MqError> {
524 let strategy = {
526 let inner = self.inner.read().await;
527 inner
528 .backpressure
529 .as_ref()
530 .map(|p| p.on_overflow)
531 .unwrap_or(OverflowStrategy::Reject)
532 };
533
534 match strategy {
535 OverflowStrategy::Block => self.publish_with_block(topic, message).await,
536 _ => self.publish_immediate(topic, message, strategy).await,
537 }
538 }
539
540 async fn consume(&self, topic: &str) -> Result<Option<Message>, MqError> {
541 let mut inner = self.inner.write().await;
542 let queue = inner
543 .queues
544 .entry(topic.to_string())
545 .or_insert_with(VecDeque::new);
546 if let Some(msg) = queue.pop_front() {
547 inner.in_flight.insert(msg.id.clone(), msg.clone());
548 if let Some(notify) = inner.notify.get(topic) {
550 notify.notify_one();
551 }
552 Ok(Some(msg))
553 } else {
554 Ok(None)
555 }
556 }
557
558 async fn ack(&self, message_id: &str) -> Result<(), MqError> {
559 let mut inner = self.inner.write().await;
560 inner.in_flight.remove(message_id).ok_or_else(|| {
561 MqError::NotSupported(format!("Message not found for ack: {}", message_id))
562 })?;
563 Ok(())
564 }
565
566 async fn subscribe(&self, topic: &str) -> Result<(), MqError> {
567 let mut inner = self.inner.write().await;
568 *inner.subscribers.entry(topic.to_string()).or_insert(0) += 1;
569 Ok(())
570 }
571
572 async fn nack(&self, message_id: &str) -> Result<(), MqError> {
573 let mut inner = self.inner.write().await;
574 let mut msg = inner.in_flight.remove(message_id).ok_or_else(|| {
575 MqError::NotSupported(format!("Message not found for nack: {}", message_id))
576 })?;
577
578 msg.retry_count = msg.retry_count.saturating_add(1);
580
581 if msg.retry_count >= inner.max_retries {
583 inner
584 .dead_letters
585 .entry(msg.topic.clone())
586 .or_insert_with(VecDeque::new)
587 .push_back(msg);
588 } else {
589 inner
592 .queues
593 .entry(msg.topic.clone())
594 .or_insert_with(VecDeque::new)
595 .push_back(msg);
596 }
597 Ok(())
598 }
599
600 async fn reject(&self, message_id: &str) -> Result<(), MqError> {
601 let mut inner = self.inner.write().await;
602 let msg = inner.in_flight.remove(message_id).ok_or_else(|| {
603 MqError::NotSupported(format!("Message not found for reject: {}", message_id))
604 })?;
605 inner
607 .dead_letters
608 .entry(msg.topic.clone())
609 .or_insert_with(VecDeque::new)
610 .push_back(msg);
611 Ok(())
612 }
613}
614
615impl InMemoryQueue {
616 async fn publish_immediate(
620 &self,
621 topic: &str,
622 message: &[u8],
623 strategy: OverflowStrategy,
624 ) -> Result<(), MqError> {
625 let mut inner = self.inner.write().await;
626 let current_count = inner.queues.get(topic).map(|q| q.len()).unwrap_or(0);
628 if current_count >= inner.max_messages_per_topic {
629 match strategy {
630 OverflowStrategy::DropOldest => {
631 let queue = inner
633 .queues
634 .entry(topic.to_string())
635 .or_insert_with(VecDeque::new);
636 queue.pop_front();
637 }
638 OverflowStrategy::DropNewest => {
639 return Ok(());
641 }
642 OverflowStrategy::Reject => {
643 return Err(MqError::Publish(format!(
644 "topic '{}' is full: {} >= {} messages (H-3 protection)",
645 topic, current_count, inner.max_messages_per_topic
646 )));
647 }
648 OverflowStrategy::Block => {
649 return Err(MqError::Publish(format!(
651 "topic '{}' overflow with Block strategy: use publish_with_block instead",
652 topic
653 )));
654 }
655 }
656 }
657 let id = format!("msg-{}", inner.next_id);
659 inner.next_id = inner
661 .next_id
662 .checked_add(1)
663 .ok_or_else(|| MqError::Publish("message id overflow: u64::MAX reached".to_string()))?;
664 let msg = Message {
665 id,
666 retry_count: 0,
667 ..Message::new(topic, message.to_vec())
668 };
669 inner
670 .queues
671 .entry(topic.to_string())
672 .or_insert_with(VecDeque::new)
673 .push_back(msg);
674 Ok(())
675 }
676
677 async fn publish_with_block(&self, topic: &str, message: &[u8]) -> Result<(), MqError> {
681 loop {
682 let notify = {
683 let mut inner = self.inner.write().await;
684 let current_count = inner.queues.get(topic).map(|q| q.len()).unwrap_or(0);
685 if current_count < inner.max_messages_per_topic {
686 let id = format!("msg-{}", inner.next_id);
688 inner.next_id = inner.next_id.checked_add(1).ok_or_else(|| {
689 MqError::Publish("message id overflow: u64::MAX reached".to_string())
690 })?;
691 let msg = Message {
692 id,
693 retry_count: 0,
694 ..Message::new(topic, message.to_vec())
695 };
696 inner
697 .queues
698 .entry(topic.to_string())
699 .or_insert_with(VecDeque::new)
700 .push_back(msg);
701 return Ok(());
702 }
703 inner
705 .notify
706 .entry(topic.to_string())
707 .or_insert_with(|| Arc::new(Notify::new()))
708 .clone()
709 };
710 notify.notified().await;
713 }
714 }
715}
716
717pub struct QueueWrapper {
722 queue: Box<dyn MessageQueue>,
723 reconnect: Option<ReconnectPolicy>,
725}
726
727impl QueueWrapper {
728 pub fn new(provider: MqProvider) -> Self {
729 let queue: Box<dyn MessageQueue> = match provider {
730 MqProvider::Kafka(_) => Box::new(crate::kafka::InMemoryKafkaQueue::new()),
731 MqProvider::RabbitMQ(_) => Box::new(crate::rabbitmq::InMemoryRabbitmqQueue::new()),
732 MqProvider::RocketMQ(_) => Box::new(crate::rocketmq::InMemoryRocketmqQueue::new()),
733 MqProvider::ActiveMQ(_) => Box::new(crate::activemq::InMemoryActivemqQueue::new()),
734 MqProvider::Nats(_) => Box::new(crate::nats::InMemoryNatsQueue::new()),
735 MqProvider::Pulsar(_) => Box::new(crate::pulsar::InMemoryPulsarQueue::new()),
736 };
737 Self {
738 queue,
739 reconnect: None,
740 }
741 }
742
743 #[cfg(test)]
745 pub(crate) fn with_queue(queue: Box<dyn MessageQueue>) -> Self {
746 Self {
747 queue,
748 reconnect: None,
749 }
750 }
751
752 pub fn with_reconnect(mut self, policy: ReconnectPolicy) -> Self {
757 self.reconnect = Some(policy);
758 self
759 }
760
761 pub async fn publish(&self, topic: &str, message: &[u8]) -> Result<(), MqError> {
762 if let Some(policy) = &self.reconnect {
763 let mut attempts = 0u32;
764 loop {
765 match self.queue.publish(topic, message).await {
766 Ok(()) => return Ok(()),
767 Err(MqError::Connection(_)) if attempts < policy.max_retries => {
768 let delay = policy.next_delay(attempts);
769 tokio::time::sleep(delay).await;
770 attempts += 1;
771 }
772 Err(e) => return Err(e),
773 }
774 }
775 } else {
776 self.queue.publish(topic, message).await
777 }
778 }
779
780 pub async fn consume(&self, topic: &str) -> Result<Option<Message>, MqError> {
781 if let Some(policy) = &self.reconnect {
782 let mut attempts = 0u32;
783 loop {
784 match self.queue.consume(topic).await {
785 Ok(msg) => return Ok(msg),
786 Err(MqError::Connection(_)) if attempts < policy.max_retries => {
787 let delay = policy.next_delay(attempts);
788 tokio::time::sleep(delay).await;
789 attempts += 1;
790 }
791 Err(e) => return Err(e),
792 }
793 }
794 } else {
795 self.queue.consume(topic).await
796 }
797 }
798
799 pub async fn ack(&self, message_id: &str) -> Result<(), MqError> {
800 self.queue.ack(message_id).await
801 }
802
803 pub async fn subscribe(&self, topic: &str) -> Result<(), MqError> {
804 self.queue.subscribe(topic).await
805 }
806
807 pub async fn nack(&self, message_id: &str) -> Result<(), MqError> {
811 self.queue.nack(message_id).await
812 }
813
814 pub async fn reject(&self, message_id: &str) -> Result<(), MqError> {
818 self.queue.reject(message_id).await
819 }
820}
821
822#[cfg(test)]
823mod tests {
824 use super::*;
825 use std::sync::atomic::{AtomicU32, Ordering};
826
827 #[tokio::test]
832 async fn test_in_memory_queue_basic() {
833 let queue = InMemoryQueue::new();
834 queue.publish("topic1", b"hello").await.unwrap();
835 let msg = queue
836 .consume("topic1")
837 .await
838 .unwrap()
839 .expect("msg should exist");
840 assert_eq!(msg.payload, b"hello");
841 queue.ack(&msg.id).await.unwrap();
842 }
843
844 #[tokio::test]
846 async fn test_l2_next_id_overflow_protection() {
847 let queue = InMemoryQueue::new();
848 {
849 let mut inner = queue.inner.write().await;
850 inner.next_id = u64::MAX;
851 }
852 let result = queue.publish("topic1", b"msg").await;
853 assert!(result.is_err());
854 match result {
855 Err(MqError::Publish(msg)) => {
856 assert!(
857 msg.contains("overflow"),
858 "expected overflow error, got: {}",
859 msg
860 );
861 }
862 _ => panic!("Expected MqError::Publish with overflow message"),
863 }
864 }
865
866 #[tokio::test]
868 async fn test_l2_next_id_near_max() {
869 let queue = InMemoryQueue::new();
870 {
871 let mut inner = queue.inner.write().await;
872 inner.next_id = u64::MAX - 1;
873 }
874 let result1 = queue.publish("topic1", b"msg1").await;
875 assert!(result1.is_ok());
876 let result2 = queue.publish("topic1", b"msg2").await;
877 assert!(result2.is_err());
878 }
879
880 #[tokio::test]
886 async fn test_nack_requeues_message_with_retry_count() {
887 let queue = InMemoryQueue::new();
888 queue.publish("topic", b"msg1").await.unwrap();
889 let msg = queue.consume("topic").await.unwrap().unwrap();
890 assert_eq!(msg.retry_count, 0);
891
892 queue.nack(&msg.id).await.unwrap();
894 assert_eq!(queue.message_count("topic").await, 1);
895 assert_eq!(queue.in_flight_count().await, 0);
896
897 let msg2 = queue.consume("topic").await.unwrap().unwrap();
899 assert_eq!(msg2.id, msg.id);
900 assert_eq!(msg2.retry_count, 1);
901 }
902
903 #[tokio::test]
912 async fn test_nack_increments_retry_count() {
913 let queue = InMemoryQueue::with_max_retries(10);
914 queue.publish("topic", b"data").await.unwrap();
915
916 for expected_retry in 0..5u32 {
917 let msg = queue.consume("topic").await.unwrap().unwrap();
918 assert_eq!(
919 msg.retry_count, expected_retry,
920 "consume should show retry_count before this iteration's nack"
921 );
922 queue.nack(&msg.id).await.unwrap();
923 }
924 assert_eq!(queue.message_count("topic").await, 1);
926 assert_eq!(queue.dead_letter_count("topic").await, 0);
927 }
928
929 #[tokio::test]
931 async fn test_nack_max_retries_sends_to_dlx() {
932 let queue = InMemoryQueue::with_max_retries(3);
934 queue.publish("topic", b"payload").await.unwrap();
935
936 let msg = queue.consume("topic").await.unwrap().unwrap();
938 queue.nack(&msg.id).await.unwrap();
939 assert_eq!(queue.message_count("topic").await, 1);
940 assert_eq!(queue.dead_letter_count("topic").await, 0);
941
942 let msg = queue.consume("topic").await.unwrap().unwrap();
944 assert_eq!(msg.retry_count, 1);
945 queue.nack(&msg.id).await.unwrap();
946 assert_eq!(queue.message_count("topic").await, 1);
947 assert_eq!(queue.dead_letter_count("topic").await, 0);
948
949 let msg = queue.consume("topic").await.unwrap().unwrap();
951 assert_eq!(msg.retry_count, 2);
952 queue.nack(&msg.id).await.unwrap();
953 assert_eq!(queue.message_count("topic").await, 0);
954 assert_eq!(queue.dead_letter_count("topic").await, 1);
955 }
956
957 #[tokio::test]
959 async fn test_nack_max_retries_zero_sends_to_dlx_immediately() {
960 let queue = InMemoryQueue::with_max_retries(0);
961 queue.publish("topic", b"msg").await.unwrap();
962 let msg = queue.consume("topic").await.unwrap().unwrap();
963 assert_eq!(msg.retry_count, 0);
964
965 queue.nack(&msg.id).await.unwrap();
967
968 assert_eq!(queue.message_count("topic").await, 0);
969 assert_eq!(queue.dead_letter_count("topic").await, 1);
970
971 let dlq_msg = queue
973 .consume_dead_letter("topic")
974 .await
975 .expect("should have dead letter");
976 assert_eq!(dlq_msg.retry_count, 1);
977 }
978
979 #[tokio::test]
981 async fn test_nack_unknown_message_id_returns_error() {
982 let queue = InMemoryQueue::new();
983 let result = queue.nack("nonexistent-id").await;
984 assert!(result.is_err());
985 match result {
986 Err(MqError::NotSupported(msg)) => {
987 assert!(msg.contains("not found for nack"));
988 }
989 _ => panic!("Expected MqError::NotSupported"),
990 }
991 }
992
993 #[tokio::test]
995 async fn test_reject_sends_to_dead_letter_queue() {
996 let queue = InMemoryQueue::new();
997 queue.publish("topic", b"bad-msg").await.unwrap();
998 let msg = queue.consume("topic").await.unwrap().unwrap();
999
1000 queue.reject(&msg.id).await.unwrap();
1001
1002 assert_eq!(queue.message_count("topic").await, 0);
1003 assert_eq!(queue.in_flight_count().await, 0);
1004 assert_eq!(queue.dead_letter_count("topic").await, 1);
1005 }
1006
1007 #[tokio::test]
1009 async fn test_reject_does_not_increment_retry_count() {
1010 let queue = InMemoryQueue::new();
1011 queue.publish("topic", b"msg").await.unwrap();
1012 let msg = queue.consume("topic").await.unwrap().unwrap();
1013 assert_eq!(msg.retry_count, 0);
1014
1015 queue.reject(&msg.id).await.unwrap();
1016
1017 let dlq_msg = queue
1018 .consume_dead_letter("topic")
1019 .await
1020 .expect("should have dead letter");
1021 assert_eq!(
1022 dlq_msg.retry_count, 0,
1023 "reject should not increment retry_count"
1024 );
1025 }
1026
1027 #[tokio::test]
1029 async fn test_reject_unknown_message_id_returns_error() {
1030 let queue = InMemoryQueue::new();
1031 let result = queue.reject("nonexistent-id").await;
1032 assert!(result.is_err());
1033 }
1034
1035 #[tokio::test]
1037 async fn test_reject_empty_in_flight_returns_error() {
1038 let queue = InMemoryQueue::new();
1039 let result = queue.reject("any-id").await;
1041 assert!(result.is_err());
1042 assert_eq!(queue.dead_letter_count("topic").await, 0);
1043 }
1044
1045 #[tokio::test]
1047 async fn test_dead_letter_count_empty_topic() {
1048 let queue = InMemoryQueue::new();
1049 assert_eq!(queue.dead_letter_count("no-such-topic").await, 0);
1050 }
1051
1052 #[tokio::test]
1054 async fn test_dead_letter_count_after_reject() {
1055 let queue = InMemoryQueue::new();
1056 queue.publish("topic", b"m1").await.unwrap();
1057 queue.publish("topic", b"m2").await.unwrap();
1058
1059 let m1 = queue.consume("topic").await.unwrap().unwrap();
1060 queue.reject(&m1.id).await.unwrap();
1061 assert_eq!(queue.dead_letter_count("topic").await, 1);
1062
1063 let m2 = queue.consume("topic").await.unwrap().unwrap();
1064 queue.reject(&m2.id).await.unwrap();
1065 assert_eq!(queue.dead_letter_count("topic").await, 2);
1066 }
1067
1068 #[tokio::test]
1070 async fn test_consume_dead_letter() {
1071 let queue = InMemoryQueue::new();
1072 queue.publish("topic", b"first").await.unwrap();
1073 queue.publish("topic", b"second").await.unwrap();
1074
1075 let m1 = queue.consume("topic").await.unwrap().unwrap();
1076 queue.reject(&m1.id).await.unwrap();
1077 let m2 = queue.consume("topic").await.unwrap().unwrap();
1078 queue.reject(&m2.id).await.unwrap();
1079
1080 let d1 = queue
1082 .consume_dead_letter("topic")
1083 .await
1084 .expect("should have dead letter");
1085 assert_eq!(d1.payload, b"first");
1086 let d2 = queue
1087 .consume_dead_letter("topic")
1088 .await
1089 .expect("should have dead letter");
1090 assert_eq!(d2.payload, b"second");
1091
1092 assert!(queue.consume_dead_letter("topic").await.is_none());
1094 }
1095
1096 #[tokio::test]
1098 async fn test_consume_dead_letter_empty() {
1099 let queue = InMemoryQueue::new();
1100 assert!(queue.consume_dead_letter("no-such-topic").await.is_none());
1101 }
1102
1103 #[tokio::test]
1105 async fn test_requeue_dead_letter_resets_retry_count() {
1106 let queue = InMemoryQueue::with_max_retries(2);
1107 queue.publish("topic", b"msg").await.unwrap();
1108
1109 let m1 = queue.consume("topic").await.unwrap().unwrap();
1111 queue.nack(&m1.id).await.unwrap();
1112 let m2 = queue.consume("topic").await.unwrap().unwrap();
1114 assert_eq!(m2.retry_count, 1);
1115 queue.nack(&m2.id).await.unwrap();
1116
1117 assert_eq!(queue.dead_letter_count("topic").await, 1);
1118 assert_eq!(queue.message_count("topic").await, 0);
1119
1120 queue.requeue_dead_letter(&m1.id).await.unwrap();
1122 assert_eq!(queue.dead_letter_count("topic").await, 0);
1123 assert_eq!(queue.message_count("topic").await, 1);
1124
1125 let m3 = queue.consume("topic").await.unwrap().unwrap();
1127 assert_eq!(m3.id, m1.id);
1128 assert_eq!(
1129 m3.retry_count, 0,
1130 "retry_count should be reset after requeue"
1131 );
1132 }
1133
1134 #[tokio::test]
1136 async fn test_requeue_dead_letter_not_found() {
1137 let queue = InMemoryQueue::new();
1138 let result = queue.requeue_dead_letter("nonexistent-id").await;
1139 assert!(result.is_err());
1140 match result {
1141 Err(MqError::NotSupported(msg)) => {
1142 assert!(msg.contains("Dead letter not found"));
1143 }
1144 _ => panic!("Expected MqError::NotSupported"),
1145 }
1146 }
1147
1148 #[test]
1154 fn test_reconnect_policy_default_values() {
1155 let policy = ReconnectPolicy::default();
1156 assert_eq!(policy.max_retries, 5);
1157 assert_eq!(policy.initial_delay_ms, 100);
1158 assert_eq!(policy.max_delay_ms, 10_000);
1159 assert!((policy.multiplier - 2.0).abs() < f64::EPSILON);
1160 }
1161
1162 #[test]
1164 fn test_reconnect_policy_next_delay_exponential() {
1165 let policy = ReconnectPolicy {
1166 max_retries: 5,
1167 initial_delay_ms: 100,
1168 max_delay_ms: 10_000,
1169 multiplier: 2.0,
1170 };
1171 assert_eq!(policy.next_delay(0), Duration::from_millis(100));
1173 assert_eq!(policy.next_delay(1), Duration::from_millis(200));
1175 assert_eq!(policy.next_delay(2), Duration::from_millis(400));
1177 assert_eq!(policy.next_delay(3), Duration::from_millis(800));
1179 }
1180
1181 #[test]
1183 fn test_reconnect_policy_next_delay_capped_at_max() {
1184 let policy = ReconnectPolicy {
1185 max_retries: 10,
1186 initial_delay_ms: 100,
1187 max_delay_ms: 1000,
1188 multiplier: 2.0,
1189 };
1190 assert_eq!(policy.next_delay(4), Duration::from_millis(1000));
1192 assert_eq!(policy.next_delay(10), Duration::from_millis(1000));
1194 }
1195
1196 #[test]
1198 fn test_reconnect_policy_zero_attempt() {
1199 let policy = ReconnectPolicy {
1200 max_retries: 3,
1201 initial_delay_ms: 500,
1202 max_delay_ms: 10_000,
1203 multiplier: 3.0,
1204 };
1205 assert_eq!(policy.next_delay(0), Duration::from_millis(500));
1206 }
1207
1208 struct FailingQueue {
1216 call_count: AtomicU32,
1217 succeed_on_attempt: u32,
1218 }
1219
1220 impl FailingQueue {
1221 fn new(succeed_on_attempt: u32) -> Self {
1222 Self {
1223 call_count: AtomicU32::new(0),
1224 succeed_on_attempt,
1225 }
1226 }
1227 }
1228
1229 #[async_trait]
1230 impl MessageQueue for FailingQueue {
1231 async fn publish(&self, _topic: &str, _message: &[u8]) -> Result<(), MqError> {
1232 let attempt = self.call_count.fetch_add(1, Ordering::SeqCst) + 1;
1233 if attempt >= self.succeed_on_attempt {
1234 Ok(())
1235 } else {
1236 Err(MqError::Connection(
1237 "simulated connection error".to_string(),
1238 ))
1239 }
1240 }
1241
1242 async fn consume(&self, _topic: &str) -> Result<Option<Message>, MqError> {
1243 Err(MqError::Connection(
1244 "simulated connection error".to_string(),
1245 ))
1246 }
1247
1248 async fn ack(&self, _message_id: &str) -> Result<(), MqError> {
1249 Err(MqError::Connection("simulated".to_string()))
1250 }
1251
1252 async fn subscribe(&self, _topic: &str) -> Result<(), MqError> {
1253 Err(MqError::Connection("simulated".to_string()))
1254 }
1255 }
1256
1257 struct PublishErrorQueue;
1259
1260 #[async_trait]
1261 impl MessageQueue for PublishErrorQueue {
1262 async fn publish(&self, _topic: &str, _message: &[u8]) -> Result<(), MqError> {
1263 Err(MqError::Publish("non-connection error".to_string()))
1264 }
1265
1266 async fn consume(&self, _topic: &str) -> Result<Option<Message>, MqError> {
1267 Err(MqError::Publish("non-connection error".to_string()))
1268 }
1269
1270 async fn ack(&self, _message_id: &str) -> Result<(), MqError> {
1271 Ok(())
1272 }
1273
1274 async fn subscribe(&self, _topic: &str) -> Result<(), MqError> {
1275 Ok(())
1276 }
1277 }
1278
1279 #[tokio::test]
1281 async fn test_reconnect_retries_on_connection_error() {
1282 let failing = FailingQueue::new(3);
1284 let wrapper = QueueWrapper::with_queue(Box::new(failing)).with_reconnect(ReconnectPolicy {
1285 max_retries: 5,
1286 initial_delay_ms: 1, max_delay_ms: 10,
1288 multiplier: 2.0,
1289 });
1290
1291 let result = wrapper.publish("topic", b"data").await;
1292 assert!(result.is_ok(), "should succeed after retries");
1293 }
1294
1295 #[tokio::test]
1297 async fn test_reconnect_gives_up_after_max_retries() {
1298 let failing = FailingQueue::new(u32::MAX);
1300 let wrapper = QueueWrapper::with_queue(Box::new(failing)).with_reconnect(ReconnectPolicy {
1301 max_retries: 2,
1302 initial_delay_ms: 1,
1303 max_delay_ms: 10,
1304 multiplier: 2.0,
1305 });
1306
1307 let result = wrapper.publish("topic", b"data").await;
1308 assert!(result.is_err());
1309 match result {
1310 Err(MqError::Connection(_)) => {}
1311 _ => panic!("Expected MqError::Connection"),
1312 }
1313 }
1314
1315 #[tokio::test]
1317 async fn test_reconnect_no_retry_on_non_connection_error() {
1318 let wrapper =
1319 QueueWrapper::with_queue(Box::new(PublishErrorQueue)).with_reconnect(ReconnectPolicy {
1320 max_retries: 5,
1321 initial_delay_ms: 1,
1322 max_delay_ms: 10,
1323 multiplier: 2.0,
1324 });
1325
1326 let result = wrapper.publish("topic", b"data").await;
1327 assert!(result.is_err());
1328 match result {
1329 Err(MqError::Publish(msg)) => {
1330 assert!(msg.contains("non-connection error"));
1331 }
1332 _ => panic!("Expected MqError::Publish"),
1333 }
1334 }
1335
1336 #[tokio::test]
1342 async fn test_backpressure_drop_oldest() {
1343 let queue = InMemoryQueue::with_backpressure(BackpressurePolicy {
1344 max_queue_size: 2,
1345 on_overflow: OverflowStrategy::DropOldest,
1346 });
1347
1348 queue.publish("topic", b"m1").await.unwrap();
1349 queue.publish("topic", b"m2").await.unwrap();
1350 queue.publish("topic", b"m3").await.unwrap();
1352
1353 assert_eq!(queue.message_count("topic").await, 2);
1354
1355 let m1 = queue.consume("topic").await.unwrap().unwrap();
1357 assert_eq!(m1.payload, b"m2", "oldest should be dropped");
1358 let m2 = queue.consume("topic").await.unwrap().unwrap();
1359 assert_eq!(m2.payload, b"m3");
1360 }
1361
1362 #[tokio::test]
1364 async fn test_backpressure_drop_newest() {
1365 let queue = InMemoryQueue::with_backpressure(BackpressurePolicy {
1366 max_queue_size: 1,
1367 on_overflow: OverflowStrategy::DropNewest,
1368 });
1369
1370 queue.publish("topic", b"m1").await.unwrap();
1371 let result = queue.publish("topic", b"m2").await;
1373 assert!(result.is_ok());
1374
1375 assert_eq!(queue.message_count("topic").await, 1);
1376 let m = queue.consume("topic").await.unwrap().unwrap();
1378 assert_eq!(m.payload, b"m1", "newest should be dropped");
1379 }
1380
1381 #[tokio::test]
1383 async fn test_backpressure_reject() {
1384 let queue = InMemoryQueue::with_backpressure(BackpressurePolicy {
1385 max_queue_size: 1,
1386 on_overflow: OverflowStrategy::Reject,
1387 });
1388
1389 queue.publish("topic", b"m1").await.unwrap();
1390 let result = queue.publish("topic", b"m2").await;
1392 assert!(result.is_err());
1393
1394 assert_eq!(queue.message_count("topic").await, 1);
1395 }
1396
1397 #[tokio::test]
1399 async fn test_backpressure_block_unblocks_on_consume() {
1400 let queue = InMemoryQueue::with_backpressure(BackpressurePolicy {
1401 max_queue_size: 1,
1402 on_overflow: OverflowStrategy::Block,
1403 });
1404 queue.publish("topic", b"m1").await.unwrap();
1405
1406 let queue_clone = queue.clone();
1408 let handle = tokio::spawn(async move { queue_clone.publish("topic", b"m2").await });
1409
1410 tokio::time::sleep(Duration::from_millis(50)).await;
1412 assert!(!handle.is_finished(), "publish should be blocked");
1413
1414 queue.consume("topic").await.unwrap();
1416
1417 let result = tokio::time::timeout(Duration::from_secs(1), handle)
1419 .await
1420 .expect("publish should complete after consume");
1421 assert!(result.is_ok(), "publish should succeed: {:?}", result);
1422
1423 assert_eq!(queue.message_count("topic").await, 1);
1425 let m = queue.consume("topic").await.unwrap().unwrap();
1426 assert_eq!(m.payload, b"m2");
1427 }
1428
1429 #[tokio::test]
1431 async fn test_backpressure_block_times_out_when_queue_stays_full() {
1432 let queue = InMemoryQueue::with_backpressure(BackpressurePolicy {
1433 max_queue_size: 1,
1434 on_overflow: OverflowStrategy::Block,
1435 });
1436 queue.publish("topic", b"m1").await.unwrap();
1437
1438 let result =
1440 tokio::time::timeout(Duration::from_millis(100), queue.publish("topic", b"m2")).await;
1441
1442 assert!(result.is_err(), "publish should block and time out");
1444
1445 assert_eq!(queue.message_count("topic").await, 1);
1447 }
1448
1449 #[tokio::test]
1451 async fn test_backpressure_block_isolated_per_topic() {
1452 let queue = InMemoryQueue::with_backpressure(BackpressurePolicy {
1453 max_queue_size: 1,
1454 on_overflow: OverflowStrategy::Block,
1455 });
1456 queue.publish("topic-a", b"a1").await.unwrap();
1457 queue.publish("topic-b", b"b1").await.unwrap();
1458
1459 let queue_clone = queue.clone();
1462 let handle = tokio::spawn(async move { queue_clone.publish("topic-a", b"a2").await });
1463
1464 tokio::time::sleep(Duration::from_millis(50)).await;
1465 assert!(!handle.is_finished(), "topic-a publish should be blocked");
1466
1467 queue.consume("topic-b").await.unwrap();
1469 tokio::time::sleep(Duration::from_millis(50)).await;
1470 assert!(
1471 !handle.is_finished(),
1472 "topic-a publish should still be blocked after topic-b consume"
1473 );
1474
1475 queue.consume("topic-a").await.unwrap();
1477 let result = tokio::time::timeout(Duration::from_secs(1), handle)
1478 .await
1479 .expect("publish should complete after topic-a consume");
1480 assert!(result.is_ok());
1481 }
1482}