ethrex_storage/backend/
in_memory.rs

1use crate::api::{
2    PrefixResult, StorageBackend, StorageLockedView, StorageReadView, StorageWriteBatch,
3};
4use crate::error::StoreError;
5use rustc_hash::FxHashMap;
6use std::path::Path;
7use std::sync::{Arc, RwLock};
8
9type Table = FxHashMap<Vec<u8>, Vec<u8>>;
10type Database = FxHashMap<&'static str, Table>;
11
12#[derive(Debug)]
13pub struct InMemoryBackend {
14    // RCU-style snapshot store: readers clone the inner Arc and then read lock-free.
15    // Writes run under the outer write lock and use Arc::make_mut for copy-on-write.
16    // If read snapshots are still alive, writes may clone the full Database.
17    inner: Arc<RwLock<Arc<Database>>>,
18}
19
20impl InMemoryBackend {
21    pub fn open() -> Result<Self, StoreError> {
22        Ok(Self {
23            inner: Arc::new(RwLock::new(Arc::new(Database::default()))),
24        })
25    }
26}
27
28impl StorageBackend for InMemoryBackend {
29    fn clear_table(&self, table: &str) -> Result<(), StoreError> {
30        let mut db = self
31            .inner
32            .write()
33            .map_err(|_| StoreError::Custom("Failed to acquire write lock".to_string()))?;
34
35        let db_mut = Arc::make_mut(&mut *db);
36        if let Some(table_ref) = db_mut.get_mut(table) {
37            table_ref.clear();
38        }
39        Ok(())
40    }
41
42    fn begin_read(&self) -> Result<Arc<dyn StorageReadView>, StoreError> {
43        let snapshot = self
44            .inner
45            .read()
46            .map_err(|_| StoreError::Custom("Failed to acquire read lock".to_string()))?
47            .clone();
48        Ok(Arc::new(InMemoryReadTx { snapshot }))
49    }
50
51    fn begin_write(&self) -> Result<Box<dyn StorageWriteBatch + 'static>, StoreError> {
52        Ok(Box::new(InMemoryWriteTx {
53            backend: self.inner.clone(),
54        }))
55    }
56
57    fn begin_locked(
58        &self,
59        table_name: &'static str,
60    ) -> Result<Box<dyn StorageLockedView>, StoreError> {
61        let snapshot = self
62            .inner
63            .read()
64            .map_err(|_| StoreError::Custom("Failed to acquire read lock".to_string()))?
65            .clone();
66        Ok(Box::new(InMemoryLocked {
67            snapshot,
68            table_name,
69        }))
70    }
71
72    fn create_checkpoint(&self, _path: &Path) -> Result<(), StoreError> {
73        // Checkpoints are not supported for the InMemory DB
74        // Silently ignoring the request to create a checkpoint is harmless
75        Ok(())
76    }
77}
78
79pub struct InMemoryLocked {
80    snapshot: Arc<Database>,
81    table_name: &'static str,
82}
83
84impl StorageLockedView for InMemoryLocked {
85    fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>, StoreError> {
86        Ok(self
87            .snapshot
88            .get(&self.table_name)
89            .and_then(|table_ref| table_ref.get(key))
90            .cloned())
91    }
92}
93
94pub struct InMemoryReadTx {
95    snapshot: Arc<Database>,
96}
97
98impl StorageReadView for InMemoryReadTx {
99    fn get(&self, table: &str, key: &[u8]) -> Result<Option<Vec<u8>>, StoreError> {
100        Ok(self
101            .snapshot
102            .get(table)
103            .and_then(|table_ref| table_ref.get(key))
104            .cloned())
105    }
106
107    fn prefix_iterator(
108        &self,
109        table: &str,
110        prefix: &[u8],
111    ) -> Result<Box<dyn Iterator<Item = PrefixResult> + '_>, StoreError> {
112        let table_data = self.snapshot.get(table).cloned().unwrap_or_default();
113
114        // Seek and walk forward, as RocksDB does: callers seek past the start of a prefix
115        // (`block_hash || start_index` for receipts) or across several prefixes (witness
116        // pruning) and stop on their own.
117        let mut entries: Vec<(Vec<u8>, Vec<u8>)> = table_data
118            .into_iter()
119            .filter(|(key, _)| key.as_slice() >= prefix)
120            .collect();
121        entries.sort_unstable_by(|(left, _), (right, _)| left.cmp(right));
122
123        let results: Vec<PrefixResult> = entries
124            .into_iter()
125            .map(|(k, v)| Ok((k.into_boxed_slice(), v.into_boxed_slice())))
126            .collect();
127
128        Ok(Box::new(results.into_iter()))
129    }
130
131    fn first_key(&self, table: &'static str) -> Result<Option<Vec<u8>>, StoreError> {
132        let Some(table_data) = self.snapshot.get(table) else {
133            return Ok(None);
134        };
135        Ok(table_data.keys().min().cloned())
136    }
137
138    fn last_key(&self, table: &'static str) -> Result<Option<Vec<u8>>, StoreError> {
139        let Some(table_data) = self.snapshot.get(table) else {
140            return Ok(None);
141        };
142        Ok(table_data.keys().max().cloned())
143    }
144}
145
146pub struct InMemoryWriteTx {
147    backend: Arc<RwLock<Arc<Database>>>,
148}
149
150impl StorageWriteBatch for InMemoryWriteTx {
151    fn put_batch(
152        &mut self,
153        table: &'static str,
154        batch: Vec<(Vec<u8>, Vec<u8>)>,
155    ) -> Result<(), StoreError> {
156        let mut db = self
157            .backend
158            .write()
159            .map_err(|_| StoreError::Custom("Failed to acquire write lock".to_string()))?;
160
161        // Copy-on-write update of the current snapshot.
162        let db_mut = Arc::make_mut(&mut *db);
163        let table_ref = db_mut.entry(table).or_default();
164
165        for (key, value) in batch {
166            table_ref.insert(key, value);
167        }
168
169        Ok(())
170    }
171
172    fn delete(&mut self, table: &str, key: &[u8]) -> Result<(), StoreError> {
173        let mut db = self
174            .backend
175            .write()
176            .map_err(|_| StoreError::Custom("Failed to acquire write lock".to_string()))?;
177
178        let db_mut = Arc::make_mut(&mut *db);
179        if let Some(table_ref) = db_mut.get_mut(table) {
180            table_ref.remove(key);
181        }
182        Ok(())
183    }
184
185    fn delete_range(
186        &mut self,
187        table: &'static str,
188        start: &[u8],
189        end: &[u8],
190    ) -> Result<(), StoreError> {
191        let mut db = self
192            .backend
193            .write()
194            .map_err(|_| StoreError::Custom("Failed to acquire write lock".to_string()))?;
195
196        let db_mut = Arc::make_mut(&mut *db);
197        if let Some(table_ref) = db_mut.get_mut(table) {
198            table_ref.retain(|k, _| !(k.as_slice() >= start && k.as_slice() < end));
199        }
200        Ok(())
201    }
202
203    fn merge(&mut self, table: &'static str, key: &[u8], operand: &[u8]) -> Result<(), StoreError> {
204        // InMemory has no native merge operator, so apply the merge inline.
205        // Only TRANSACTION_LOCATIONS uses merge today; dispatch by table.
206        if table != crate::api::tables::TRANSACTION_LOCATIONS {
207            return Err(StoreError::Custom(format!(
208                "merge not supported for table {table}"
209            )));
210        }
211        let mut db = self
212            .backend
213            .write()
214            .map_err(|_| StoreError::Custom("Failed to acquire write lock".to_string()))?;
215        let db_mut = Arc::make_mut(&mut *db);
216        let table_ref = db_mut.entry(table).or_default();
217        let existing = table_ref.get(key).map(|v| v.as_slice());
218        let merged = crate::store::tx_locations_merge(existing, std::iter::once(operand))
219            .ok_or_else(|| StoreError::Custom("tx_locations_merge returned None".to_string()))?;
220        table_ref.insert(key.to_vec(), merged);
221        Ok(())
222    }
223
224    fn commit(&mut self) -> Result<(), StoreError> {
225        // NOTE: every `put`, `delete`, and `delete_range` above mutates the live
226        // `Arc<Database>` immediately under the write lock, so `commit` is a no-op
227        // and multi-op sequences (e.g. the journal entry + trie writes in
228        // `commit_to_disk`, or the `delete_range` + finalized-number update in
229        // `forkchoice_update_inner`) are not atomic. That's acceptable here: this
230        // backend is RAM-backed (dev/test only), and atomicity only guards crash
231        // recovery — a process death loses all in-memory state anyway, so a
232        // half-applied batch is never observable.
233        Ok(())
234    }
235}
236
237#[cfg(test)]
238mod tests {
239    use super::*;
240
241    /// Like the RocksDB backend, which configures no prefix extractor, the iterator starts
242    /// at the first key not less than the one given and runs through the rest of the
243    /// table in key order.
244    #[test]
245    fn prefix_iterator_seeks_and_walks_forward() {
246        let backend = InMemoryBackend::open().unwrap();
247        let mut write = backend.begin_write().unwrap();
248        let keys = ["b1", "a2", "c", "a1"].map(|key| (key.as_bytes().to_vec(), vec![]));
249        write.put_batch("table", keys.to_vec()).unwrap();
250        write.commit().unwrap();
251
252        let read = backend.begin_read().unwrap();
253        let keys_from = |start: &[u8]| {
254            read.prefix_iterator("table", start)
255                .unwrap()
256                .map(|item| String::from_utf8(item.unwrap().0.to_vec()).unwrap())
257                .collect::<Vec<_>>()
258        };
259        assert_eq!(keys_from(b""), ["a1", "a2", "b1", "c"]);
260        assert_eq!(keys_from(b"a2"), ["a2", "b1", "c"]);
261        assert_eq!(keys_from(b"a3"), ["b1", "c"]);
262        assert_eq!(keys_from(b"d"), Vec::<String>::new());
263    }
264}
ethrex_storage/backend/in_memory.rs

ethrex_storage/backend/
in_memory.rs
