Skip to main content

hotfix_store/
in_memory.rs

1//! An in-memory message store that loses its state on restart.
2
3use crate::{MessageStore, Result};
4use chrono::{DateTime, Utc};
5use std::collections::HashMap;
6
7/// An in-memory message store implementation.
8///
9/// This store keeps all messages in memory and does not persist them.
10/// Use this only for testing or when persistence is not required.
11#[derive(Debug)]
12pub struct InMemoryMessageStore {
13    sender_seq_number: u64,
14    target_seq_number: u64,
15    creation_time: DateTime<Utc>,
16    messages: HashMap<u64, Vec<u8>>,
17}
18
19impl Default for InMemoryMessageStore {
20    fn default() -> Self {
21        Self {
22            sender_seq_number: 0,
23            target_seq_number: 0,
24            creation_time: Utc::now(),
25            messages: HashMap::new(),
26        }
27    }
28}
29
30#[async_trait::async_trait]
31impl MessageStore for InMemoryMessageStore {
32    async fn add(&mut self, sequence_number: u64, message: &[u8]) -> Result<()> {
33        self.messages.insert(sequence_number, message.to_vec());
34        Ok(())
35    }
36
37    async fn get_slice(&self, begin: usize, end: usize) -> Result<Vec<Vec<u8>>> {
38        let mut msgs = Vec::with_capacity(end - begin + 1);
39        for idx in begin..=end {
40            if let Some(msg) = self.messages.get(&(idx as u64)) {
41                msgs.push(msg.to_vec());
42            }
43        }
44        Ok(msgs)
45    }
46
47    fn next_sender_seq_number(&self) -> u64 {
48        self.sender_seq_number + 1
49    }
50
51    fn next_target_seq_number(&self) -> u64 {
52        self.target_seq_number + 1
53    }
54
55    async fn increment_sender_seq_number(&mut self) -> Result<()> {
56        self.sender_seq_number += 1;
57        Ok(())
58    }
59
60    async fn increment_target_seq_number(&mut self) -> Result<()> {
61        self.target_seq_number += 1;
62        Ok(())
63    }
64
65    async fn set_target_seq_number(&mut self, seq_number: u64) -> Result<()> {
66        self.target_seq_number = seq_number;
67        Ok(())
68    }
69
70    async fn reset(&mut self) -> Result<()> {
71        self.sender_seq_number = 0;
72        self.target_seq_number = 0;
73        self.messages.clear();
74        self.creation_time = Utc::now();
75        Ok(())
76    }
77
78    fn creation_time(&self) -> DateTime<Utc> {
79        self.creation_time
80    }
81}