1use super::{EthMappingsStore, SettingsStore, SettingsStoreExt};
5use crate::blocks::TipsetKey;
6use crate::db::PersistentStore;
7use crate::libp2p_bitswap::{BitswapStoreRead, BitswapStoreReadWrite};
8use crate::rpc::eth::types::EthHash;
9use crate::utils::db::car_stream::CarBlock;
10use crate::utils::multihash::prelude::*;
11use ahash::HashMap;
12use anyhow::Context as _;
13use cid::Cid;
14use fvm_ipld_blockstore::Blockstore;
15use itertools::Itertools;
16use nunny::Vec as NonEmpty;
17use parking_lot::RwLock;
18
19#[derive(Debug, Default)]
20pub struct MemoryDB {
21 blockchain_db: RwLock<HashMap<Cid, Vec<u8>>>,
22 blockchain_persistent_db: RwLock<HashMap<Cid, Vec<u8>>>,
23 settings_db: RwLock<HashMap<String, Vec<u8>>>,
24 pub eth_mappings_db: RwLock<HashMap<EthHash, Vec<u8>>>,
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 .sorted_by_key(|&(&cid, _)| cid)
65 .map(|(&cid, data)| {
66 anyhow::Ok(CarBlock {
67 cid,
68 data: data.clone(),
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
136impl Blockstore for MemoryDB {
137 fn get(&self, k: &Cid) -> anyhow::Result<Option<Vec<u8>>> {
138 Ok(self.blockchain_db.read().get(k).cloned().or(self
139 .blockchain_persistent_db
140 .read()
141 .get(k)
142 .cloned()))
143 }
144
145 fn put_keyed(&self, k: &Cid, block: &[u8]) -> anyhow::Result<()> {
146 self.blockchain_db.write().insert(*k, block.to_vec());
147 Ok(())
148 }
149}
150
151impl PersistentStore for MemoryDB {
152 fn put_keyed_persistent(&self, k: &Cid, block: &[u8]) -> anyhow::Result<()> {
153 self.blockchain_persistent_db
154 .write()
155 .insert(*k, block.to_vec());
156 Ok(())
157 }
158}
159
160impl BitswapStoreRead for MemoryDB {
161 fn contains(&self, cid: &Cid) -> anyhow::Result<bool> {
162 Ok(self.blockchain_db.read().contains_key(cid))
163 }
164
165 fn get(&self, cid: &Cid) -> anyhow::Result<Option<Vec<u8>>> {
166 Blockstore::get(self, cid)
167 }
168}
169
170impl BitswapStoreReadWrite for MemoryDB {
171 type Hashes = MultihashCode;
172
173 fn insert(&self, block: &crate::libp2p_bitswap::Block64<Self::Hashes>) -> anyhow::Result<()> {
174 self.put_keyed(block.cid(), block.data())
175 }
176}
177
178impl super::HeaviestTipsetKeyProvider for MemoryDB {
179 fn heaviest_tipset_key(&self) -> anyhow::Result<Option<TipsetKey>> {
180 SettingsStoreExt::read_obj::<TipsetKey>(self, crate::db::setting_keys::HEAD_KEY)
181 }
182
183 fn set_heaviest_tipset_key(&self, tsk: &TipsetKey) -> anyhow::Result<()> {
184 SettingsStoreExt::write_obj(self, crate::db::setting_keys::HEAD_KEY, tsk)
185 }
186}
187
188#[cfg(test)]
189mod tests {
190 use super::*;
191 use crate::db::{car::ForestCar, setting_keys::HEAD_KEY};
192 use fvm_ipld_encoding::DAG_CBOR;
193 use multihash_codetable::Code::Blake2b256;
194 use nunny::vec as nonempty;
195
196 #[tokio::test]
197 async fn test_export_forest_car() {
198 let db = MemoryDB::default();
199 let record1 = b"non-persistent";
200 let key1 = Cid::new_v1(DAG_CBOR, Blake2b256.digest(record1.as_slice()));
201 db.put_keyed(&key1, record1.as_slice()).unwrap();
202
203 let record2 = b"persistent";
204 let key2 = Cid::new_v1(DAG_CBOR, Blake2b256.digest(record2.as_slice()));
205 db.put_keyed_persistent(&key2, record2.as_slice()).unwrap();
206
207 let mut car_db_bytes = vec![];
208 assert!(
209 db.export_forest_car(&mut car_db_bytes)
210 .await
211 .unwrap_err()
212 .to_string()
213 .contains("chain head is not tracked and cannot be exported")
214 );
215
216 db.write_obj(HEAD_KEY, &TipsetKey::from(nonempty![key1]))
217 .unwrap();
218
219 car_db_bytes.clear();
220 db.export_forest_car(&mut car_db_bytes).await.unwrap();
221
222 let car = ForestCar::new(car_db_bytes).unwrap();
223 assert_eq!(car.head_tipset_key(), &nonempty![key1]);
224 assert!(car.has(&key1).unwrap());
225 assert!(car.has(&key2).unwrap());
226 }
227}