capp_queue/backend/
memory.rs

1//! In-memory implementation of TaskStorage trait. The storage allows tasks to be
2//! pushed to and popped from a queue, and also allows tasks to be set and
3//! retrieved by their UUID.
4use crate::task::{Task, TaskId};
5use crate::{TaskQueue, TaskQueueError, TaskSerializer};
6use async_trait::async_trait;
7use serde::Serialize;
8use serde::de::DeserializeOwned;
9use std::collections::{HashMap, VecDeque};
10use std::marker::PhantomData;
11use std::sync::Mutex;
12
13pub struct InMemoryTaskQueue<D, S>
14where
15    S: TaskSerializer,
16{
17    pub hashmap: Mutex<HashMap<TaskId, Vec<u8>>>,
18    pub list: Mutex<VecDeque<TaskId>>,
19    pub inflight: Mutex<HashMap<TaskId, ()>>,
20    pub dlq: Mutex<HashMap<TaskId, Vec<u8>>>,
21    _marker: PhantomData<(D, S)>,
22}
23
24impl<D, S> InMemoryTaskQueue<D, S>
25where
26    S: TaskSerializer,
27{
28    pub fn new() -> Self {
29        Self {
30            hashmap: Mutex::new(HashMap::new()),
31            list: Mutex::new(VecDeque::new()),
32            inflight: Mutex::new(HashMap::new()),
33            dlq: Mutex::new(HashMap::new()),
34            _marker: PhantomData,
35        }
36    }
37}
38
39impl<D, S> Default for InMemoryTaskQueue<D, S>
40where
41    S: TaskSerializer,
42{
43    fn default() -> Self {
44        Self::new()
45    }
46}
47
48#[async_trait]
49impl<D, S> TaskQueue<D> for InMemoryTaskQueue<D, S>
50where
51    D: std::fmt::Debug
52        + Clone
53        + Serialize
54        + DeserializeOwned
55        + Send
56        + Sync
57        + 'static,
58    S: TaskSerializer + Send + Sync,
59{
60    async fn push(&self, task: &Task<D>) -> Result<(), TaskQueueError> {
61        let mut list = self
62            .list
63            .lock()
64            .map_err(|e| TaskQueueError::QueueError(e.to_string()))?;
65        let mut hashmap = self
66            .hashmap
67            .lock()
68            .map_err(|e| TaskQueueError::QueueError(e.to_string()))?;
69        let mut inflight = self
70            .inflight
71            .lock()
72            .map_err(|e| TaskQueueError::QueueError(e.to_string()))?;
73
74        let task_bytes = S::serialize_task(task)?;
75        hashmap.insert(task.task_id, task_bytes);
76        list.push_back(task.task_id);
77        inflight.remove(&task.task_id);
78        Ok(())
79    }
80
81    async fn pop(&self) -> Result<Task<D>, TaskQueueError> {
82        let mut list = self
83            .list
84            .lock()
85            .map_err(|e| TaskQueueError::QueueError(e.to_string()))?;
86        let hashmap = self
87            .hashmap
88            .lock()
89            .map_err(|e| TaskQueueError::QueueError(e.to_string()))?;
90        let mut inflight = self
91            .inflight
92            .lock()
93            .map_err(|e| TaskQueueError::QueueError(e.to_string()))?;
94
95        if let Some(task_id) = list.pop_front() {
96            let task_bytes = hashmap
97                .get(&task_id)
98                .ok_or(TaskQueueError::TaskNotFound(task_id))?;
99            inflight.insert(task_id, ());
100            S::deserialize_task(task_bytes)
101        } else {
102            Err(TaskQueueError::QueueEmpty)
103        }
104    }
105
106    async fn ack(&self, task_id: &TaskId) -> Result<(), TaskQueueError> {
107        let mut hashmap = self
108            .hashmap
109            .lock()
110            .map_err(|e| TaskQueueError::QueueError(e.to_string()))?;
111        let mut inflight = self
112            .inflight
113            .lock()
114            .map_err(|e| TaskQueueError::QueueError(e.to_string()))?;
115        hashmap
116            .remove(task_id)
117            .ok_or(TaskQueueError::TaskNotFound(*task_id))?;
118        inflight.remove(task_id);
119        Ok(())
120    }
121
122    async fn nack(&self, task: &Task<D>) -> Result<(), TaskQueueError> {
123        let mut dlq = self
124            .dlq
125            .lock()
126            .map_err(|e| TaskQueueError::QueueError(e.to_string()))?;
127        let task_bytes = S::serialize_task(task)?;
128        dlq.insert(task.task_id, task_bytes);
129
130        let mut hashmap = self
131            .hashmap
132            .lock()
133            .map_err(|e| TaskQueueError::QueueError(e.to_string()))?;
134        let mut inflight = self
135            .inflight
136            .lock()
137            .map_err(|e| TaskQueueError::QueueError(e.to_string()))?;
138        hashmap
139            .remove(&task.task_id)
140            .ok_or(TaskQueueError::TaskNotFound(task.task_id))?;
141        inflight.remove(&task.task_id);
142        Ok(())
143    }
144
145    async fn set(&self, task: &Task<D>) -> Result<(), TaskQueueError> {
146        let mut hashmap = self
147            .hashmap
148            .lock()
149            .map_err(|e| TaskQueueError::QueueError(e.to_string()))?;
150        let task_bytes = S::serialize_task(task)?;
151        hashmap.insert(task.task_id, task_bytes);
152        Ok(())
153    }
154
155    async fn recover_inflight(&self) -> Result<u64, TaskQueueError> {
156        let mut list = self
157            .list
158            .lock()
159            .map_err(|e| TaskQueueError::QueueError(e.to_string()))?;
160        let mut inflight = self
161            .inflight
162            .lock()
163            .map_err(|e| TaskQueueError::QueueError(e.to_string()))?;
164        let recovered = inflight.len() as u64;
165        for task_id in inflight.drain().map(|(id, _)| id) {
166            list.push_back(task_id);
167        }
168        Ok(recovered)
169    }
170}
171
172impl<D, S> std::fmt::Debug for InMemoryTaskQueue<D, S>
173where
174    S: TaskSerializer,
175{
176    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
177        let hashmap = self.hashmap.lock().unwrap();
178        let list = self.list.lock().unwrap();
179        let inflight = self.inflight.lock().unwrap();
180        let dlq = self.dlq.lock().unwrap();
181
182        f.debug_struct("InMemoryTaskQueue")
183            .field("hashmap_size", &hashmap.len())
184            .field("list", &*list)
185            .field("inflight_size", &inflight.len())
186            .field("dlq_size", &dlq.len())
187            .finish()
188    }
189}
190
191#[cfg(test)]
192mod tests {
193    use super::*;
194    use crate::JsonSerializer;
195    use serde::{Deserialize, Serialize};
196    use tokio;
197
198    #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
199    struct TestData {
200        value: u32,
201    }
202
203    #[tokio::test]
204    async fn test_push_and_pop() {
205        let queue = InMemoryTaskQueue::<TestData, JsonSerializer>::new();
206        let task = Task::new(TestData { value: 42 });
207
208        queue.push(&task).await.unwrap();
209        let popped_task = queue.pop().await.unwrap();
210
211        assert_eq!(popped_task.payload, TestData { value: 42 });
212    }
213
214    #[tokio::test]
215    async fn test_queue_empty() {
216        let queue = InMemoryTaskQueue::<TestData, JsonSerializer>::new();
217
218        match queue.pop().await {
219            Err(TaskQueueError::QueueEmpty) => (),
220            _ => panic!("Expected QueueEmpty error"),
221        }
222    }
223
224    #[tokio::test]
225    async fn test_ack() {
226        let queue = InMemoryTaskQueue::<TestData, JsonSerializer>::new();
227        let task = Task::new(TestData { value: 42 });
228
229        queue.push(&task).await.unwrap();
230        let popped_task = queue.pop().await.unwrap();
231        queue.ack(&popped_task.task_id).await.unwrap();
232
233        // The queue should be empty after ack
234        assert!(matches!(queue.pop().await, Err(TaskQueueError::QueueEmpty)));
235    }
236
237    #[tokio::test]
238    async fn test_nack() {
239        let queue = InMemoryTaskQueue::<TestData, JsonSerializer>::new();
240        let task = Task::new(TestData { value: 42 });
241
242        queue.push(&task).await.unwrap();
243        let popped_task = queue.pop().await.unwrap();
244        queue.nack(&popped_task).await.unwrap();
245
246        // The queue should be empty after nack
247        assert!(matches!(queue.pop().await, Err(TaskQueueError::QueueEmpty)));
248
249        // The task should be in the DLQ
250        let dlq = queue.dlq.lock().unwrap();
251        assert_eq!(dlq.len(), 1);
252    }
253
254    #[tokio::test]
255    async fn test_set() {
256        let queue = InMemoryTaskQueue::<TestData, JsonSerializer>::new();
257        let mut task = Task::new(TestData { value: 42 });
258
259        queue.push(&task).await.unwrap();
260
261        // Modify the task
262        task.payload.value = 43;
263        queue.set(&task).await.unwrap();
264
265        let updated_task = queue.pop().await.unwrap();
266        assert_eq!(updated_task.payload, TestData { value: 43 });
267    }
268
269    #[tokio::test]
270    async fn test_multiple_tasks() {
271        let queue = InMemoryTaskQueue::<TestData, JsonSerializer>::new();
272        let tasks = vec![
273            Task::new(TestData { value: 1 }),
274            Task::new(TestData { value: 2 }),
275            Task::new(TestData { value: 3 }),
276        ];
277
278        for task in &tasks {
279            queue.push(task).await.unwrap();
280        }
281
282        for expected_task in tasks {
283            let popped_task = queue.pop().await.unwrap();
284            assert_eq!(popped_task.payload, expected_task.payload);
285        }
286
287        // Queue should be empty now
288        assert!(matches!(queue.pop().await, Err(TaskQueueError::QueueEmpty)));
289    }
290
291    #[tokio::test]
292    async fn test_task_not_found() {
293        let queue = InMemoryTaskQueue::<TestData, JsonSerializer>::new();
294        let non_existent_task_id = TaskId::new();
295
296        match queue.ack(&non_existent_task_id).await {
297            Err(TaskQueueError::TaskNotFound(_)) => (),
298            _ => panic!("Expected TaskNotFound error"),
299        }
300    }
301
302    #[tokio::test]
303    async fn pop_moves_to_inflight() {
304        let queue = InMemoryTaskQueue::<TestData, JsonSerializer>::new();
305        let task = Task::new(TestData { value: 1 });
306
307        queue.push(&task).await.unwrap();
308        let popped = queue.pop().await.unwrap();
309
310        assert!(queue.inflight.lock().unwrap().contains_key(&popped.task_id));
311    }
312
313    #[tokio::test]
314    async fn ack_clears_inflight() {
315        let queue = InMemoryTaskQueue::<TestData, JsonSerializer>::new();
316        let task = Task::new(TestData { value: 2 });
317
318        queue.push(&task).await.unwrap();
319        let popped = queue.pop().await.unwrap();
320        queue.ack(&popped.task_id).await.unwrap();
321
322        assert!(queue.inflight.lock().unwrap().is_empty());
323    }
324
325    #[tokio::test]
326    async fn nack_clears_inflight() {
327        let queue = InMemoryTaskQueue::<TestData, JsonSerializer>::new();
328        let task = Task::new(TestData { value: 3 });
329
330        queue.push(&task).await.unwrap();
331        let popped = queue.pop().await.unwrap();
332        queue.nack(&popped).await.unwrap();
333
334        assert!(queue.inflight.lock().unwrap().is_empty());
335        assert_eq!(queue.dlq.lock().unwrap().len(), 1);
336    }
337
338    #[tokio::test]
339    async fn push_retry_clears_inflight() {
340        let queue = InMemoryTaskQueue::<TestData, JsonSerializer>::new();
341        let mut task = Task::new(TestData { value: 4 });
342
343        queue.push(&task).await.unwrap();
344        let popped = queue.pop().await.unwrap();
345        assert!(queue.inflight.lock().unwrap().contains_key(&popped.task_id));
346
347        task.set_retry("retry");
348        queue.push(&task).await.unwrap();
349
350        assert!(queue.inflight.lock().unwrap().is_empty());
351        assert_eq!(queue.list.lock().unwrap().len(), 1);
352    }
353
354    #[tokio::test]
355    async fn recover_inflight_moves_back_to_queue() {
356        let queue = InMemoryTaskQueue::<TestData, JsonSerializer>::new();
357        let task = Task::new(TestData { value: 5 });
358
359        queue.push(&task).await.unwrap();
360        let popped = queue.pop().await.unwrap();
361        assert!(queue.inflight.lock().unwrap().contains_key(&popped.task_id));
362
363        let recovered = queue.recover_inflight().await.unwrap();
364        assert_eq!(recovered, 1);
365        assert!(queue.inflight.lock().unwrap().is_empty());
366
367        let re_popped = queue.pop().await.unwrap();
368        assert_eq!(re_popped.task_id, popped.task_id);
369    }
370}
capp_queue/backend/memory.rs

capp_queue/backend/
memory.rs
