fix_get_right_most_child

api/doc/CHANGELOG.md

@@ -7,6 +7,7 @@ All notable changes to the Aptos Node API will be captured in this file. This ch
 ## Unreleased
 - OpenAPI layout changed slightly in some enum cases, see [#13929](https://github.com/aptos-labs/aptos-core/pull/13929) for more information.
 
+
 ## 1.2.0 (2022-09-29)
 - **[Breaking Changes]** Following the deprecation notice from the previous release, the following breaking changes have landed in this release. Please see the notes from last release for information on the new endpoints you must migrate to:
     - The `/events/{event_key}` endpoint has been removed.

storage/aptosdb/src/db/test_helper.rs

@@ -67,6 +67,7 @@ pub(crate) fn update_store(
     store: &StateStore,
     input: impl Iterator<Item = (StateKey, Option<StateValue>)>,
     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,
             )

storage/aptosdb/src/pruner/state_merkle_pruner/test.rs

@@ -375,7 +375,7 @@ fn verify_state_value_pruner(inputs: Vec<Vec<(StateKey, Option<StateValue>)>>) {
         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()))

storage/aptosdb/src/state_merkle_db.rs

@@ -27,7 +27,7 @@ use aptos_rocksdb_options::gen_rocksdb_options;
 use aptos_schemadb::{SchemaBatch, DB};
 #[cfg(test)]
 use aptos_scratchpad::get_state_shard_id;
-use aptos_storage_interface::{db_ensure as ensure, AptosDbError, Result};
+use aptos_storage_interface::{db_ensure as ensure, db_other_bail as bail, AptosDbError, Result};
 use aptos_types::{
     nibble::{nibble_path::NibblePath, ROOT_NIBBLE_HEIGHT},
     proof::{SparseMerkleProofExt, SparseMerkleRangeProof},
@@ -760,7 +760,11 @@ impl StateMerkleDb {
 
             if let Some((node_key, node)) = iter.next().transpose()? {
                 if node_key.version() != version {
-                    continue;
+                    bail!(
+                        "node key version {} != version {}",
+                        node_key.version(),
+                        version
+                    );
                 }
                 if let Node::Leaf(leaf_node) = node {
                     match ret {

storage/aptosdb/src/state_store/state_store_test.rs

@@ -465,6 +465,61 @@ proptest! {
         );
     }
 
+    #[test]
+    fn test_get_rightmost_leaf_with_sharding(
+        (input, batch1_size) in hash_map(any::<StateKey>(), any::<StateValue>(), 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 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 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
+        assert!(actual.is_some());
+        prop_assert_eq!(actual, expected);
+    }
+
     #[test]
     fn test_get_rightmost_leaf(
         (input, batch1_size) in hash_map(any::<StateKey>(), any::<StateValue>(), 2..1000)
@@ -512,6 +567,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 +582,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 +630,14 @@ proptest! {
 
 // Initializes the state store by inserting one key at each version.
 fn init_store(store: &StateStore, input: impl Iterator<Item = (StateKey, StateValue)>) {
-    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<Item = (StateKey, StateValue)>) {
+    update_store(store, input.into_iter().map(|(k, v)| (k, Some(v))), 0, true);
 }