From 98cf706d9c064dfe6f7147345ed9ffc8887619a8 Mon Sep 17 00:00:00 2001 From: Bo Wu Date: Mon, 7 Oct 2024 11:47:14 -0700 Subject: [PATCH] fix_get_right_most_child (cherry picked from commit 5150f8317c26973cd2f22fe304cb97b46262b039) --- storage/aptosdb/src/db/test_helper.rs | 3 +- .../src/pruner/state_merkle_pruner/test.rs | 2 +- storage/aptosdb/src/state_merkle_db.rs | 32 +-------- .../src/state_store/state_store_test.rs | 71 +++++++++++++++++-- 4 files changed, 70 insertions(+), 38 deletions(-) diff --git a/storage/aptosdb/src/db/test_helper.rs b/storage/aptosdb/src/db/test_helper.rs index 570bdf8bd4433..68265f72c0e34 100644 --- a/storage/aptosdb/src/db/test_helper.rs +++ b/storage/aptosdb/src/db/test_helper.rs @@ -67,6 +67,7 @@ pub(crate) fn update_store( store: &StateStore, input: impl Iterator)>, first_version: Version, + enable_sharding: bool, ) -> HashValue { use aptos_storage_interface::{jmt_update_refs, jmt_updates}; let mut root_hash = *aptos_crypto::hash::SPARSE_MERKLE_PLACEHOLDER_HASH; @@ -94,7 +95,7 @@ pub(crate) fn update_store( None, &ledger_batch, &sharded_state_kv_batches, - /*put_state_value_indices=*/ false, + /*put_state_value_indices=*/ enable_sharding, /*skip_usage=*/ false, /*last_checkpoint_index=*/ None, ) diff --git a/storage/aptosdb/src/pruner/state_merkle_pruner/test.rs b/storage/aptosdb/src/pruner/state_merkle_pruner/test.rs index a660a6082975e..edf59efde854f 100644 --- a/storage/aptosdb/src/pruner/state_merkle_pruner/test.rs +++ b/storage/aptosdb/src/pruner/state_merkle_pruner/test.rs @@ -375,7 +375,7 @@ fn verify_state_value_pruner(inputs: Vec)>>) { user_pruning_window_offset: 0, }); for batch in inputs { - update_store(store, batch.clone().into_iter(), version); + update_store(store, batch.clone().into_iter(), version, false); for (k, v) in batch.iter() { if let Some((old_version, old_v_opt)) = current_state_values.insert(k.clone(), (version, v.clone())) diff --git a/storage/aptosdb/src/state_merkle_db.rs b/storage/aptosdb/src/state_merkle_db.rs index 4c42a61042712..b1f5e757d9fb5 100644 --- a/storage/aptosdb/src/state_merkle_db.rs +++ b/storage/aptosdb/src/state_merkle_db.rs @@ -19,8 +19,7 @@ use aptos_config::config::{RocksdbConfig, RocksdbConfigs, StorageDirPaths}; use aptos_crypto::{hash::CryptoHash, HashValue}; use aptos_experimental_runtimes::thread_manager::{optimal_min_len, THREAD_MANAGER}; use aptos_jellyfish_merkle::{ - node_type::{NodeKey, NodeType}, - JellyfishMerkleTree, TreeReader, TreeUpdateBatch, TreeWriter, + node_type::NodeKey, JellyfishMerkleTree, TreeReader, TreeUpdateBatch, TreeWriter, }; use aptos_logger::prelude::*; use aptos_rocksdb_options::gen_rocksdb_options; @@ -677,20 +676,6 @@ impl StateMerkleDb { ) -> Result> { let mut ret = None; - if self.enable_sharding { - let mut iter = self.metadata_db().iter::()?; - iter.seek(&(version, 0)).unwrap(); - // early exit if no node is found for the target version - match iter.next().transpose()? { - Some((node_key, node)) => { - if node.node_type() == NodeType::Null || node_key.version() != version { - return Ok(None); - } - }, - None => return Ok(None), - }; - } - // traverse all shards in a naive way let shards = 0..self.hack_num_real_shards(); let start_num_of_nibbles = if self.enable_sharding { 1 } else { 0 }; @@ -822,21 +807,6 @@ impl TreeReader for StateMerkleDb { } fn get_rightmost_leaf(&self, version: Version) -> Result> { - // Since everything has the same version during restore, we seek to the first node and get - // its version. - - let mut iter = self.metadata_db().iter::()?; - // get the root node corresponding to the version - iter.seek(&(version, 0))?; - match iter.next().transpose()? { - Some((node_key, node)) => { - if node.node_type() == NodeType::Null || node_key.version() != version { - return Ok(None); - } - }, - None => return Ok(None), - }; - let ret = None; let shards = 0..self.hack_num_real_shards(); diff --git a/storage/aptosdb/src/state_store/state_store_test.rs b/storage/aptosdb/src/state_store/state_store_test.rs index b76d396b1b630..736e788d443f4 100644 --- a/storage/aptosdb/src/state_store/state_store_test.rs +++ b/storage/aptosdb/src/state_store/state_store_test.rs @@ -465,6 +465,59 @@ proptest! { ); } + #[test] + fn test_get_rightmost_leaf_with_sharding( + (input, batch1_size) in hash_map(any::(), any::(), 2..1000) + .prop_flat_map(|input| { + let len = input.len(); + (Just(input), 1..len) + }) + ) { + let tmp_dir1 = TempPath::new(); + let db1 = AptosDB::new_for_test_with_sharding(&tmp_dir1, 1000); + let store1 = &db1.state_store; + init_sharded_store(store1, input.clone().into_iter()); + + let version = (input.len() - 1) as Version; + let expected_root_hash = store1.get_root_hash(version).unwrap(); + + let tmp_dir2 = TempPath::new(); + let db2 = AptosDB::new_for_test_with_sharding(&tmp_dir2, 1000); + + + let store2 = &db2.state_store; + let mut restore = + StateSnapshotRestore::new(&store2.state_merkle_db, store2, version, expected_root_hash, true, /* async_commit */ StateSnapshotRestoreMode::Default).unwrap(); + let max_hash = HashValue::new([0xff; HashValue::LENGTH]); + let dummy_state_key = StateKey::raw(&[]); + let (top_levels_batch, sharded_batches, _) = store2.state_merkle_db.merklize_value_set(vec![(max_hash, Some(&(HashValue::random(), dummy_state_key)))], 0, None, None).unwrap(); + store2.state_merkle_db.commit(version, top_levels_batch, sharded_batches).unwrap(); + assert!(store2.state_merkle_db.get_rightmost_leaf(version).unwrap().is_none()); + let mut ordered_input: Vec<_> = input + .into_iter() + .collect(); + ordered_input.sort_unstable_by_key(|(key, _value)| key.hash()); + + let batch1: Vec<_> = ordered_input + .into_iter() + .take(batch1_size) + .collect(); + let rightmost_of_batch1 = batch1.last().map(|(key, _value)| key.hash()).unwrap(); + let proof_of_batch1 = store1 + .get_value_range_proof(rightmost_of_batch1, version) + .unwrap(); + + restore.add_chunk(batch1, proof_of_batch1).unwrap(); + restore.wait_for_async_commit().unwrap(); + + let expected = store2.state_merkle_db.get_rightmost_leaf_naive(version).unwrap(); + // When re-initializing the store, the rightmost leaf should exist indicating the progress + let actual = store2.state_merkle_db.get_rightmost_leaf(version).unwrap(); + // ensure the rightmost leaf is not None + prop_assert!(actual.is_some()); + prop_assert_eq!(actual, expected); + } + #[test] fn test_get_rightmost_leaf( (input, batch1_size) in hash_map(any::(), any::(), 2..1000) @@ -484,15 +537,13 @@ proptest! { let tmp_dir2 = TempPath::new(); let db2 = AptosDB::new_for_test(&tmp_dir2); let store2 = &db2.state_store; - let max_hash = HashValue::new([0xff; HashValue::LENGTH]); let mut restore = StateSnapshotRestore::new(&store2.state_merkle_db, store2, version, expected_root_hash, true, /* async_commit */ StateSnapshotRestoreMode::Default).unwrap(); - + let max_hash = HashValue::new([0xff; HashValue::LENGTH]); let dummy_state_key = StateKey::raw(&[]); let (top_levels_batch, sharded_batches, _) = store2.state_merkle_db.merklize_value_set(vec![(max_hash, Some(&(HashValue::random(), dummy_state_key)))], 0, None, None).unwrap(); store2.state_merkle_db.commit(version, top_levels_batch, sharded_batches).unwrap(); assert!(store2.state_merkle_db.get_rightmost_leaf(version).unwrap().is_none()); - let mut ordered_input: Vec<_> = input .into_iter() .collect(); @@ -512,6 +563,7 @@ proptest! { let expected = store2.state_merkle_db.get_rightmost_leaf_naive(version).unwrap(); let actual = store2.state_merkle_db.get_rightmost_leaf(version).unwrap(); + prop_assert_eq!(actual, expected); } @@ -526,7 +578,7 @@ proptest! { let mut version = 0; for batch in input { let next_version = version + batch.len() as Version; - let root_hash = update_store(store, batch.into_iter(), version); + let root_hash = update_store(store, batch.into_iter(), version, false); let last_version = next_version - 1; let snapshot = db @@ -574,5 +626,14 @@ proptest! { // Initializes the state store by inserting one key at each version. fn init_store(store: &StateStore, input: impl Iterator) { - update_store(store, input.into_iter().map(|(k, v)| (k, Some(v))), 0); + update_store( + store, + input.into_iter().map(|(k, v)| (k, Some(v))), + 0, + false, + ); +} + +fn init_sharded_store(store: &StateStore, input: impl Iterator) { + update_store(store, input.into_iter().map(|(k, v)| (k, Some(v))), 0, true); }