forest/db/
blockstore_with_write_buffer.rs

1// Copyright 2019-2026 ChainSafe Systems
2// SPDX-License-Identifier: Apache-2.0, MIT
3
4use crate::prelude::*;
5use ahash::HashMap;
6use parking_lot::RwLock;
7
8pub struct BlockstoreWithWriteBuffer<DB: Blockstore> {
9    inner: DB,
10    buffer: RwLock<HashMap<Cid, Vec<u8>>>,
11    buffer_capacity: usize,
12}
13
14impl<DB: Blockstore> Blockstore for BlockstoreWithWriteBuffer<DB> {
15    fn get(&self, k: &Cid) -> anyhow::Result<Option<Vec<u8>>> {
16        if let Some(v) = self.buffer.read().get(k) {
17            return Ok(Some(v.clone()));
18        }
19        self.inner.get(k)
20    }
21
22    fn has(&self, k: &Cid) -> anyhow::Result<bool> {
23        Ok(self.buffer.read().contains_key(k) || self.inner.has(k)?)
24    }
25
26    fn put_keyed(&self, k: &Cid, block: &[u8]) -> anyhow::Result<()> {
27        {
28            let mut buffer = self.buffer.write();
29            buffer.insert(*k, block.to_vec());
30        }
31        self.flush_buffer_if_needed()
32    }
33}
34
35impl<DB: Blockstore> BlockstoreWithWriteBuffer<DB> {
36    pub fn new_with_capacity(inner: DB, buffer_capacity: usize) -> Self {
37        Self {
38            inner,
39            buffer_capacity,
40            buffer: RwLock::new(HashMap::with_capacity(buffer_capacity)),
41        }
42    }
43
44    fn flush_buffer(&self) -> anyhow::Result<()> {
45        let records = {
46            let mut buffer = self.buffer.write();
47            buffer.drain().collect_vec()
48        };
49        self.inner.put_many_keyed(records)
50    }
51
52    fn flush_buffer_if_needed(&self) -> anyhow::Result<()> {
53        if self.buffer.read().len() >= self.buffer_capacity {
54            self.flush_buffer()
55        } else {
56            Ok(())
57        }
58    }
59}
60
61impl<DB: Blockstore> Drop for BlockstoreWithWriteBuffer<DB> {
62    fn drop(&mut self) {
63        if let Err(e) = self.flush_buffer() {
64            tracing::warn!("{e:#}");
65        }
66    }
67}
68
69#[cfg(test)]
70mod tests {
71    use super::*;
72    use crate::{db::MemoryDB, utils::rand::forest_rng};
73    use fvm_ipld_encoding::DAG_CBOR;
74    use multihash_codetable::Code::Blake2b256;
75    use multihash_codetable::MultihashDigest as _;
76    use rand::Rng as _;
77    use std::sync::Arc;
78
79    #[test]
80    fn test_buffer_flush() {
81        const BUFFER_SIZE: usize = 10;
82        const N_RECORDS: usize = 15;
83        let mem_db = Arc::new(MemoryDB::default());
84        let buf_db = BlockstoreWithWriteBuffer::new_with_capacity(mem_db.clone(), BUFFER_SIZE);
85        let mut records = Vec::with_capacity(N_RECORDS);
86        for _ in 0..N_RECORDS {
87            let mut record = [0; 1024];
88            forest_rng().fill(&mut record);
89            let key = Cid::new_v1(DAG_CBOR, Blake2b256.digest(record.as_slice()));
90            records.push((key, record));
91        }
92
93        buf_db.put_many_keyed(records.clone()).unwrap();
94
95        for (i, (k, v)) in records.iter().enumerate() {
96            assert!(buf_db.has(k).unwrap());
97            assert_eq!(buf_db.get(k).unwrap().unwrap().as_slice(), v);
98            if i < BUFFER_SIZE {
99                assert!(mem_db.has(k).unwrap());
100                assert_eq!(mem_db.get(k).unwrap().unwrap().as_slice(), v);
101            } else {
102                assert!(!mem_db.has(k).unwrap());
103            }
104        }
105
106        drop(buf_db);
107
108        for (k, v) in records.iter() {
109            assert!(mem_db.has(k).unwrap());
110            assert_eq!(mem_db.get(k).unwrap().unwrap().as_slice(), v);
111        }
112    }
113}
forest/db/blockstore_with_write_buffer.rs

forest/db/
blockstore_with_write_buffer.rs
