Skip to main content

truefix_store/
memory.rs

1//! Volatile in-memory message store.
2
3use std::collections::BTreeMap;
4use std::sync::Mutex;
5
6use async_trait::async_trait;
7
8use crate::{MessageStore, StoreError};
9
10pub(crate) struct State {
11    pub(crate) sender: u64,
12    pub(crate) target: u64,
13    pub(crate) messages: BTreeMap<u64, Vec<u8>>,
14}
15
16impl Default for State {
17    fn default() -> Self {
18        Self {
19            sender: 1,
20            target: 1,
21            messages: BTreeMap::new(),
22        }
23    }
24}
25
26impl State {
27    pub(crate) fn range(&self, begin: u64, end: u64) -> Vec<(u64, Vec<u8>)> {
28        self.messages
29            .range(begin..=end)
30            .map(|(seq, bytes)| (*seq, bytes.clone()))
31            .collect()
32    }
33}
34
35/// An in-memory store (does not survive process restart).
36#[derive(Default)]
37pub struct MemoryStore {
38    state: Mutex<State>,
39}
40
41impl MemoryStore {
42    /// Create an empty in-memory store.
43    pub fn new() -> Self {
44        Self::default()
45    }
46}
47
48fn lock<T>(m: &Mutex<T>) -> Result<std::sync::MutexGuard<'_, T>, StoreError> {
49    m.lock()
50        .map_err(|_| StoreError::Backend("poisoned lock".into()))
51}
52
53#[async_trait]
54impl MessageStore for MemoryStore {
55    async fn next_sender_seq(&self) -> Result<u64, StoreError> {
56        Ok(lock(&self.state)?.sender)
57    }
58    async fn next_target_seq(&self) -> Result<u64, StoreError> {
59        Ok(lock(&self.state)?.target)
60    }
61    async fn set_next_sender_seq(&self, seq: u64) -> Result<(), StoreError> {
62        lock(&self.state)?.sender = seq;
63        Ok(())
64    }
65    async fn set_next_target_seq(&self, seq: u64) -> Result<(), StoreError> {
66        lock(&self.state)?.target = seq;
67        Ok(())
68    }
69    async fn save(&self, seq: u64, message: &[u8]) -> Result<(), StoreError> {
70        lock(&self.state)?.messages.insert(seq, message.to_vec());
71        Ok(())
72    }
73    async fn get(&self, begin: u64, end: u64) -> Result<Vec<(u64, Vec<u8>)>, StoreError> {
74        Ok(lock(&self.state)?.range(begin, end))
75    }
76    async fn reset(&self) -> Result<(), StoreError> {
77        let mut s = lock(&self.state)?;
78        *s = State::default();
79        Ok(())
80    }
81    async fn contains(&self, seq: u64) -> Result<bool, StoreError> {
82        Ok(lock(&self.state)?.messages.contains_key(&seq))
83    }
84}