Skip to main content

forest/db/
memory.rs

1// Copyright 2019-2026 ChainSafe Systems
2// SPDX-License-Identifier: Apache-2.0, MIT
3
4use super::{EthMappingsStore, SettingsStore, SettingsStoreExt};
5use crate::blocks::TipsetKey;
6use crate::db::PersistentStore;
7use crate::libp2p_bitswap::{BitswapStoreRead, BitswapStoreReadWrite};
8use crate::prelude::*;
9use crate::rpc::eth::types::EthHash;
10use crate::shim::clock::ChainEpoch;
11use crate::utils::db::car_stream::CarBlock;
12use crate::utils::multihash::prelude::*;
13use ahash::HashMap;
14use indexmap::IndexMap;
15use nunny::Vec as NonEmpty;
16use parking_lot::RwLock;
17
18#[derive(Debug, Default)]
19pub struct MemoryDB {
20    blockchain_db: RwLock<HashMap<Cid, Vec<u8>>>,
21    blockchain_persistent_db: RwLock<HashMap<Cid, Vec<u8>>>,
22    settings_db: RwLock<HashMap<String, Vec<u8>>>,
23    pub eth_mappings_db: RwLock<HashMap<EthHash, Vec<u8>>>,
24    pub ts_lookup_db: RwLock<HashMap<ChainEpoch, TipsetKey>>,
25}
26
27impl MemoryDB {
28    pub fn blockstore_len(&self) -> usize {
29        self.blockchain_db.read().len() + self.blockchain_persistent_db.read().len()
30    }
31
32    pub fn blockstore_size_bytes(&self) -> usize {
33        self.blockchain_db
34            .read()
35            .iter()
36            .chain(self.blockchain_persistent_db.read().iter())
37            .map(|(k, v)| k.to_bytes().len() + v.len())
38            .sum()
39    }
40
41    pub async fn export_forest_car<W: tokio::io::AsyncWrite + Unpin>(
42        &self,
43        writer: &mut W,
44    ) -> anyhow::Result<()> {
45        let roots =
46            SettingsStoreExt::read_obj::<TipsetKey>(self, crate::db::setting_keys::HEAD_KEY)?
47                .context("chain head is not tracked and cannot be exported")?
48                .into_cids();
49        self.export_forest_car_with_roots(roots, writer).await
50    }
51
52    pub async fn export_forest_car_with_roots<W: tokio::io::AsyncWrite + Unpin>(
53        &self,
54        roots: NonEmpty<Cid>,
55        writer: &mut W,
56    ) -> anyhow::Result<()> {
57        let blocks = {
58            let blockchain_db = self.blockchain_db.read();
59            let blockchain_persistent_db = self.blockchain_persistent_db.read();
60            blockchain_db
61                .iter()
62                .chain(blockchain_persistent_db.iter())
63                // Sort to make the result CAR deterministic
64                .sorted_by_key(|&(&cid, _)| cid)
65                .map(|(&cid, data)| {
66                    anyhow::Ok(CarBlock {
67                        cid,
68                        data: data.clone().into(),
69                    })
70                })
71                .collect_vec()
72        };
73        let frames =
74            crate::db::car::forest::Encoder::compress_stream_default(futures::stream::iter(blocks));
75        crate::db::car::forest::Encoder::write(writer, roots, frames).await
76    }
77}
78
79impl SettingsStore for MemoryDB {
80    fn read_bin(&self, key: &str) -> anyhow::Result<Option<Vec<u8>>> {
81        Ok(self.settings_db.read().get(key).cloned())
82    }
83
84    fn write_bin(&self, key: &str, value: &[u8]) -> anyhow::Result<()> {
85        self.settings_db
86            .write()
87            .insert(key.to_owned(), value.to_vec());
88        Ok(())
89    }
90
91    fn exists(&self, key: &str) -> anyhow::Result<bool> {
92        Ok(self.settings_db.read().contains_key(key))
93    }
94
95    fn setting_keys(&self) -> anyhow::Result<Vec<String>> {
96        Ok(self.settings_db.read().keys().cloned().collect_vec())
97    }
98}
99
100impl EthMappingsStore for MemoryDB {
101    fn read_bin(&self, key: &EthHash) -> anyhow::Result<Option<Vec<u8>>> {
102        Ok(self.eth_mappings_db.read().get(key).cloned())
103    }
104
105    fn write_bin(&self, key: &EthHash, value: &[u8]) -> anyhow::Result<()> {
106        self.eth_mappings_db
107            .write()
108            .insert(key.to_owned(), value.to_vec());
109        Ok(())
110    }
111
112    fn exists(&self, key: &EthHash) -> anyhow::Result<bool> {
113        Ok(self.eth_mappings_db.read().contains_key(key))
114    }
115
116    fn get_message_cids(&self) -> anyhow::Result<Vec<(Cid, u64)>> {
117        let cids = self
118            .eth_mappings_db
119            .read()
120            .values()
121            .filter_map(|value| fvm_ipld_encoding::from_slice::<(Cid, u64)>(value).ok())
122            .collect();
123
124        Ok(cids)
125    }
126
127    fn delete(&self, keys: Vec<EthHash>) -> anyhow::Result<()> {
128        let mut lock = self.eth_mappings_db.write();
129        for hash in keys.iter() {
130            lock.remove(hash);
131        }
132        Ok(())
133    }
134
135    fn tipset_key_by_epoch(&self, epoch: ChainEpoch) -> anyhow::Result<Option<TipsetKey>> {
136        Ok(self.ts_lookup_db.read().get(&epoch).cloned())
137    }
138
139    fn delete_tipset_key_at_epoch(&self, epoch: ChainEpoch) -> anyhow::Result<()> {
140        self.ts_lookup_db.write().remove(&epoch);
141        Ok(())
142    }
143
144    fn set_tipset_key_at_epoch_raw(
145        &self,
146        epoch: ChainEpoch,
147        tsk: &TipsetKey,
148    ) -> anyhow::Result<()> {
149        self.ts_lookup_db.write().insert(epoch, tsk.clone());
150        Ok(())
151    }
152}
153
154impl Blockstore for MemoryDB {
155    fn get(&self, k: &Cid) -> anyhow::Result<Option<Vec<u8>>> {
156        Ok(self.blockchain_db.read().get(k).cloned().or(self
157            .blockchain_persistent_db
158            .read()
159            .get(k)
160            .cloned()))
161    }
162
163    fn put_keyed(&self, k: &Cid, block: &[u8]) -> anyhow::Result<()> {
164        self.blockchain_db.write().insert(*k, block.to_vec());
165        Ok(())
166    }
167}
168
169impl PersistentStore for MemoryDB {
170    fn put_keyed_persistent(&self, k: &Cid, block: &[u8]) -> anyhow::Result<()> {
171        self.blockchain_persistent_db
172            .write()
173            .insert(*k, block.to_vec());
174        Ok(())
175    }
176}
177
178impl BitswapStoreRead for MemoryDB {
179    fn contains(&self, cid: &Cid) -> anyhow::Result<bool> {
180        Ok(self.blockchain_db.read().contains_key(cid))
181    }
182
183    fn get(&self, cid: &Cid) -> anyhow::Result<Option<Vec<u8>>> {
184        Blockstore::get(self, cid)
185    }
186}
187
188impl BitswapStoreReadWrite for MemoryDB {
189    type Hashes = MultihashCode;
190
191    fn insert(&self, block: &crate::libp2p_bitswap::Block64<Self::Hashes>) -> anyhow::Result<()> {
192        self.put_keyed(block.cid(), block.data())
193    }
194}
195
196impl super::HeaviestTipsetKeyProvider for MemoryDB {
197    fn heaviest_tipset_key(&self) -> anyhow::Result<Option<TipsetKey>> {
198        SettingsStoreExt::read_obj::<TipsetKey>(self, crate::db::setting_keys::HEAD_KEY)
199    }
200
201    fn set_heaviest_tipset_key(&self, tsk: &TipsetKey) -> anyhow::Result<()> {
202        SettingsStoreExt::write_obj(self, crate::db::setting_keys::HEAD_KEY, tsk)
203    }
204}
205
206#[derive(Debug, Default, derive_more::Deref)]
207/// A memory blockstore that preserves the insertion order
208pub struct IndexMapBlockstore {
209    inner: RwLock<IndexMap<Cid, Vec<u8>>>,
210}
211
212impl IndexMapBlockstore {
213    #[allow(dead_code)]
214    pub async fn export_forest_car<W: tokio::io::AsyncWrite + Unpin>(
215        &self,
216        roots: NonEmpty<Cid>,
217        writer: &mut W,
218    ) -> anyhow::Result<()> {
219        let blocks = {
220            let inner = self.inner.read();
221            let invalid_roots = roots
222                .iter()
223                .filter(|&c| !inner.contains_key(c))
224                .collect_vec();
225            anyhow::ensure!(
226                invalid_roots.is_empty(),
227                "All roots should present in the blockstore, invalid roots: {invalid_roots:?}"
228            );
229            inner
230                .iter()
231                .map(|(&cid, data)| {
232                    anyhow::Ok(CarBlock {
233                        cid,
234                        data: data.clone().into(),
235                    })
236                })
237                .collect_vec()
238        };
239        let frames =
240            crate::db::car::forest::Encoder::compress_stream_default(futures::stream::iter(blocks));
241        crate::db::car::forest::Encoder::write(writer, roots, frames).await
242    }
243}
244
245impl Blockstore for IndexMapBlockstore {
246    fn get(&self, k: &Cid) -> anyhow::Result<Option<Vec<u8>>> {
247        Ok(self.read().get(k).cloned())
248    }
249
250    fn put_keyed(&self, k: &Cid, block: &[u8]) -> anyhow::Result<()> {
251        self.write().insert(*k, block.to_vec());
252        Ok(())
253    }
254}
255
256#[cfg(test)]
257mod tests {
258    use super::*;
259    use crate::{
260        db::{car::ForestCar, setting_keys::HEAD_KEY},
261        utils::cid::CidCborExt as _,
262    };
263    use fil_actors_shared::fvm_ipld_hamt::Hamt;
264    use fvm_ipld_encoding::DAG_CBOR;
265    use multihash_codetable::Code::Blake2b256;
266    use nunny::vec as nonempty;
267
268    #[tokio::test]
269    async fn test_export_forest_car() {
270        let db = MemoryDB::default();
271        let record1 = b"non-persistent";
272        let key1 = Cid::new_v1(DAG_CBOR, Blake2b256.digest(record1.as_slice()));
273        db.put_keyed(&key1, record1.as_slice()).unwrap();
274
275        let record2 = b"persistent";
276        let key2 = Cid::new_v1(DAG_CBOR, Blake2b256.digest(record2.as_slice()));
277        db.put_keyed_persistent(&key2, record2.as_slice()).unwrap();
278
279        let mut car_db_bytes = vec![];
280        assert!(
281            db.export_forest_car(&mut car_db_bytes)
282                .await
283                .unwrap_err()
284                .to_string()
285                .contains("chain head is not tracked and cannot be exported")
286        );
287
288        db.write_obj(HEAD_KEY, &TipsetKey::from(nonempty![key1]))
289            .unwrap();
290
291        car_db_bytes.clear();
292        db.export_forest_car(&mut car_db_bytes).await.unwrap();
293
294        let car = ForestCar::new(car_db_bytes).unwrap();
295        assert_eq!(car.head_tipset_key(), &nonempty![key1]);
296        assert!(car.has(&key1).unwrap());
297        assert!(car.has(&key2).unwrap());
298    }
299
300    #[tokio::test]
301    async fn test_index_map_blockstore() {
302        const BIT_WIDTH: u32 = 5;
303
304        let db = IndexMapBlockstore::default();
305        // similate tipset lookup table
306        let mut hamt: Hamt<_, TipsetKey, ChainEpoch> = Hamt::new_with_bit_width(&db, BIT_WIDTH);
307        let checkpoints = [
308            (
309                0,
310                TipsetKey::from(nunny::vec![Cid::from_cbor_blake2b256(&"1").unwrap()]),
311            ),
312            (
313                5,
314                TipsetKey::from(nunny::vec![Cid::from_cbor_blake2b256(&"5").unwrap()]),
315            ),
316            (
317                10,
318                TipsetKey::from(nunny::vec![Cid::from_cbor_blake2b256(&"10").unwrap()]),
319            ),
320        ];
321        for (epoch, tsk) in checkpoints.iter().cloned() {
322            hamt.set(epoch, tsk).unwrap();
323        }
324        let hamt_root = hamt.flush().unwrap();
325        assert!(db.has(&hamt_root).unwrap(), "hamt root should present");
326
327        // export with invalid roots should fail
328        let mut car = vec![];
329        db.export_forest_car(
330            nunny::vec![Cid::from_cbor_blake2b256(&"invalid").unwrap()],
331            &mut car,
332        )
333        .await
334        .unwrap_err();
335
336        // export with hamt root should succeed
337        let mut car_bytes = vec![];
338        let car_roots = nunny::vec![hamt_root];
339        db.export_forest_car(car_roots.clone(), &mut car_bytes)
340            .await
341            .unwrap();
342
343        let car: ForestCar<Vec<u8>> = ForestCar::new(car_bytes).unwrap();
344        let hamt_from_car: Hamt<_, TipsetKey, ChainEpoch> =
345            Hamt::load_with_bit_width(&hamt_root, &car, BIT_WIDTH).unwrap();
346
347        let checkpoints_from_memdb_hamt = {
348            let mut v = vec![];
349            hamt.for_each_cacheless(|epoch, tsk| {
350                v.push((*epoch, tsk.clone()));
351                anyhow::Ok(())
352            })
353            .unwrap();
354            v
355        };
356        let checkpoints_from_car_hamt = {
357            let mut v = vec![];
358            hamt_from_car
359                .for_each_cacheless(|epoch, tsk| {
360                    v.push((*epoch, tsk.clone()));
361                    anyhow::Ok(())
362                })
363                .unwrap();
364            v
365        };
366        // Cannot compare the 2 hamts as they use different DB types
367        assert_eq!(checkpoints_from_memdb_hamt, checkpoints_from_car_hamt);
368        // The hamt iteration order is different from the insertion order
369        assert_eq!(
370            HashMap::from_iter(checkpoints),
371            HashMap::from_iter(checkpoints_from_car_hamt)
372        );
373    }
374}