Skip to main content

fedimint_cursed_redb/
lib.rs

1//! Uses immutable data structures and saves to redb on commit.
2
3use std::fmt::Debug;
4use std::ops::Range;
5use std::path::Path;
6use std::sync::{Arc, Mutex};
7
8use fedimint_core::db::{
9    DatabaseError, DatabaseResult, IDatabaseTransactionOps, IDatabaseTransactionOpsCore,
10    IRawDatabase, IRawDatabaseTransaction, PrefixStream,
11};
12use fedimint_core::{apply, async_trait_maybe_send};
13use futures::stream;
14use imbl::OrdMap;
15use redb::{Database, ReadableDatabase, ReadableTable, TableDefinition};
16
17pub use crate::error::MemAndRedbOpenError;
18
19const KV_TABLE: TableDefinition<&[u8], &[u8]> = TableDefinition::new("fedimint_kv");
20
21#[derive(Debug, Default)]
22pub struct DatabaseInsertOperation {
23    pub key: Vec<u8>,
24    pub value: Vec<u8>,
25    pub old_value: Option<Vec<u8>>,
26}
27
28#[derive(Debug, Default)]
29pub struct DatabaseDeleteOperation {
30    pub key: Vec<u8>,
31    pub old_value: Option<Vec<u8>>,
32}
33
34#[derive(Debug)]
35pub enum DatabaseOperation {
36    Insert(DatabaseInsertOperation),
37    Delete(DatabaseDeleteOperation),
38}
39
40#[derive(Clone)]
41pub struct MemAndRedb {
42    data: Arc<Mutex<OrdMap<Vec<u8>, Vec<u8>>>>,
43    db: Arc<Database>,
44}
45
46impl Debug for MemAndRedb {
47    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
48        f.debug_struct("MemDatabase").finish_non_exhaustive()
49    }
50}
51
52#[derive(Debug)]
53pub struct MemAndRedbTransaction<'a> {
54    operations: Vec<DatabaseOperation>,
55    tx_data: OrdMap<Vec<u8>, Vec<u8>>,
56    db: &'a MemAndRedb,
57}
58
59mod error;
60
61#[cfg(not(target_family = "wasm"))]
62mod native;
63
64#[cfg(any(test, target_family = "wasm"))]
65mod read_exact;
66
67#[cfg(target_family = "wasm")]
68mod wasm;
69
70impl MemAndRedb {
71    fn new_from_redb(db: Database) -> DatabaseResult<Self> {
72        let db = Arc::new(db);
73        let mut data = OrdMap::new();
74
75        // Load existing data from redb
76        let read_txn = db.begin_read().map_err(DatabaseError::backend)?;
77        if let Ok(table) = read_txn.open_table(KV_TABLE) {
78            for entry in table.iter().map_err(DatabaseError::backend)? {
79                let (key, value) = entry.map_err(DatabaseError::backend)?;
80                data.insert(key.value().to_vec(), value.value().to_vec());
81            }
82        }
83        // Table might not exist on first run, which is fine
84
85        Ok(Self {
86            data: Arc::new(Mutex::new(data)),
87            db,
88        })
89    }
90}
91
92#[apply(async_trait_maybe_send!)]
93impl IRawDatabase for MemAndRedb {
94    type Transaction<'a> = MemAndRedbTransaction<'a>;
95
96    async fn begin_transaction<'a>(&'a self) -> MemAndRedbTransaction<'a> {
97        MemAndRedbTransaction {
98            operations: Vec::new(),
99            tx_data: {
100                let data_lock = self.data.lock().expect("poison");
101                data_lock.clone()
102            },
103            db: self,
104        }
105    }
106
107    fn checkpoint(&self, _: &Path) -> DatabaseResult<()> {
108        unimplemented!()
109    }
110}
111
112#[apply(async_trait_maybe_send!)]
113impl<'a> IDatabaseTransactionOpsCore for MemAndRedbTransaction<'a> {
114    async fn raw_insert_bytes(
115        &mut self,
116        key: &[u8],
117        value: &[u8],
118    ) -> DatabaseResult<Option<Vec<u8>>> {
119        let val = IDatabaseTransactionOpsCore::raw_get_bytes(self, key).await;
120        // Insert data from copy so we can read our own writes
121        let old_value = self.tx_data.insert(key.to_vec(), value.to_vec());
122        self.operations
123            .push(DatabaseOperation::Insert(DatabaseInsertOperation {
124                key: key.to_vec(),
125                value: value.to_vec(),
126                old_value,
127            }));
128        val
129    }
130
131    async fn raw_get_bytes(&mut self, key: &[u8]) -> DatabaseResult<Option<Vec<u8>>> {
132        Ok(self.tx_data.get(key).cloned())
133    }
134
135    async fn raw_remove_entry(&mut self, key: &[u8]) -> DatabaseResult<Option<Vec<u8>>> {
136        // Remove data from copy so we can read our own writes
137        let old_value = self.tx_data.remove(&key.to_vec());
138        self.operations
139            .push(DatabaseOperation::Delete(DatabaseDeleteOperation {
140                key: key.to_vec(),
141                old_value: old_value.clone(),
142            }));
143        Ok(old_value)
144    }
145
146    async fn raw_find_by_range(&mut self, range: Range<&[u8]>) -> DatabaseResult<PrefixStream<'_>> {
147        let data = self
148            .tx_data
149            .range::<_, Vec<u8>>(Range {
150                start: range.start.to_vec(),
151                end: range.end.to_vec(),
152            })
153            .map(|(key, value)| (key.clone(), value.clone()))
154            .collect::<Vec<_>>();
155
156        Ok(Box::pin(stream::iter(data)))
157    }
158
159    async fn raw_find_by_prefix(&mut self, key_prefix: &[u8]) -> DatabaseResult<PrefixStream<'_>> {
160        let data = self
161            .tx_data
162            .range::<_, Vec<u8>>((key_prefix.to_vec())..)
163            .take_while(|(key, _)| key.starts_with(key_prefix))
164            .map(|(key, value)| (key.clone(), value.clone()))
165            .collect::<Vec<_>>();
166
167        Ok(Box::pin(stream::iter(data)))
168    }
169
170    async fn raw_remove_by_prefix(&mut self, key_prefix: &[u8]) -> DatabaseResult<()> {
171        let keys = self
172            .tx_data
173            .range::<_, Vec<u8>>((key_prefix.to_vec())..)
174            .take_while(|(key, _)| key.starts_with(key_prefix))
175            .map(|(key, _)| key.clone())
176            .collect::<Vec<_>>();
177        for key in keys.iter() {
178            let old_value = self.tx_data.remove(&key.to_vec());
179            self.operations
180                .push(DatabaseOperation::Delete(DatabaseDeleteOperation {
181                    key: key.to_vec(),
182                    old_value,
183                }));
184        }
185        Ok(())
186    }
187
188    async fn raw_find_by_prefix_sorted_descending(
189        &mut self,
190        key_prefix: &[u8],
191    ) -> DatabaseResult<PrefixStream<'_>> {
192        let mut data = self
193            .tx_data
194            .range::<_, Vec<u8>>((key_prefix.to_vec())..)
195            .take_while(|(key, _)| key.starts_with(key_prefix))
196            .map(|(key, value)| (key.clone(), value.clone()))
197            .collect::<Vec<_>>();
198        data.sort_by(|a, b| a.cmp(b).reverse());
199
200        Ok(Box::pin(stream::iter(data)))
201    }
202}
203
204impl<'a> IDatabaseTransactionOps for MemAndRedbTransaction<'a> {}
205
206// In-memory database transaction should only be used for test code and never
207// for production as it doesn't properly implement MVCC
208#[apply(async_trait_maybe_send!)]
209impl<'a> IRawDatabaseTransaction for MemAndRedbTransaction<'a> {
210    async fn commit_tx(self) -> DatabaseResult<()> {
211        let mut data_locked = self.db.data.lock().expect("poison");
212        let write_txn = self.db.db.begin_write().map_err(DatabaseError::backend)?;
213        let operations = self.operations;
214        let mut data_new = data_locked.clone();
215        {
216            let mut table = write_txn
217                .open_table(KV_TABLE)
218                .map_err(DatabaseError::backend)?;
219
220            // Apply all operations
221            for op in operations {
222                match op {
223                    DatabaseOperation::Insert(insert_op) => {
224                        table
225                            .insert(&insert_op.key[..], &insert_op.value[..])
226                            .map_err(DatabaseError::backend)?;
227                        let old_value = data_new.insert(insert_op.key, insert_op.value);
228                        if old_value != insert_op.old_value {
229                            return Err(DatabaseError::WriteConflict);
230                        }
231                    }
232                    DatabaseOperation::Delete(delete_op) => {
233                        table
234                            .remove(&delete_op.key[..])
235                            .map_err(DatabaseError::backend)?;
236                        let old_value = data_new.remove(&delete_op.key);
237                        if old_value != delete_op.old_value {
238                            return Err(DatabaseError::WriteConflict);
239                        }
240                    }
241                }
242            }
243        }
244        // Commit redb transaction
245        write_txn.commit().map_err(DatabaseError::backend)?;
246
247        // Update in-memory data
248        *data_locked = data_new;
249        Ok(())
250    }
251}
252
253#[cfg(all(test, not(target_family = "wasm")))]
254mod tests;