From 8d4ea519b35004f9d73391ad7177ad71c2b5d815 Mon Sep 17 00:00:00 2001
From: aldenhu <msmouse@gmail.com>
Date: Thu, 23 Jan 2025 11:56:41 -0800
Subject: [PATCH] Fix state merkle truncate

in case of a partial commit on a unsharded state db, two things were
wrong:
1. when we see a partial version, we need to look at the previous version before jumping the the epoch boundary (we forgot to substract 1 from the version)
2. the ledger meta db were not passed in as intended (instead, we passed
   the state merkle meta db)

The test_db_restart smoke test does catch this sometimes but it's not
easy to trigger. I bumpped the number of restarts on that test.
---
 .../src/ledger_db/ledger_metadata_db.rs        |  2 +-
 storage/aptosdb/src/state_store/mod.rs         | 18 ++++++++++--------
 storage/aptosdb/src/utils/truncation_helper.rs | 11 +++++++----
 testsuite/smoke-test/src/storage.rs            |  2 +-
 4 files changed, 19 insertions(+), 14 deletions(-)

diff --git a/storage/aptosdb/src/ledger_db/ledger_metadata_db.rs b/storage/aptosdb/src/ledger_db/ledger_metadata_db.rs
index 39494dae4dee9..9cbe3651b3b10 100644
--- a/storage/aptosdb/src/ledger_db/ledger_metadata_db.rs
+++ b/storage/aptosdb/src/ledger_db/ledger_metadata_db.rs
@@ -61,7 +61,7 @@ impl LedgerMetadataDb {
         )
     }
 
-    pub(super) fn db(&self) -> &DB {
+    pub(crate) fn db(&self) -> &DB {
         &self.db
     }
 
diff --git a/storage/aptosdb/src/state_store/mod.rs b/storage/aptosdb/src/state_store/mod.rs
index d5c68b382e00a..17151854c6f7d 100644
--- a/storage/aptosdb/src/state_store/mod.rs
+++ b/storage/aptosdb/src/state_store/mod.rs
@@ -397,7 +397,7 @@ impl StateStore {
             if crash_if_difference_is_too_large {
                 assert_le!(difference, MAX_COMMIT_PROGRESS_DIFFERENCE);
             }
-            truncate_ledger_db(ledger_db, overall_commit_progress)
+            truncate_ledger_db(ledger_db.clone(), overall_commit_progress)
                 .expect("Failed to truncate ledger db.");
 
             // State K/V commit progress isn't (can't be) written atomically with the data,
@@ -427,16 +427,18 @@ impl StateStore {
                     assert_le!(difference, MAX_COMMIT_PROGRESS_DIFFERENCE);
                 }
             }
-            let db = state_merkle_db.metadata_db();
-            let state_merkle_target_version =
-                find_tree_root_at_or_before(db, &state_merkle_db, overall_commit_progress)
-                    .expect("DB read failed.")
-                    .unwrap_or_else(|| {
-                        panic!(
+            let state_merkle_target_version = find_tree_root_at_or_before(
+                ledger_metadata_db,
+                &state_merkle_db,
+                overall_commit_progress,
+            )
+            .expect("DB read failed.")
+            .unwrap_or_else(|| {
+                panic!(
                     "Could not find a valid root before or at version {}, maybe it was pruned?",
                     overall_commit_progress
                 )
-                    });
+            });
             if state_merkle_target_version < state_merkle_max_version {
                 info!(
                     state_merkle_max_version = state_merkle_max_version,
diff --git a/storage/aptosdb/src/utils/truncation_helper.rs b/storage/aptosdb/src/utils/truncation_helper.rs
index a3edf7580ab9c..68f1e02c3e88c 100644
--- a/storage/aptosdb/src/utils/truncation_helper.rs
+++ b/storage/aptosdb/src/utils/truncation_helper.rs
@@ -4,7 +4,7 @@
 #![allow(dead_code)]
 
 use crate::{
-    ledger_db::{LedgerDb, LedgerDbSchemaBatches},
+    ledger_db::{ledger_metadata_db::LedgerMetadataDb, LedgerDb, LedgerDbSchemaBatches},
     schema::{
         db_metadata::{DbMetadataKey, DbMetadataSchema, DbMetadataValue},
         epoch_by_version::EpochByVersionSchema,
@@ -201,7 +201,7 @@ pub(crate) fn truncate_state_merkle_db_single_shard(
 }
 
 pub(crate) fn find_tree_root_at_or_before(
-    ledger_metadata_db: &DB,
+    ledger_metadata_db: &LedgerMetadataDb,
     state_merkle_db: &StateMerkleDb,
     version: Version,
 ) -> Result<Option<Version>> {
@@ -214,8 +214,11 @@ pub(crate) fn find_tree_root_at_or_before(
 
         // It's possible that it's a partial commit when sharding is not enabled,
         // look again for the previous version:
+        if version == 0 {
+            return Ok(None);
+        }
         if let Some(closest_version) =
-            find_closest_node_version_at_or_before(state_merkle_db.metadata_db(), version)?
+            find_closest_node_version_at_or_before(state_merkle_db.metadata_db(), version - 1)?
         {
             if root_exists_at_version(state_merkle_db, closest_version)? {
                 return Ok(Some(closest_version));
@@ -223,7 +226,7 @@ pub(crate) fn find_tree_root_at_or_before(
 
             // Now we are probably looking at a pruned version in this epoch, look for the previous
             // epoch ending:
-            let mut iter = ledger_metadata_db.iter::<EpochByVersionSchema>()?;
+            let mut iter = ledger_metadata_db.db().iter::<EpochByVersionSchema>()?;
             iter.seek_for_prev(&version)?;
             if let Some((closest_epoch_version, _)) = iter.next().transpose()? {
                 if root_exists_at_version(state_merkle_db, closest_epoch_version)? {
diff --git a/testsuite/smoke-test/src/storage.rs b/testsuite/smoke-test/src/storage.rs
index db753d0c38ecb..c01da934cdb77 100644
--- a/testsuite/smoke-test/src/storage.rs
+++ b/testsuite/smoke-test/src/storage.rs
@@ -530,7 +530,7 @@ async fn test_db_restart() {
         quit_flag.clone(),
     ));
 
-    for round in 0..3 {
+    for round in 0..10 {
         info!("{LINE} Restart round {round}");
         for (v, vid) in restarting_validator_ids.iter().enumerate() {
             let validator = swarm.validator_mut(*vid).unwrap();