serai_db/
mem.rs

1use core::fmt::Debug;
2use alloc::{sync::Arc, vec::Vec};
3use std::{sync::RwLock, collections::HashMap};
4
5use crate::{Get, Db};
6
7/// A transaction for [`MemDb`].
8#[must_use]
9#[derive(PartialEq, Eq, Debug)]
10pub struct MemDbTxn<'db> {
11  db: &'db MemDb,
12  queued_changes: HashMap<Vec<u8>, Option<Vec<u8>>>,
13}
14
15impl Get for MemDbTxn<'_> {
16  fn get(&self, key: impl AsRef<[u8]>) -> Option<impl AsRef<[u8]>> {
17    match self.queued_changes.get(key.as_ref()) {
18      Some(Some(value)) => Some(value.clone()),
19      Some(None) => None?,
20      None => self.db.get(key.as_ref()).map(|bytes| bytes.as_ref().to_vec()),
21    }
22  }
23}
24
25impl crate::Transaction for MemDbTxn<'_> {
26  fn set(&mut self, key: impl AsRef<[u8]>, value: impl AsRef<[u8]>) {
27    self.queued_changes.insert(key.as_ref().to_vec(), Some(value.as_ref().to_vec()));
28  }
29  fn del(&mut self, key: impl AsRef<[u8]>) {
30    self.queued_changes.insert(key.as_ref().to_vec(), None);
31  }
32  fn commit(mut self) {
33    let mut db = self.db.0.write().unwrap();
34    for (key, value) in self.queued_changes.drain() {
35      match value {
36        Some(value) => {
37          db.insert(key, value);
38        }
39        None => {
40          db.remove(&key);
41        }
42      }
43    }
44  }
45}
46
47/// An in-memory database.
48#[derive(Clone, Debug)]
49pub struct MemDb(Arc<RwLock<HashMap<Vec<u8>, Vec<u8>>>>);
50
51impl PartialEq for MemDb {
52  fn eq(&self, other: &MemDb) -> bool {
53    *self.0.read().unwrap() == *other.0.read().unwrap()
54  }
55}
56impl Eq for MemDb {}
57
58impl Default for MemDb {
59  fn default() -> MemDb {
60    MemDb(Arc::new(RwLock::new(HashMap::new())))
61  }
62}
63
64impl MemDb {
65  /// Create a new in-memory database.
66  pub fn new() -> MemDb {
67    MemDb::default()
68  }
69}
70
71impl Get for MemDb {
72  fn get(&self, key: impl AsRef<[u8]>) -> Option<impl AsRef<[u8]>> {
73    self.0.read().unwrap().get(key.as_ref()).cloned()
74  }
75}
76impl Db for MemDb {
77  type Transaction<'db> = MemDbTxn<'db>;
78  fn txn(&mut self) -> MemDbTxn<'_> {
79    MemDbTxn { db: self, queued_changes: HashMap::new() }
80  }
81}
serai_db/mem.rs

serai_db/
mem.rs
