1use 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#[derive(Default)]
37pub struct MemoryStore {
38 state: Mutex<State>,
39}
40
41impl MemoryStore {
42 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}