smeb_y/0g-storage-node
1
0
Fork 0
mirror of https://github.com/0glabs/0g-storage-node.git synced 2024-11-15 04:25:19 +00:00
Code Issues Actions 4 Packages Projects Releases Wiki Activity

Add trait.

Browse Source
This commit is contained in:
Peilun Li 2024-10-08 15:40:47 +08:00
parent b76fa7be9b
commit 5100c22933
6 changed files with 152 additions and 23 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

43
common/append_merkle/src/lib.rs
View File

@ -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<E: HashElement, A: Algorithm<E>> {
/// Keep all the nodes in the latest version. `layers[0]` is the layer of leaves.
layers: Vec<Vec<E>>,
node_manager: NodeManager<E>,
/// 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<u64, DeltaNodes<E>>,
@ -33,9 +37,15 @@ pub struct AppendMerkleTree<E: HashElement, A: Algorithm<E>> {
}
impl<E: HashElement, A: Algorithm<E>> AppendMerkleTree<E, A> {
pub fn new(leaves: Vec<E>, leaf_height: usize, start_tx_seq: Option<u64>) -> Self {
pub fn new(
node_db: Arc<dyn NodeDatabase<E>>,
leaves: Vec<E>,
leaf_height: usize,
start_tx_seq: Option<u64>,
) -> 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<E: HashElement, A: Algorithm<E>> AppendMerkleTree<E, A> {
}
pub fn new_with_subtrees(
node_db: Arc<dyn NodeDatabase<E>>,
initial_data: MerkleTreeInitialData<E>,
leaf_height: usize,
start_tx_seq: Option<u64>,
) -> Result<Self> {
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<E: HashElement, A: Algorithm<E>> AppendMerkleTree<E, A> {
// 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<E: HashElement, A: Algorithm<E>> AppendMerkleTree<E, A> {
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::<H256, Sha3Algorithm>::new(vec![H256::zero()], 0, None);
let mut merkle = AppendMerkleTree::<H256, Sha3Algorithm>::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::<H256, Sha3Algorithm>::new(vec![H256::zero()], 0, None);
let mut merkle = AppendMerkleTree::<H256, Sha3Algorithm>::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::<H256, Sha3Algorithm>::new(vec![H256::zero()], 0, None);
let mut merkle = AppendMerkleTree::<H256, Sha3Algorithm>::new(
Arc::new(EmptyNodeDatabase {}),
vec![H256::zero()],
0,
None,
);
let mut start_pos = 0;
for (tx_seq, &entry_len) in n.iter().enumerate() {

53
common/append_merkle/src/node_manager.rs Normal file
View File

@ -0,0 +1,53 @@
use crate::HashElement;
use anyhow::Result;
use std::collections::HashMap;
use std::sync::Arc;
use tracing::error;
pub struct NodeManager<E: HashElement> {
cache: HashMap<(usize, usize), E>,
db: Arc<dyn NodeDatabase<E>>,
}
impl<E: HashElement> NodeManager<E> {
pub fn new(db: Arc<dyn NodeDatabase<E>>) -> Self {
Self {
cache: HashMap::new(),
db,
}
}
pub fn get_node(&self, layer: usize, pos: usize) -> Option<E> {
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<E: HashElement>: Send + Sync {
fn get_node(&self, layer: usize, pos: usize) -> Result<Option<E>>;
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<E: HashElement> NodeDatabase<E> for EmptyNodeDatabase {
fn get_node(&self, _layer: usize, _pos: usize) -> Result<Option<E>> {
Ok(None)
}
fn save_node(&self, _layer: usize, _pos: usize, _node: &E) -> Result<()> {
Ok(())
}
}

35
node/storage/src/log_store/flow_store.rs
View File

@ -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<FlowDBStore>,
seal_manager: SealTaskManager,
config: FlowConfig,
}
impl FlowStore {
pub fn new(db: Arc<dyn ZgsKeyValueDB>, config: FlowConfig) -> Self {
pub fn new(db: Arc<FlowDBStore>, 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::<u64>()..])?;
Ok((layer_index, position))
}
impl NodeDatabase<DataRoot> for FlowDBStore {
fn get_node(&self, layer: usize, pos: usize) -> Result<Option<DataRoot>> {
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)?)
}
}

8
node/storage/src/log_store/load_chunk/mod.rs
View File

@ -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() {

23
node/storage/src/log_store/log_manager.rs
View File

@ -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<Self> {
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;

13
node/storage/src/log_store/tests.rs
View File

@ -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::<H256, Sha3Algorithm>::new(vec![H256::zero()], 0, None);
let mut merkle = AppendMerkleTree::<H256, Sha3Algorithm>::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::<H256, Sha3Algorithm>::new(
Arc::new(EmptyNodeDatabase {}),
data_to_merkle_leaves(&data).unwrap(),
0,
None,