diff --git a/common/append_merkle/src/lib.rs b/common/append_merkle/src/lib.rs index 55fe806..32677d3 100644 --- a/common/append_merkle/src/lib.rs +++ b/common/append_merkle/src/lib.rs @@ -1,4 +1,5 @@ mod merkle_tree; +mod node_manager; mod proof; mod sha3; @@ -7,17 +8,20 @@ use std::cmp::Ordering; use std::collections::{BTreeMap, HashMap}; use std::fmt::Debug; use std::marker::PhantomData; +use std::sync::Arc; use tracing::{trace, warn}; pub use crate::merkle_tree::{ Algorithm, HashElement, MerkleTreeInitialData, MerkleTreeRead, ZERO_HASHES, }; +pub use crate::node_manager::{EmptyNodeDatabase, NodeDatabase, NodeManager}; pub use proof::{Proof, RangeProof}; pub use sha3::Sha3Algorithm; pub struct AppendMerkleTree> { /// Keep all the nodes in the latest version. `layers[0]` is the layer of leaves. layers: Vec>, + node_manager: NodeManager, /// Keep the delta nodes that can be used to construct a history tree. /// The key is the root node of that version. delta_nodes_map: BTreeMap>, @@ -33,9 +37,15 @@ pub struct AppendMerkleTree> { } impl> AppendMerkleTree { - pub fn new(leaves: Vec, leaf_height: usize, start_tx_seq: Option) -> Self { + pub fn new( + node_db: Arc>, + leaves: Vec, + leaf_height: usize, + start_tx_seq: Option, + ) -> Self { let mut merkle = Self { layers: vec![leaves], + node_manager: NodeManager::new(node_db), delta_nodes_map: BTreeMap::new(), root_to_tx_seq_map: HashMap::new(), min_depth: None, @@ -62,12 +72,14 @@ impl> AppendMerkleTree { } pub fn new_with_subtrees( + node_db: Arc>, initial_data: MerkleTreeInitialData, leaf_height: usize, start_tx_seq: Option, ) -> Result { let mut merkle = Self { layers: vec![vec![]], + node_manager: NodeManager::new(node_db), delta_nodes_map: BTreeMap::new(), root_to_tx_seq_map: HashMap::new(), min_depth: None, @@ -103,6 +115,8 @@ impl> AppendMerkleTree { // Create an empty merkle tree with `depth`. let mut merkle = Self { layers: vec![vec![]; depth], + // dummy node manager for the last chunk. + node_manager: NodeManager::new(Arc::new(EmptyNodeDatabase {})), delta_nodes_map: BTreeMap::new(), root_to_tx_seq_map: HashMap::new(), min_depth: Some(depth), @@ -123,6 +137,8 @@ impl> AppendMerkleTree { layers[0] = leaves; let mut merkle = Self { layers, + // dummy node manager for the last chunk. + node_manager: NodeManager::new(Arc::new(EmptyNodeDatabase {})), delta_nodes_map: BTreeMap::new(), root_to_tx_seq_map: HashMap::new(), min_depth: Some(depth), @@ -699,9 +715,11 @@ macro_rules! ensure_eq { #[cfg(test)] mod tests { use crate::merkle_tree::MerkleTreeRead; + use crate::node_manager::EmptyNodeDatabase; use crate::sha3::Sha3Algorithm; use crate::AppendMerkleTree; use ethereum_types::H256; + use std::sync::Arc; #[test] fn test_proof() { @@ -711,8 +729,12 @@ mod tests { for _ in 0..entry_len { data.push(H256::random()); } - let mut merkle = - AppendMerkleTree::::new(vec![H256::zero()], 0, None); + let mut merkle = AppendMerkleTree::::new( + Arc::new(EmptyNodeDatabase {}), + vec![H256::zero()], + 0, + None, + ); merkle.append_list(data.clone()); merkle.commit(Some(0)); verify(&data, &mut merkle); @@ -739,8 +761,12 @@ mod tests { for _ in 0..entry_len { data.push(H256::random()); } - let mut merkle = - AppendMerkleTree::::new(vec![H256::zero()], 0, None); + let mut merkle = AppendMerkleTree::::new( + Arc::new(EmptyNodeDatabase {}), + vec![H256::zero()], + 0, + None, + ); merkle.append_list(data.clone()); merkle.commit(Some(0)); @@ -764,7 +790,12 @@ mod tests { #[test] fn test_proof_at_version() { let n = [2, 255, 256, 257]; - let mut merkle = AppendMerkleTree::::new(vec![H256::zero()], 0, None); + let mut merkle = AppendMerkleTree::::new( + Arc::new(EmptyNodeDatabase {}), + vec![H256::zero()], + 0, + None, + ); let mut start_pos = 0; for (tx_seq, &entry_len) in n.iter().enumerate() { diff --git a/common/append_merkle/src/node_manager.rs b/common/append_merkle/src/node_manager.rs new file mode 100644 index 0000000..dbef6b7 --- /dev/null +++ b/common/append_merkle/src/node_manager.rs @@ -0,0 +1,53 @@ +use crate::HashElement; +use anyhow::Result; +use std::collections::HashMap; +use std::sync::Arc; +use tracing::error; + +pub struct NodeManager { + cache: HashMap<(usize, usize), E>, + db: Arc>, +} + +impl NodeManager { + pub fn new(db: Arc>) -> Self { + Self { + cache: HashMap::new(), + db, + } + } + + pub fn get_node(&self, layer: usize, pos: usize) -> Option { + match self.cache.get(&(layer, pos)) { + Some(node) => Some(node.clone()), + None => self.db.get_node(layer, pos).unwrap_or_else(|e| { + error!("Failed to get node: {}", e); + None + }), + } + } + + pub fn add_node(&mut self, layer: usize, pos: usize, node: E) { + if let Err(e) = self.db.save_node(layer, pos, &node) { + error!("Failed to save node: {}", e); + } + self.cache.insert((layer, pos), node); + } +} + +pub trait NodeDatabase: Send + Sync { + fn get_node(&self, layer: usize, pos: usize) -> Result>; + fn save_node(&self, layer: usize, pos: usize, node: &E) -> Result<()>; +} + +/// A dummy database structure for in-memory merkle tree that will not read/write db. +pub struct EmptyNodeDatabase {} +impl NodeDatabase for EmptyNodeDatabase { + fn get_node(&self, _layer: usize, _pos: usize) -> Result> { + Ok(None) + } + + fn save_node(&self, _layer: usize, _pos: usize, _node: &E) -> Result<()> { + Ok(()) + } +} diff --git a/node/storage/src/log_store/flow_store.rs b/node/storage/src/log_store/flow_store.rs index 37ee1e2..b99aa58 100644 --- a/node/storage/src/log_store/flow_store.rs +++ b/node/storage/src/log_store/flow_store.rs @@ -10,7 +10,7 @@ use crate::log_store::log_manager::{ use crate::log_store::{FlowRead, FlowSeal, FlowWrite}; use crate::{try_option, ZgsKeyValueDB}; use anyhow::{anyhow, bail, Result}; -use append_merkle::{MerkleTreeInitialData, MerkleTreeRead}; +use append_merkle::{EmptyNodeDatabase, MerkleTreeInitialData, MerkleTreeRead, NodeDatabase}; use itertools::Itertools; use parking_lot::RwLock; use shared_types::{ChunkArray, DataRoot, FlowProof, Merkle}; @@ -25,15 +25,15 @@ use tracing::{debug, error, trace}; use zgs_spec::{BYTES_PER_SECTOR, SEALS_PER_LOAD, SECTORS_PER_LOAD, SECTORS_PER_SEAL}; pub struct FlowStore { - db: FlowDBStore, + db: Arc, seal_manager: SealTaskManager, config: FlowConfig, } impl FlowStore { - pub fn new(db: Arc, config: FlowConfig) -> Self { + pub fn new(db: Arc, config: FlowConfig) -> Self { Self { - db: FlowDBStore::new(db), + db, seal_manager: Default::default(), config, } @@ -436,7 +436,13 @@ impl FlowDBStore { let mut expected_index = 0; let empty_data = vec![0; PORA_CHUNK_SIZE * ENTRY_SIZE]; - let empty_root = *Merkle::new(data_to_merkle_leaves(&empty_data)?, 0, None).root(); + let empty_root = *Merkle::new( + Arc::new(EmptyNodeDatabase {}), + data_to_merkle_leaves(&empty_data)?, + 0, + None, + ) + .root(); for r in self.kvdb.iter(COL_ENTRY_BATCH_ROOT) { let (index_bytes, root_bytes) = r?; @@ -666,3 +672,22 @@ fn decode_mpt_node_key(data: &[u8]) -> Result<(usize, usize)> { let position = try_decode_usize(&data[mem::size_of::()..])?; Ok((layer_index, position)) } + +impl NodeDatabase for FlowDBStore { + fn get_node(&self, layer: usize, pos: usize) -> Result> { + Ok(self + .kvdb + .get(COL_FLOW_MPT_NODES, &encode_mpt_node_key(layer, pos))? + .map(|v| DataRoot::from_slice(&v))) + } + + fn save_node(&self, layer: usize, pos: usize, node: &DataRoot) -> Result<()> { + let mut tx = self.kvdb.transaction(); + tx.put( + COL_FLOW_MPT_NODES, + &encode_mpt_node_key(layer, pos), + node.as_bytes(), + ); + Ok(self.kvdb.write(tx)?) + } +} diff --git a/node/storage/src/log_store/load_chunk/mod.rs b/node/storage/src/log_store/load_chunk/mod.rs index 3daff69..c27f13c 100644 --- a/node/storage/src/log_store/load_chunk/mod.rs +++ b/node/storage/src/log_store/load_chunk/mod.rs @@ -4,15 +4,15 @@ mod seal; mod serde; use ::serde::{Deserialize, Serialize}; -use std::cmp::min; - use anyhow::Result; use ethereum_types::H256; use ssz_derive::{Decode, Encode}; +use std::cmp::min; +use std::sync::Arc; use crate::log_store::log_manager::data_to_merkle_leaves; use crate::try_option; -use append_merkle::{Algorithm, MerkleTreeRead, Sha3Algorithm}; +use append_merkle::{Algorithm, EmptyNodeDatabase, MerkleTreeRead, Sha3Algorithm}; use shared_types::{ChunkArray, DataRoot, Merkle}; use tracing::trace; use zgs_spec::{ @@ -248,7 +248,7 @@ impl EntryBatch { } else { vec![] }; - let mut merkle = Merkle::new(initial_leaves, 0, None); + let mut merkle = Merkle::new(Arc::new(EmptyNodeDatabase {}), initial_leaves, 0, None); for subtree in self.data.get_subtree_list() { trace!(?subtree, "get subtree, leaves={}", merkle.leaves()); if subtree.start_sector != merkle.leaves() { diff --git a/node/storage/src/log_store/log_manager.rs b/node/storage/src/log_store/log_manager.rs index 20b72ec..4fe4bbb 100644 --- a/node/storage/src/log_store/log_manager.rs +++ b/node/storage/src/log_store/log_manager.rs @@ -1,12 +1,12 @@ use crate::config::ShardConfig; -use crate::log_store::flow_store::{batch_iter_sharded, FlowConfig, FlowStore}; +use crate::log_store::flow_store::{batch_iter_sharded, FlowConfig, FlowDBStore, FlowStore}; use crate::log_store::tx_store::TransactionStore; use crate::log_store::{ FlowRead, FlowWrite, LogStoreChunkRead, LogStoreChunkWrite, LogStoreRead, LogStoreWrite, }; use crate::{try_option, ZgsKeyValueDB}; use anyhow::{anyhow, bail, Result}; -use append_merkle::{Algorithm, MerkleTreeRead, Sha3Algorithm}; +use append_merkle::{Algorithm, EmptyNodeDatabase, MerkleTreeRead, Sha3Algorithm}; use ethereum_types::H256; use kvdb_rocksdb::{Database, DatabaseConfig}; use merkle_light::merkle::{log2_pow2, MerkleTree}; @@ -626,7 +626,8 @@ impl LogManager { executor: task_executor::TaskExecutor, ) -> Result { let tx_store = TransactionStore::new(db.clone())?; - let flow_store = Arc::new(FlowStore::new(db.clone(), config.flow)); + let flow_db = Arc::new(FlowDBStore::new(db.clone())); + let flow_store = Arc::new(FlowStore::new(flow_db.clone(), config.flow)); let mut initial_data = flow_store.get_chunk_root_list()?; // If the last tx `put_tx` does not complete, we will revert it in `initial_data.subtree_list` // first and call `put_tx` later. The known leaves in its data will be saved in `extra_leaves` @@ -705,8 +706,12 @@ impl LogManager { } } - let mut pora_chunks_merkle = - Merkle::new_with_subtrees(initial_data, log2_pow2(PORA_CHUNK_SIZE), start_tx_seq)?; + let mut pora_chunks_merkle = Merkle::new_with_subtrees( + flow_db, + initial_data, + log2_pow2(PORA_CHUNK_SIZE), + start_tx_seq, + )?; let last_chunk_merkle = match start_tx_seq { Some(tx_seq) => { tx_store.rebuild_last_chunk_merkle(pora_chunks_merkle.leaves(), tx_seq)? @@ -977,7 +982,13 @@ impl LogManager { let data = pad_data[start_index * ENTRY_SIZE ..(start_index + PORA_CHUNK_SIZE) * ENTRY_SIZE] .to_vec(); - let root = *Merkle::new(data_to_merkle_leaves(&data)?, 0, None).root(); + let root = *Merkle::new( + Arc::new(EmptyNodeDatabase {}), + data_to_merkle_leaves(&data)?, + 0, + None, + ) + .root(); merkle.pora_chunks_merkle.append(root); root_map.insert(merkle.pora_chunks_merkle.leaves() - 1, (root, 1)); start_index += PORA_CHUNK_SIZE; diff --git a/node/storage/src/log_store/tests.rs b/node/storage/src/log_store/tests.rs index d56613d..aef7d55 100644 --- a/node/storage/src/log_store/tests.rs +++ b/node/storage/src/log_store/tests.rs @@ -3,11 +3,14 @@ use crate::log_store::log_manager::{ PORA_CHUNK_SIZE, }; use crate::log_store::{LogStoreChunkRead, LogStoreChunkWrite, LogStoreRead, LogStoreWrite}; -use append_merkle::{Algorithm, AppendMerkleTree, MerkleTreeRead, Sha3Algorithm}; +use append_merkle::{ + Algorithm, AppendMerkleTree, EmptyNodeDatabase, MerkleTreeRead, Sha3Algorithm, +}; use ethereum_types::H256; use rand::random; use shared_types::{compute_padded_chunk_size, ChunkArray, Transaction, CHUNK_SIZE}; use std::cmp; +use std::sync::Arc; use task_executor::test_utils::TestRuntime; #[test] @@ -26,7 +29,12 @@ fn test_put_get() { data[i * CHUNK_SIZE] = random(); } let (padded_chunks, _) = compute_padded_chunk_size(data_size); - let mut merkle = AppendMerkleTree::::new(vec![H256::zero()], 0, None); + let mut merkle = AppendMerkleTree::::new( + Arc::new(EmptyNodeDatabase {}), + vec![H256::zero()], + 0, + None, + ); merkle.append_list(data_to_merkle_leaves(&LogManager::padding_raw(start_offset - 1)).unwrap()); let mut data_padded = data.clone(); data_padded.append(&mut vec![0u8; CHUNK_SIZE]); @@ -124,6 +132,7 @@ fn test_root() { let mt = sub_merkle_tree(&data).unwrap(); println!("{:?} {}", mt.root(), hex::encode(mt.root())); let append_mt = AppendMerkleTree::::new( + Arc::new(EmptyNodeDatabase {}), data_to_merkle_leaves(&data).unwrap(), 0, None,