ave_actors_store/
memory.rs

1//! In-memory [`DbManager`] backend, intended for tests and ephemeral usage.
2
3use crate::{
4    database::{Collection, DbManager, State},
5    error::{Error, StoreOperation},
6};
7
8use std::{
9    collections::{BTreeMap, HashMap},
10    sync::{Arc, RwLock},
11};
12
13type MemoryData = Arc<
14    RwLock<HashMap<(String, String), Arc<RwLock<BTreeMap<String, Vec<u8>>>>>>,
15>;
16
17/// In-memory database manager backed by a shared `Arc<RwLock<...>>`.
18///
19/// All collections and state stores created by the same `MemoryManager` instance
20/// share the underlying `HashMap`, so data survives across handle clones just
21/// as it would with a real database.
22#[derive(Default, Clone)]
23pub struct MemoryManager {
24    data: MemoryData,
25}
26
27impl DbManager<MemoryStore, MemoryStore> for MemoryManager {
28    fn create_state(
29        &self,
30        name: &str,
31        prefix: &str,
32    ) -> Result<MemoryStore, Error> {
33        let mut data_lock = self.data.write().map_err(|e| Error::Store {
34            operation: StoreOperation::LockManagerData,
35            reason: format!("{}", e),
36        })?;
37        let data = data_lock
38            .entry((name.to_owned(), prefix.to_owned()))
39            .or_insert_with(|| Arc::new(RwLock::new(BTreeMap::new())))
40            .clone();
41        drop(data_lock);
42
43        Ok(MemoryStore {
44            name: name.to_owned(),
45            prefix: prefix.to_owned(),
46            data,
47        })
48    }
49
50    fn stop(&mut self) -> Result<(), Error> {
51        Ok(())
52    }
53
54    fn create_collection(
55        &self,
56        name: &str,
57        prefix: &str,
58    ) -> Result<MemoryStore, Error> {
59        let mut data_lock = self.data.write().map_err(|e| Error::Store {
60            operation: StoreOperation::LockManagerData,
61            reason: format!("{}", e),
62        })?;
63        let data = data_lock
64            .entry((name.to_owned(), prefix.to_owned()))
65            .or_insert_with(|| Arc::new(RwLock::new(BTreeMap::new())))
66            .clone();
67        drop(data_lock);
68
69        Ok(MemoryStore {
70            name: name.to_owned(),
71            prefix: prefix.to_owned(),
72            data,
73        })
74    }
75}
76
77/// In-memory implementation of both [`Collection`] and [`State`].
78///
79/// Data is stored in a `BTreeMap` behind an `Arc<RwLock<...>>` so it can be
80/// shared across clones. All keys in collection mode are prefixed with
81/// `"<prefix>."` to avoid collisions with state-mode entries.
82#[derive(Default, Clone)]
83pub struct MemoryStore {
84    name: String,
85    prefix: String,
86    data: Arc<RwLock<BTreeMap<String, Vec<u8>>>>,
87}
88
89impl MemoryStore {
90    fn collection_prefix(&self) -> String {
91        format!("{}.", self.prefix)
92    }
93}
94
95impl State for MemoryStore {
96    fn name(&self) -> &str {
97        &self.name
98    }
99
100    fn get(&self) -> Result<Vec<u8>, Error> {
101        let lock = self.data.read().map_err(|e| Error::Store {
102            operation: StoreOperation::LockData,
103            reason: format!("{}", e),
104        })?;
105
106        lock.get(&self.prefix).map_or_else(
107            || {
108                Err(Error::EntryNotFound {
109                    key: self.prefix.clone(),
110                })
111            },
112            |value| Ok(value.clone()),
113        )
114    }
115
116    fn put(&mut self, data: &[u8]) -> Result<(), Error> {
117        self.data
118            .write()
119            .map_err(|e| Error::Store {
120                operation: StoreOperation::LockData,
121                reason: format!("{}", e),
122            })?
123            .insert(self.prefix.clone(), data.to_vec());
124
125        Ok(())
126    }
127
128    fn del(&mut self) -> Result<(), Error> {
129        let mut lock = self.data.write().map_err(|e| Error::Store {
130            operation: StoreOperation::LockData,
131            reason: format!("{}", e),
132        })?;
133        match lock.remove(&self.prefix) {
134            Some(_) => Ok(()),
135            None => Err(Error::EntryNotFound {
136                key: self.prefix.clone(),
137            }),
138        }
139    }
140
141    fn purge(&mut self) -> Result<(), Error> {
142        self.data
143            .write()
144            .map_err(|e| Error::Store {
145                operation: StoreOperation::LockData,
146                reason: format!("{}", e),
147            })?
148            .remove(&self.prefix);
149        Ok(())
150    }
151}
152
153impl Collection for MemoryStore {
154    fn last(&self) -> Result<Option<(String, Vec<u8>)>, Error> {
155        let mut iter = self.iter(true)?;
156        iter.next().transpose()
157    }
158
159    fn name(&self) -> &str {
160        &self.name
161    }
162
163    fn get(&self, key: &str) -> Result<Vec<u8>, Error> {
164        let key = format!("{}.{}", self.prefix, key);
165        let lock = self.data.read().map_err(|e| Error::Store {
166            operation: StoreOperation::LockData,
167            reason: format!("{}", e),
168        })?;
169
170        lock.get(&key).map_or_else(
171            || Err(Error::EntryNotFound { key: key.clone() }),
172            |value| Ok(value.clone()),
173        )
174    }
175
176    fn put(&mut self, key: &str, data: &[u8]) -> Result<(), Error> {
177        let key = format!("{}.{}", self.prefix, key);
178        self.data
179            .write()
180            .map_err(|e| Error::Store {
181                operation: StoreOperation::LockData,
182                reason: format!("{}", e),
183            })?
184            .insert(key, data.to_vec());
185
186        Ok(())
187    }
188
189    fn del(&mut self, key: &str) -> Result<(), Error> {
190        let key = format!("{}.{}", self.prefix, key);
191        let mut lock = self.data.write().map_err(|e| Error::Store {
192            operation: StoreOperation::LockData,
193            reason: format!("{}", e),
194        })?;
195        match lock.remove(&key) {
196            Some(_) => Ok(()),
197            None => Err(Error::EntryNotFound { key }),
198        }
199    }
200
201    fn purge(&mut self) -> Result<(), Error> {
202        let mut lock = self.data.write().map_err(|e| Error::Store {
203            operation: StoreOperation::LockData,
204            reason: format!("{}", e),
205        })?;
206        let collection_prefix = self.collection_prefix();
207
208        let keys_to_remove: Vec<String> = lock
209            .keys()
210            .filter(|key| key.starts_with(&collection_prefix))
211            .cloned()
212            .collect();
213        for key in keys_to_remove {
214            lock.remove(&key);
215        }
216        drop(lock);
217        Ok(())
218    }
219
220    fn iter<'a>(
221        &'a self,
222        reverse: bool,
223    ) -> Result<
224        Box<dyn Iterator<Item = Result<(String, Vec<u8>), Error>> + 'a>,
225        Error,
226    > {
227        let lock = self.data.read().map_err(|e| Error::Store {
228            operation: StoreOperation::LockData,
229            reason: format!("{}", e),
230        })?;
231        let collection_prefix = self.collection_prefix();
232        let prefix_len = collection_prefix.len();
233
234        let items: Vec<(String, Vec<u8>)> = if reverse {
235            lock.iter()
236                .rev()
237                .filter(|(key, _)| key.starts_with(&collection_prefix))
238                .map(|(key, value)| {
239                    let key = &key[prefix_len..];
240                    (key.to_owned(), value.clone())
241                })
242                .collect()
243        } else {
244            lock.iter()
245                .filter(|(key, _)| key.starts_with(&collection_prefix))
246                .map(|(key, value)| {
247                    let key = &key[prefix_len..];
248                    (key.to_owned(), value.clone())
249                })
250                .collect()
251        };
252
253        Ok(Box::new(items.into_iter().map(Ok)))
254    }
255}
256
257#[cfg(test)]
258mod tests {
259    use super::*;
260    use crate::test_store_trait;
261    test_store_trait! {
262        unit_test_memory_manager:crate::memory::MemoryManager:MemoryStore
263    }
264}
ave_actors_store/memory.rs

ave_actors_store/
memory.rs
