snkrj/
lib.rs

1// https://docs.rs/sanakirja/latest/sanakirja/index.html
2// https://pijul.org/posts/2021-02-06-rethinking-sanakirja/
3
4use std::{
5    collections::{BTreeMap, BTreeSet},
6    fmt::Debug,
7    fs, io, mem,
8    path::PathBuf,
9    result::Result,
10};
11
12use sanakirja::btree::{page, Db_};
13pub use sanakirja::*;
14
15///
16/// A simple wrapper around Sanakirja aatabase that acts as a very fast on disk BTreeMap.
17///
18/// The state of the tree is uncommited until `.export()` is called during which it is unsafe to stop the program.
19///
20pub struct Database<Key, Value>
21where
22    Key: Ord + Clone + Debug + Storable,
23    Value: Storable + PartialEq,
24{
25    path: PathBuf,
26    puts: BTreeMap<Key, Value>,
27    dels: BTreeSet<Key>,
28    db: Db_<Key, Value, page::Page<Key, Value>>,
29    txn: MutTxn<Env, ()>,
30}
31
32const ROOT_DB: usize = 0;
33const PAGE_SIZE: u64 = 4096;
34
35impl<Key, Value> Database<Key, Value>
36where
37    Key: Ord + Clone + Debug + Storable,
38    Value: Storable + PartialEq,
39{
40    /// Open a database without a lock file where only one instance is safe to open.
41    pub fn open(path: PathBuf) -> Result<Self, Error> {
42        let env = unsafe { Env::new_nolock(&path, PAGE_SIZE, 1)? };
43
44        let mut txn = Env::mut_txn_begin(env)?;
45
46        let db = txn
47            .root_db(ROOT_DB)
48            .unwrap_or_else(|| unsafe { btree::create_db_(&mut txn).unwrap() });
49
50        Ok(Self {
51            path,
52            puts: BTreeMap::default(),
53            dels: BTreeSet::default(),
54            db,
55            txn,
56        })
57    }
58
59    #[inline]
60    pub fn get(&self, key: &Key) -> Option<&Value> {
61        if let Some(cached_put) = self.get_from_ram(key) {
62            return Some(cached_put);
63        }
64
65        self.get_from_disk(key)
66    }
67
68    /// Get only from the uncommited tree (ram) without checking the database (disk)
69    #[inline]
70    pub fn get_from_ram(&self, key: &Key) -> Option<&Value> {
71        self.puts.get(key)
72    }
73
74    /// Get mut only from the uncommited tree (ram) without checking the database (disk)
75    #[inline]
76    pub fn get_mut_from_ram(&mut self, key: &Key) -> Option<&mut Value> {
77        self.puts.get_mut(key)
78    }
79
80    /// Get only from the database (disk) without checking the uncommited tree (ram)
81    #[inline]
82    pub fn get_from_disk(&self, key: &Key) -> Option<&Value> {
83        let option = btree::get(&self.txn, &self.db, key, None).unwrap();
84
85        if let Some((key_found, v)) = option {
86            if key == key_found {
87                return Some(v);
88            }
89        }
90
91        None
92    }
93
94    #[inline]
95    pub fn insert(&mut self, key: Key, value: Value) -> Option<Value> {
96        self.dels.remove(&key);
97        self.insert_to_ram(key, value)
98    }
99
100    /// Insert without removing the key to the dels tree, so be sure that it hasn't added to the delete set
101    #[inline]
102    pub fn insert_to_ram(&mut self, key: Key, value: Value) -> Option<Value> {
103        self.puts.insert(key, value)
104    }
105
106    #[inline]
107    pub fn update(&mut self, key: Key, value: Value) -> Option<Value> {
108        self.dels.insert(key.clone());
109        self.puts.insert(key, value)
110    }
111
112    #[inline]
113    pub fn remove(&mut self, key: &Key) -> Option<Value> {
114        self.remove_from_ram(key).or_else(|| {
115            self.remove_later_from_disk(key);
116
117            None
118        })
119    }
120
121    /// Get only from the uncommited tree (ram) without checking the database (disk)
122    #[inline]
123    pub fn remove_from_ram(&mut self, key: &Key) -> Option<Value> {
124        self.puts.remove(key)
125    }
126
127    /// Add the key only to the dels tree without checking if it's present in the puts tree, only use if you are positive that you neither added nor updated an entry with this key
128    #[inline]
129    pub fn remove_later_from_disk(&mut self, key: &Key) {
130        self.dels.insert(key.clone());
131    }
132
133    #[inline]
134    pub fn is_empty(&self) -> bool {
135        self.iter_disk().next().is_none()
136    }
137
138    /// Iterate over key/value pairs from the uncommited tree (ram)
139    #[inline]
140    pub fn iter_ram(&self) -> std::collections::btree_map::Iter<'_, Key, Value> {
141        self.puts.iter()
142    }
143
144    /// Iterate over key/value pairs from the database (disk)
145    #[inline]
146    pub fn iter_disk(
147        &self,
148    ) -> btree::Iter<'_, MutTxn<Env, ()>, Key, Value, page::Page<Key, Value>> {
149        btree::iter(&self.txn, &self.db, None).unwrap()
150    }
151
152    /// Iterate over key/value pairs
153    #[inline]
154    pub fn iter_ram_then_disk(&self) -> impl Iterator<Item = (&Key, &Value)> {
155        self.iter_ram().chain(self.iter_disk().map(|r| r.unwrap()))
156    }
157
158    /// Collect a **clone** of all uncommited key/value pairs (ram)
159    pub fn collect_ram(&self) -> BTreeMap<Key, Value>
160    where
161        Value: Clone,
162    {
163        self.puts.clone()
164    }
165
166    /// Collect a **clone** of all key/value pairs from the database (disk)
167    pub fn collect_disk(&self) -> BTreeMap<Key, Value>
168    where
169        Value: Clone,
170    {
171        self.iter_disk()
172            .map(|r| r.unwrap())
173            .map(|(key, value)| (key.clone(), value.clone()))
174            .collect::<_>()
175    }
176}
177
178pub trait AnyDatabase {
179    #[allow(unused)]
180    fn export(self, defragment: bool) -> Result<(), Error>;
181    fn boxed_export(self: Box<Self>, defragment: bool) -> Result<(), Error>;
182    #[allow(unused)]
183    fn destroy(self) -> io::Result<()>;
184}
185
186impl<Key, Value> AnyDatabase for Database<Key, Value>
187where
188    Key: Ord + Clone + Debug + Storable,
189    Value: Storable + PartialEq + Clone,
190{
191    /// Flush all puts and dels from the ram to disk with an option to defragment the database to save some disk space
192    ///
193    /// /!\ Do not kill the program while this function is runnning  /!\
194    fn export(self, defragment: bool) -> Result<(), Error> {
195        Box::new(self).boxed_export(defragment)
196    }
197
198    /// Flush all puts and dels from the ram to disk with an option to defragment the database to save some disk space
199    ///
200    /// /!\ Do not kill the program while this function is runnning  /!\
201    fn boxed_export(mut self: Box<Self>, defragment: bool) -> Result<(), Error> {
202        if defragment {
203            let mut btree = self.as_ref().collect_disk();
204
205            let path = self.path.to_owned();
206            self.dels.iter().for_each(|key| {
207                btree.remove(key);
208            });
209            btree.append(&mut self.puts);
210
211            self.destroy()?;
212
213            *self = Self::open(path).unwrap();
214
215            if !self.is_empty() {
216                panic!()
217            }
218
219            self.puts = btree;
220        }
221
222        if self.dels.is_empty() && self.puts.is_empty() {
223            return Ok(());
224        }
225
226        mem::take(&mut self.dels)
227            .into_iter()
228            .try_for_each(|key| -> Result<(), Error> {
229                btree::del(&mut self.txn, &mut self.db, &key, None)?;
230                Ok(())
231            })?;
232
233        mem::take(&mut self.puts).into_iter().try_for_each(
234            |(key, value)| -> Result<(), Error> {
235                btree::put(&mut self.txn, &mut self.db, &key, &value)?;
236                Ok(())
237            },
238        )?;
239
240        self.txn.set_root(ROOT_DB, self.db.db.into());
241
242        self.txn.commit()
243    }
244
245    fn destroy(self) -> io::Result<()> {
246        let path = self.path.to_owned();
247
248        drop(self);
249
250        fs::remove_file(&path)
251    }
252}
snkrj/lib.rs

snkrj/
lib.rs
