[storage] get_rightmost_leaf take a version

storage/aptosdb/src/state_merkle_db.rs

@@ -199,14 +199,20 @@ impl StateMerkleDb {
 
     /// Finds the rightmost leaf by scanning the entire DB.
     #[cfg(test)]
-    pub fn get_rightmost_leaf_naive(&self) -> Result<Option<(NodeKey, LeafNode)>> {
+    pub fn get_rightmost_leaf_naive(
+        &self,
+        version: Version,
+    ) -> Result<Option<(NodeKey, LeafNode)>> {
         let mut ret = None;
 
         let mut iter = self.iter::<JellyfishMerkleNodeSchema>(Default::default())?;
-        iter.seek_to_first();
+        iter.seek(&(version, 0)).unwrap();
 
         while let Some((node_key, node)) = iter.next().transpose()? {
             if let Node::Leaf(leaf_node) = node {
+                if node_key.version() != version {
+                    break;
+                }
                 match ret {
                     None => ret = Some((node_key, leaf_node)),
                     Some(ref other) => {
@@ -257,17 +263,17 @@ impl TreeReader<StateKey> for StateMerkleDb {
         Ok(node_opt)
     }
 
-    fn get_rightmost_leaf(&self) -> Result<Option<(NodeKey, LeafNode)>> {
+    fn get_rightmost_leaf(&self, version: Version) -> Result<Option<(NodeKey, LeafNode)>> {
         // Since everything has the same version during restore, we seek to the first node and get
         // its version.
         let mut iter = self.iter::<JellyfishMerkleNodeSchema>(Default::default())?;
-        iter.seek_to_first();
+        iter.seek(&(version, 0))?;
         let version = match iter.next().transpose()? {
             Some((node_key, node)) => {
-                if node.node_type() == NodeType::Null {
+                if node.node_type() == NodeType::Null || node_key.version() != version {
                     return Ok(None);
                 }
-                node_key.version()
+                version
             }
             None => return Ok(None),
         };

storage/aptosdb/src/state_restore/restore_test.rs

@@ -108,8 +108,8 @@ where
         self.tree_store.get_node_option(node_key)
     }
 
-    fn get_rightmost_leaf(&self) -> Result<Option<(NodeKey, LeafNode<K>)>> {
-        self.tree_store.get_rightmost_leaf()
+    fn get_rightmost_leaf(&self, version: Version) -> Result<Option<(NodeKey, LeafNode<K>)>> {
+        self.tree_store.get_rightmost_leaf(version)
     }
 }
 

storage/aptosdb/src/state_store/mod.rs

@@ -715,11 +715,12 @@ impl StateStore {
         version: Version,
         expected_root_hash: HashValue,
     ) -> Result<Box<dyn StateSnapshotReceiver<StateKey, StateValue>>> {
-        Ok(Box::new(StateSnapshotRestore::new_overwrite(
+        Ok(Box::new(StateSnapshotRestore::new(
             &self.state_merkle_db,
             self,
             version,
             expected_root_hash,
+            true, /* async_commit */
         )?))
     }
 

storage/aptosdb/src/state_store/state_store_test.rs

@@ -430,8 +430,8 @@ proptest! {
 
         restore.add_chunk(batch1, proof_of_batch1).unwrap();
 
-        let expected = store2.state_merkle_db.get_rightmost_leaf_naive().unwrap();
-        let actual = store2.state_merkle_db.get_rightmost_leaf().unwrap();
+        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);
     }
 

storage/jellyfish-merkle/src/lib.rs

@@ -136,9 +136,9 @@ pub trait TreeReader<K> {
     /// Gets node given a node key. Returns `None` if the node does not exist.
     fn get_node_option(&self, node_key: &NodeKey) -> Result<Option<Node<K>>>;
 
-    /// Gets the rightmost leaf. Note that this assumes we are in the process of restoring the tree
-    /// and all nodes are at the same version.
-    fn get_rightmost_leaf(&self) -> Result<Option<(NodeKey, LeafNode<K>)>>;
+    /// Gets the rightmost leaf at a version. Note that this assumes we are in the process of
+    /// restoring the tree and all nodes are at the same version.
+    fn get_rightmost_leaf(&self, version: Version) -> Result<Option<(NodeKey, LeafNode<K>)>>;
 }
 
 pub trait TreeWriter<K>: Send + Sync {

storage/jellyfish-merkle/src/mock_tree_store.rs

@@ -32,14 +32,16 @@ where
         Ok(self.data.read().0.get(node_key).cloned())
     }
 
-    fn get_rightmost_leaf(&self) -> Result<Option<(NodeKey, LeafNode<K>)>> {
+    fn get_rightmost_leaf(&self, version: Version) -> Result<Option<(NodeKey, LeafNode<K>)>> {
         let locked = self.data.read();
         let mut node_key_and_node: Option<(NodeKey, LeafNode<K>)> = None;
 
         for (key, value) in locked.0.iter() {
             if let Node::Leaf(leaf_node) = value {
-                if node_key_and_node.is_none()
-                    || leaf_node.account_key() > node_key_and_node.as_ref().unwrap().1.account_key()
+                if key.version() == version
+                    && (node_key_and_node.is_none()
+                        || leaf_node.account_key()
+                            > node_key_and_node.as_ref().unwrap().1.account_key())
                 {
                     node_key_and_node.replace((key.clone(), leaf_node.clone()));
                 }

storage/jellyfish-merkle/src/node_type/node_type_test.rs

@@ -30,7 +30,10 @@ impl TreeReader<StateKey> for DummyReader {
         unimplemented!()
     }
 
-    fn get_rightmost_leaf(&self) -> anyhow::Result<Option<(NodeKey, LeafNode<StateKey>)>> {
+    fn get_rightmost_leaf(
+        &self,
+        _version: Version,
+    ) -> anyhow::Result<Option<(NodeKey, LeafNode<StateKey>)>> {
         unimplemented!()
     }
 }

storage/jellyfish-merkle/src/restore/mod.rs

@@ -186,7 +186,7 @@ where
                 expected_root_hash,
             );
             (true, vec![], None)
-        } else if let Some((node_key, leaf_node)) = tree_reader.get_rightmost_leaf()? {
+        } else if let Some((node_key, leaf_node)) = tree_reader.get_rightmost_leaf(version)? {
             // If the system crashed in the middle of the previous restoration attempt, we need
             // to recover the partial nodes to the state right before the crash.
             (