feat(dmq): make consumer server queue drain expired messages

Author: jpraynaud

internal/mithril-dmq/src/consumer/server/queue.rs

@@ -1,13 +1,18 @@
-use std::collections::VecDeque;
+use std::{collections::VecDeque, sync::Arc};
 
+use mithril_common::StdResult;
 use tokio::sync::{Mutex, Notify};
 
-use crate::DmqMessage;
+use crate::{
+    DmqMessage,
+    model::{SystemUnixTimestampProvider, UnixTimestampProvider},
+};
 
 /// A queue for storing DMQ messages.
 pub(crate) struct MessageQueue {
     messages: Mutex<VecDeque<DmqMessage>>,
     new_message_notify: Notify,
+    timestamp_provider: Arc<dyn UnixTimestampProvider>,
     max_size: usize,
 }
 
@@ -16,33 +21,62 @@ impl MessageQueue {
     const MAX_SIZE_DEFAULT: usize = 10000;
 
     /// Creates a new instance of [BlockingNonBlockingQueue].
-    pub fn new(max_size: usize) -> Self {
+    pub fn new(max_size: usize, timestamp_provider: Arc<dyn UnixTimestampProvider>) -> Self {
         Self {
             messages: Mutex::new(VecDeque::new()),
             new_message_notify: Notify::new(),
+            timestamp_provider,
             max_size,
         }
     }
 
-    /// Enqueues a new message into the queue.
-    pub async fn enqueue(&self, message: DmqMessage) {
+    /// Cleans the queue
+    ///
+    /// Removes expired messages and ensures the queue does not exceed the maximum size.
+    async fn clean_queue(&self) {
         let mut message_queue_guard = self.messages.lock().await;
-        (*message_queue_guard).push_back(message);
+        // Remove expired messages from the front of the queue
+        // There may be other expired messages in the queue, but they will be removed on dequeue
+        // This avoids full scan of the queue.
+        while let Some(message) = message_queue_guard.front()
+            && self.has_message_expired(message).unwrap_or(false)
+        {
+            message_queue_guard.pop_front();
+        }
 
         while message_queue_guard.len() > self.max_size {
             message_queue_guard.pop_front();
         }
+    }
+
+    /// Checks if a message has expired.
+    fn has_message_expired(&self, message: &DmqMessage) -> StdResult<bool> {
+        let current_timestamp: u32 = self.timestamp_provider.current_timestamp()?.try_into()?;
+        Ok(message.msg_payload.expires_at < current_timestamp)
+    }
+
+    /// Enqueues a new message into the queue.
+    pub async fn enqueue(&self, message: DmqMessage) {
+        {
+            // Run in a block to avoid Mutex deadlock in clean_queue
+            let mut message_queue_guard = self.messages.lock().await;
+            (*message_queue_guard).push_back(message);
+        }
+        self.clean_queue().await;
 
         self.new_message_notify.notify_waiters();
     }
 
     /// Returns the messages from the queue in a non blocking way, if available.
     pub async fn dequeue_non_blocking(&self, limit: Option<usize>) -> Vec<DmqMessage> {
+        self.clean_queue().await;
         let mut message_queue_guard = self.messages.lock().await;
         let limit = limit.unwrap_or((*message_queue_guard).len());
         let mut messages = Vec::new();
         for _ in 0..limit {
-            if let Some(message) = (*message_queue_guard).pop_front() {
+            if let Some(message) = (*message_queue_guard).pop_front()
+                && !self.has_message_expired(&message).unwrap_or(false)
+            {
                 messages.push(message);
             }
         }
@@ -76,7 +110,10 @@ impl MessageQueue {
 
 impl Default for MessageQueue {
     fn default() -> Self {
-        Self::new(Self::MAX_SIZE_DEFAULT)
+        Self::new(
+            Self::MAX_SIZE_DEFAULT,
+            Arc::new(SystemUnixTimestampProvider),
+        )
     }
 }
 
@@ -88,6 +125,8 @@ mod tests {
     use pallas_network::miniprotocols::localmsgsubmission::{DmqMsg, DmqMsgPayload};
     use tokio::time::sleep;
 
+    use crate::model::MockUnixTimestampProvider;
+
     use super::*;
 
     fn fake_msg() -> DmqMsg {
@@ -104,20 +143,38 @@ mod tests {
         }
     }
 
-    fn fake_messages(range: RangeInclusive<u8>) -> Vec<DmqMessage> {
+    fn fake_messages(range: RangeInclusive<u8>, expires_at: u32) -> Vec<DmqMessage> {
         range
             .map(|i| {
                 let mut message = fake_msg();
                 message.msg_payload.msg_id = vec![i];
+                message.msg_payload.expires_at = expires_at;
                 message.into()
             })
             .collect::<Vec<_>>()
     }
 
+    fn create_queue(max_size: usize, current_timestamp: u64) -> MessageQueue {
+        MessageQueue::new(
+            max_size,
+            Arc::new({
+                let mut mock_timestamp_provider = MockUnixTimestampProvider::new();
+                mock_timestamp_provider
+                    .expect_current_timestamp()
+                    .returning(move || Ok(current_timestamp));
+
+                mock_timestamp_provider
+            }),
+        )
+    }
+
     #[tokio::test]
     async fn enqueue_and_dequeue_non_blocking_no_limit() {
-        let queue = MessageQueue::default();
-        let messages = fake_messages(1..=5);
+        let max_size = 100;
+        let current_timestamp = 10;
+        let expires_at = 100;
+        let queue = create_queue(max_size, current_timestamp);
+        let messages = fake_messages(1..=5, expires_at);
         for message in messages.clone() {
             queue.enqueue(message).await;
         }
@@ -130,8 +187,11 @@ mod tests {
 
     #[tokio::test]
     async fn enqueue_and_dequeue_non_blocking_with_limit() {
-        let queue = MessageQueue::default();
-        let messages = fake_messages(1..=5);
+        let max_size = 100;
+        let current_timestamp = 10;
+        let expires_at = 100;
+        let queue = create_queue(max_size, current_timestamp);
+        let messages = fake_messages(1..=5, expires_at);
         for message in messages.clone() {
             queue.enqueue(message).await;
         }
@@ -144,8 +204,11 @@ mod tests {
 
     #[tokio::test]
     async fn enqueue_and_dequeue_blocking_no_limit() {
-        let queue = MessageQueue::default();
-        let messages = fake_messages(1..=5);
+        let max_size = 100;
+        let current_timestamp = 10;
+        let expires_at = 100;
+        let queue = create_queue(max_size, current_timestamp);
+        let messages = fake_messages(1..=5, expires_at);
         for message in messages.clone() {
             queue.enqueue(message).await;
         }
@@ -158,8 +221,11 @@ mod tests {
 
     #[tokio::test]
     async fn enqueue_and_dequeue_blocking_with_limit() {
-        let queue = MessageQueue::default();
-        let messages = fake_messages(1..=5);
+        let max_size = 100;
+        let current_timestamp = 10;
+        let expires_at = 100;
+        let queue = create_queue(max_size, current_timestamp);
+        let messages = fake_messages(1..=5, expires_at);
         for message in messages.clone() {
             queue.enqueue(message).await;
         }
@@ -172,7 +238,9 @@ mod tests {
 
     #[tokio::test]
     async fn dequeue_blocking_blocks_when_no_message_available() {
-        let queue = MessageQueue::default();
+        let max_size = 100;
+        let current_timestamp = 10;
+        let queue = create_queue(max_size, current_timestamp);
 
         let result = tokio::select!(
             _res = sleep(Duration::from_millis(100)) => {Err(anyhow!("Timeout"))},
@@ -183,10 +251,12 @@ mod tests {
     }
 
     #[tokio::test]
-    async fn enqueue_blocks_over_max_size_drains_oldest_messages() {
-        let max_queue_size = 3;
-        let queue = MessageQueue::new(max_queue_size);
-        let messages = fake_messages(1..=5);
+    async fn queue_drains_oldest_messages_when_full() {
+        let max_size = 3;
+        let current_timestamp = 10;
+        let expires_at = 100;
+        let queue = create_queue(max_size, current_timestamp);
+        let messages = fake_messages(1..=5, expires_at);
         for message in messages.clone() {
             queue.enqueue(message).await;
         }
@@ -196,4 +266,26 @@ mod tests {
 
         assert_eq!(messages[2..=4].to_vec(), dequeued_messages);
     }
+
+    #[tokio::test]
+    async fn queue_drains_expired_message() {
+        let max_size = 3;
+        let current_timestamp = 10;
+        let expires_at_expired = 1;
+        let expires_at_non_expired = 100;
+        let queue = create_queue(max_size, current_timestamp);
+        let mut messages = fake_messages(1..=10, expires_at_expired);
+        for (index, mut message) in messages.clone().into_iter().enumerate() {
+            if index >= 5 {
+                message.msg_payload.expires_at = expires_at_non_expired;
+                messages[index] = message.clone();
+            }
+            queue.enqueue(message).await;
+        }
+        let limit = None;
+
+        let dequeued_messages = queue.dequeue_blocking(limit).await;
+
+        assert_eq!(messages[7..=9].to_vec(), dequeued_messages);
+    }
 }