[CP][Fix] Return the correct version for state snapshot when the root…

… is not committed. (#8671)

storage/aptosdb/src/state_merkle_db.rs

@@ -255,9 +255,18 @@ impl StateMerkleDb {
                 .rev_iter::<JellyfishMerkleNodeSchema>(Default::default())?;
             iter.seek_for_prev(&NodeKey::new_empty_path(max_possible_version))?;
             if let Some((key, _node)) = iter.next().transpose()? {
-                // TODO: If we break up a single update batch to multiple commits, we would need to
-                // deal with a partial version, which hasn't got the root committed.
-                return Ok(Some(key.version()));
+                let version = key.version();
+                if self
+                    .metadata_db()
+                    .get::<JellyfishMerkleNodeSchema>(&NodeKey::new_empty_path(version))?
+                    .is_some()
+                {
+                    return Ok(Some(version));
+                }
+                // Since we split state merkle commit into multiple batches, it's possible that
+                // the root is not committed yet. In this case we need to look at the previous
+                // root.
+                return self.get_state_snapshot_version_before(version);
             }
         }
         // No version before genesis.

storage/aptosdb/src/state_store/state_store_test.rs

@@ -4,18 +4,23 @@
 
 use super::*;
 use crate::{
+    jellyfish_merkle_node::JellyfishMerkleNodeSchema,
     new_sharded_kv_schema_batch,
     state_restore::StateSnapshotRestore,
     test_helper::{arb_state_kv_sets, update_store},
     AptosDB,
 };
-use aptos_jellyfish_merkle::TreeReader;
+use aptos_jellyfish_merkle::{
+    node_type::{Node, NodeKey},
+    TreeReader,
+};
 use aptos_storage_interface::{
     jmt_update_refs, jmt_updates, DbReader, DbWriter, StateSnapshotReceiver,
 };
 use aptos_temppath::TempPath;
 use aptos_types::{
-    access_path::AccessPath, account_address::AccountAddress, state_store::state_key::StateKeyTag,
+    access_path::AccessPath, account_address::AccountAddress, nibble::nibble_path::NibblePath,
+    state_store::state_key::StateKeyTag,
 };
 use proptest::{collection::hash_map, prelude::*};
 
@@ -276,7 +281,46 @@ pub fn test_get_state_snapshot_before() {
     assert_eq!(store.get_state_snapshot_before(3).unwrap(), Some((2, hash)));
     assert_eq!(store.get_state_snapshot_before(2).unwrap(), Some((0, hash)));
     assert_eq!(store.get_state_snapshot_before(1).unwrap(), Some((0, hash)));
-    assert_eq!(store.get_state_snapshot_before(0).unwrap(), None,);
+    assert_eq!(store.get_state_snapshot_before(0).unwrap(), None);
+
+    // Only consider a version as available when the root node is there.
+    // Here we are adding another non-root node, and removing the root node, to verify if there is
+    // a node at version X but the root node at version X doesn't exist, we shouldn't return
+    // version X.
+    let batch = SchemaBatch::new();
+    batch
+        .put::<JellyfishMerkleNodeSchema>(
+            &NodeKey::new(2, NibblePath::new_odd(vec![0])),
+            &Node::Null,
+        )
+        .unwrap();
+    db.state_merkle_db
+        .metadata_db()
+        .write_schemas(batch)
+        .unwrap();
+
+    assert_eq!(
+        db.state_merkle_db
+            .get_state_snapshot_version_before(4)
+            .unwrap(),
+        Some(2)
+    );
+
+    let batch = SchemaBatch::new();
+    batch
+        .delete::<JellyfishMerkleNodeSchema>(&NodeKey::new_empty_path(2))
+        .unwrap();
+    db.state_merkle_db
+        .metadata_db()
+        .write_schemas(batch)
+        .unwrap();
+
+    assert_eq!(
+        db.state_merkle_db
+            .get_state_snapshot_version_before(4)
+            .unwrap(),
+        Some(0)
+    );
 }
 
 proptest! {