keyvaluedb_memorydb/
lib.rs

1//! A key-value in-memory database that implements the `KeyValueDB` trait
2
3#![deny(clippy::all)]
4
5use keyvaluedb::{
6    DBKey, DBKeyRef, DBKeyValueRef, DBOp, DBTransaction, DBTransactionError, DBValue, KeyValueDB,
7    KeyValueDBPinBoxFuture,
8};
9use parking_lot::RwLock;
10use std::{
11    collections::{BTreeMap, HashMap},
12    io,
13    sync::Arc,
14};
15
16/// A key-value database fulfilling the `KeyValueDB` trait, living in memory.
17/// This is generally intended for tests and is not particularly optimized.
18#[derive(Clone)]
19pub struct InMemory {
20    columns: Arc<RwLock<HashMap<u32, BTreeMap<DBKey, DBValue>>>>,
21}
22
23/// Create an in-memory database with the given number of columns.
24/// Columns will be indexable by 0..`num_cols`
25pub fn create(num_cols: u32) -> InMemory {
26    let mut cols = HashMap::new();
27
28    for idx in 0..num_cols {
29        cols.insert(idx, BTreeMap::new());
30    }
31
32    InMemory {
33        columns: Arc::new(RwLock::new(cols)),
34    }
35}
36
37impl KeyValueDB for InMemory {
38    fn get<'a>(
39        &'a self,
40        col: u32,
41        key: &'a [u8],
42    ) -> KeyValueDBPinBoxFuture<'a, io::Result<Option<DBValue>>> {
43        let this = self.clone();
44        Box::pin(async move {
45            let columns = this.columns.read();
46            match columns.get(&col) {
47                None => Err(io::Error::from(io::ErrorKind::NotFound)),
48                Some(map) => Ok(map.get(key).cloned()),
49            }
50        })
51    }
52
53    /// Remove a value by key, returning the old value
54    fn delete<'a>(
55        &'a self,
56        col: u32,
57        key: &'a [u8],
58    ) -> KeyValueDBPinBoxFuture<'a, io::Result<Option<DBValue>>> {
59        let this = self.clone();
60        Box::pin(async move {
61            let mut columns = this.columns.write();
62            match columns.get_mut(&col) {
63                None => Err(io::Error::from(io::ErrorKind::NotFound)),
64                Some(map) => Ok(map.remove(key)),
65            }
66        })
67    }
68
69    fn write(
70        &self,
71        transaction: DBTransaction,
72    ) -> KeyValueDBPinBoxFuture<'_, Result<(), DBTransactionError>> {
73        let this = self.clone();
74        Box::pin(async move {
75            let mut columns = this.columns.write();
76            let ops = transaction.ops;
77            for op in ops {
78                match op {
79                    DBOp::Insert { col, key, value } => {
80                        if let Some(col) = columns.get_mut(&col) {
81                            col.insert(key, value);
82                        }
83                    }
84                    DBOp::Delete { col, key } => {
85                        if let Some(col) = columns.get_mut(&col) {
86                            col.remove(&*key);
87                        }
88                    }
89                    DBOp::DeletePrefix { col, prefix } => {
90                        if let Some(col) = columns.get_mut(&col) {
91                            use std::ops::Bound;
92                            if prefix.is_empty() {
93                                col.clear();
94                            } else {
95                                let start_range = Bound::Included(prefix.to_vec());
96                                let keys: Vec<_> =
97                                    if let Some(end_range) = keyvaluedb::end_prefix(&prefix[..]) {
98                                        col.range((start_range, Bound::Excluded(end_range)))
99                                            .map(|(k, _)| k.clone())
100                                            .collect()
101                                    } else {
102                                        col.range((start_range, Bound::Unbounded))
103                                            .map(|(k, _)| k.clone())
104                                            .collect()
105                                    };
106                                for key in keys.into_iter() {
107                                    col.remove(&key[..]);
108                                }
109                            }
110                        }
111                    }
112                }
113            }
114            Ok(())
115        })
116    }
117
118    fn iter<
119        'a,
120        T: Send + 'static,
121        C: Send + 'static,
122        F: FnMut(&mut C, DBKeyValueRef) -> io::Result<Option<T>> + Send + Sync + 'static,
123    >(
124        &'a self,
125        col: u32,
126        prefix: Option<&'a [u8]>,
127        mut context: C,
128        mut f: F,
129    ) -> KeyValueDBPinBoxFuture<'a, io::Result<(C, Option<T>)>> {
130        let this = self.clone();
131        Box::pin(async move {
132            match this.columns.read().get(&col) {
133                Some(map) => {
134                    for (k, v) in map {
135                        if let Some(p) = prefix {
136                            if !k.starts_with(p) {
137                                continue;
138                            }
139                        }
140                        match f(&mut context, (k, v)) {
141                            Ok(None) => (),
142                            Ok(Some(v)) => return Ok((context, Some(v))),
143                            Err(e) => return Err(e),
144                        }
145                    }
146                    Ok((context, None))
147                }
148                None => Err(io::Error::from(io::ErrorKind::NotFound)),
149            }
150        })
151    }
152
153    fn iter_keys<
154        'a,
155        T: Send + 'static,
156        C: Send + 'static,
157        F: FnMut(&mut C, DBKeyRef) -> io::Result<Option<T>> + Send + Sync + 'static,
158    >(
159        &'a self,
160        col: u32,
161        prefix: Option<&'a [u8]>,
162        mut context: C,
163        mut f: F,
164    ) -> KeyValueDBPinBoxFuture<'a, io::Result<(C, Option<T>)>> {
165        let this = self.clone();
166        Box::pin(async move {
167            match this.columns.read().get(&col) {
168                Some(map) => {
169                    for k in map.keys() {
170                        if let Some(p) = prefix {
171                            if !k.starts_with(p) {
172                                continue;
173                            }
174                        }
175                        match f(&mut context, k) {
176                            Ok(None) => (),
177                            Ok(Some(v)) => return Ok((context, Some(v))),
178                            Err(e) => return Err(e),
179                        }
180                    }
181                    Ok((context, None))
182                }
183                None => Err(io::Error::from(io::ErrorKind::NotFound)),
184            }
185        })
186    }
187
188    fn num_columns(&self) -> io::Result<u32> {
189        Ok(self.columns.read().len() as u32)
190    }
191
192    fn num_keys(&self, col: u32) -> KeyValueDBPinBoxFuture<'_, io::Result<u64>> {
193        let this = self.clone();
194        Box::pin(async move {
195            let c = this.columns.read();
196            let Some(column) = c.get(&col) else {
197                return Err(io::Error::from(io::ErrorKind::NotFound));
198            };
199            Ok(column.len() as u64)
200        })
201    }
202}
203
204#[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))]
205#[cfg(test)]
206mod tests {
207    use super::create;
208    use keyvaluedb_shared_tests as st;
209    use std::io;
210
211    #[tokio::test]
212    async fn get_fails_with_non_existing_column() -> io::Result<()> {
213        let db = create(1);
214        st::test_get_fails_with_non_existing_column(db).await
215    }
216
217    #[tokio::test]
218    async fn put_and_get() -> io::Result<()> {
219        let db = create(1);
220        st::test_put_and_get(db).await
221    }
222
223    #[tokio::test]
224    async fn num_keys() -> io::Result<()> {
225        let db = create(1);
226        st::test_num_keys(db).await
227    }
228
229    #[tokio::test]
230    async fn delete_and_get() -> io::Result<()> {
231        let db = create(1);
232        st::test_delete_and_get(db).await
233    }
234
235    #[tokio::test]
236    async fn delete_prefix() -> io::Result<()> {
237        let db = create(st::DELETE_PREFIX_NUM_COLUMNS);
238        st::test_delete_prefix(db).await
239    }
240
241    #[tokio::test]
242    async fn iter() -> io::Result<()> {
243        let db = create(1);
244        st::test_iter(db).await
245    }
246
247    #[tokio::test]
248    async fn iter_keys() -> io::Result<()> {
249        let db = create(1);
250        st::test_iter_keys(db).await
251    }
252
253    #[tokio::test]
254    async fn iter_with_prefix() -> io::Result<()> {
255        let db = create(1);
256        st::test_iter_with_prefix(db).await
257    }
258
259    #[tokio::test]
260    async fn complex() -> io::Result<()> {
261        let db = create(1);
262        st::test_complex(db).await
263    }
264
265    #[tokio::test]
266    async fn cleanup() -> io::Result<()> {
267        let db = create(1);
268        st::test_cleanup(db).await
269    }
270}
keyvaluedb_memorydb/lib.rs

keyvaluedb_memorydb/
lib.rs
