[Hot State] Add a separate StateMerkleDb for hot state

In order to compute the hot state root hash for each block (including the
speculative ones), we have to be able to return Merkle proof for any given key
and version. The most straightforward way to do this is to add another instance
of `StateMerkleDb` in `StateStore`.

For now, we always delete and reset the DB on restart. This will give us more
flexibility to make changes.

Author: wqfish

api/test-context/src/test_context.rs

@@ -176,6 +176,7 @@ pub fn new_test_context_inner(
             BUFFERED_STATE_TARGET_ITEMS_FOR_TEST,
             DEFAULT_MAX_NUM_NODES_PER_LRU_CACHE_SHARD,
             None,
+            /* reset_hot_state = */ true,
         )
         .unwrap();
         if node_config

aptos-move/aptos-validator-interface/src/storage_interface.rs

@@ -32,6 +32,7 @@ impl DBDebuggerInterface {
                 BUFFERED_STATE_TARGET_ITEMS,
                 DEFAULT_MAX_NUM_NODES_PER_LRU_CACHE_SHARD,
                 None,
+                /* reset_hot_state = */ false,
             )
             .map_err(anyhow::Error::from)?,
         )))

config/src/config/storage_config.rs

@@ -34,6 +34,7 @@ struct DbPathConfig {
     state_kv_db_path: Option<ShardedDbPathConfig>,
     state_merkle_db_path: Option<ShardedDbPathConfig>,
     hot_state_kv_db_path: Option<ShardedDbPathConfig>,
+    hot_state_merkle_db_path: Option<ShardedDbPathConfig>,
 }
 
 #[derive(Clone, Debug, Deserialize, PartialEq, Eq, Serialize)]
@@ -242,12 +243,15 @@ impl Default for RocksdbConfigs {
 pub struct HotStateConfig {
     /// Max number of items in each shard.
     pub max_items_per_shard: usize,
+    /// Whether to delete persisted data on disk on restart. Used during development.
+    pub delete_on_restart: bool,
 }
 
 impl Default for HotStateConfig {
     fn default() -> Self {
         Self {
             max_items_per_shard: 250_000,
+            delete_on_restart: true,
         }
     }
 }
@@ -458,6 +462,7 @@ impl StorageConfig {
         let mut state_kv_db_paths = ShardedDbPaths::default();
         let mut state_merkle_db_paths = ShardedDbPaths::default();
         let mut hot_state_kv_db_paths = ShardedDbPaths::default();
+        let mut hot_state_merkle_db_paths = ShardedDbPaths::default();
 
         if let Some(db_path_overrides) = self.db_path_overrides.as_ref() {
             db_path_overrides
@@ -475,6 +480,12 @@ impl StorageConfig {
             if let Some(hot_state_kv_db_path) = db_path_overrides.hot_state_kv_db_path.as_ref() {
                 hot_state_kv_db_paths = ShardedDbPaths::new(hot_state_kv_db_path);
             }
+
+            if let Some(hot_state_merkle_db_path) =
+                db_path_overrides.hot_state_merkle_db_path.as_ref()
+            {
+                hot_state_merkle_db_paths = ShardedDbPaths::new(hot_state_merkle_db_path);
+            }
         }
 
         StorageDirPaths::new(
@@ -483,6 +494,7 @@ impl StorageConfig {
             state_kv_db_paths,
             state_merkle_db_paths,
             hot_state_kv_db_paths,
+            hot_state_merkle_db_paths,
         )
     }
 
@@ -503,6 +515,7 @@ pub struct StorageDirPaths {
     state_kv_db_paths: ShardedDbPaths,
     state_merkle_db_paths: ShardedDbPaths,
     hot_state_kv_db_paths: ShardedDbPaths,
+    hot_state_merkle_db_paths: ShardedDbPaths,
 }
 
 impl StorageDirPaths {
@@ -548,13 +561,26 @@ impl StorageDirPaths {
             .unwrap_or(&self.default_path)
     }
 
+    pub fn hot_state_merkle_db_metadata_root_path(&self) -> &PathBuf {
+        self.hot_state_merkle_db_paths
+            .metadata_path()
+            .unwrap_or(&self.default_path)
+    }
+
+    pub fn hot_state_merkle_db_shard_root_path(&self, shard_id: usize) -> &PathBuf {
+        self.hot_state_merkle_db_paths
+            .shard_path(shard_id)
+            .unwrap_or(&self.default_path)
+    }
+
     pub fn from_path<P: AsRef<Path>>(path: P) -> Self {
         Self {
             default_path: path.as_ref().to_path_buf(),
             ledger_db_path: None,
             state_kv_db_paths: Default::default(),
             state_merkle_db_paths: Default::default(),
             hot_state_kv_db_paths: Default::default(),
+            hot_state_merkle_db_paths: Default::default(),
         }
     }
 
@@ -564,13 +590,15 @@ impl StorageDirPaths {
         state_kv_db_paths: ShardedDbPaths,
         state_merkle_db_paths: ShardedDbPaths,
         hot_state_kv_db_paths: ShardedDbPaths,
+        hot_state_merkle_db_paths: ShardedDbPaths,
     ) -> Self {
         Self {
             default_path,
             ledger_db_path,
             state_kv_db_paths,
             state_merkle_db_paths,
             hot_state_kv_db_paths,
+            hot_state_merkle_db_paths,
         }
     }
 }
@@ -732,6 +760,23 @@ impl ConfigSanitizer for StorageConfig {
                     return Err(Error::ConfigSanitizerFailed(sanitizer_name, e.to_string()));
                 }
             }
+
+            if let Some(hot_state_merkle_db_path) =
+                db_path_overrides.hot_state_merkle_db_path.as_ref()
+            {
+                if let Some(metadata_path) = hot_state_merkle_db_path.metadata_path.as_ref() {
+                    if !metadata_path.is_absolute() {
+                        return Err(Error::ConfigSanitizerFailed(
+                            sanitizer_name,
+                            format!("Path {metadata_path:?} in db_path_overrides is not an absolute path."),
+                        ));
+                    }
+                }
+
+                if let Err(e) = hot_state_merkle_db_path.get_shard_paths() {
+                    return Err(Error::ConfigSanitizerFailed(sanitizer_name, e.to_string()));
+                }
+            }
         }
 
         Ok(())

crates/aptos-genesis/src/lib.rs

@@ -177,6 +177,7 @@ impl GenesisInfo {
             BUFFERED_STATE_TARGET_ITEMS,
             DEFAULT_MAX_NUM_NODES_PER_LRU_CACHE_SHARD,
             None,
+            /* reset_hot_state = */ true,
         )?;
         let db_rw = DbReaderWriter::new(aptosdb);
         aptos_executor::db_bootstrapper::generate_waypoint::<AptosVMBlockExecutor>(&db_rw, genesis)

crates/aptos-genesis/src/mainnet.rs

@@ -161,6 +161,7 @@ impl MainnetGenesisInfo {
             BUFFERED_STATE_TARGET_ITEMS,
             DEFAULT_MAX_NUM_NODES_PER_LRU_CACHE_SHARD,
             None,
+            /* reset_hot_state = */ true,
         )?;
         let db_rw = DbReaderWriter::new(aptosdb);
         aptos_executor::db_bootstrapper::generate_waypoint::<AptosVMBlockExecutor>(&db_rw, genesis)

execution/executor-benchmark/src/db_generator.rs

@@ -104,6 +104,7 @@ pub(crate) fn bootstrap_with_genesis(
             BUFFERED_STATE_TARGET_ITEMS,
             DEFAULT_MAX_NUM_NODES_PER_LRU_CACHE_SHARD,
             None,
+            /* reset_hot_state = */ true,
         )
         .expect("DB should open."),
     );

execution/executor-benchmark/src/lib.rs

@@ -117,6 +117,7 @@ pub fn init_db(config: &NodeConfig) -> DbReaderWriter {
             config.storage.buffered_state_target_items,
             config.storage.max_num_nodes_per_lru_cache_shard,
             None,
+            /* reset_hot_state = */ true,
         )
         .expect("DB should open."),
     )

state-sync/state-sync-driver/src/tests/driver_factory.rs

@@ -41,6 +41,7 @@ fn test_new_initialized_configs() {
         BUFFERED_STATE_TARGET_ITEMS,
         DEFAULT_MAX_NUM_NODES_PER_LRU_CACHE_SHARD,
         None,
+        /* reset_hot_state = */ true,
     )
     .unwrap();
     let (_, db_rw) = DbReaderWriter::wrap(db);

storage/aptosdb/src/db/aptosdb_internal.rs

@@ -41,6 +41,7 @@ use std::{
 impl AptosDB {
     fn new_with_dbs(
         ledger_db: LedgerDb,
+        hot_state_merkle_db: Option<StateMerkleDb>,
         state_merkle_db: StateMerkleDb,
         state_kv_db: StateKvDb,
         pruner_config: PrunerConfig,
@@ -51,6 +52,7 @@ impl AptosDB {
         internal_indexer_db: Option<InternalIndexerDB>,
     ) -> Self {
         let ledger_db = Arc::new(ledger_db);
+        let hot_state_merkle_db = hot_state_merkle_db.map(Arc::new);
         let state_merkle_db = Arc::new(state_merkle_db);
         let state_kv_db = Arc::new(state_kv_db);
         let state_merkle_pruner = StateMerklePrunerManager::new(
@@ -65,6 +67,7 @@ impl AptosDB {
             StateKvPrunerManager::new(Arc::clone(&state_kv_db), pruner_config.ledger_pruner_config);
         let state_store = Arc::new(StateStore::new(
             Arc::clone(&ledger_db),
+            hot_state_merkle_db,
             Arc::clone(&state_merkle_db),
             Arc::clone(&state_kv_db),
             state_merkle_pruner,
@@ -113,6 +116,7 @@ impl AptosDB {
         max_num_nodes_per_lru_cache_shard: usize,
         empty_buffered_state_for_restore: bool,
         internal_indexer_db: Option<InternalIndexerDB>,
+        reset_hot_state: bool,
     ) -> Result<Self> {
         ensure!(
             pruner_config.eq(&NO_OP_STORAGE_PRUNER_CONFIG) || !readonly,
@@ -128,17 +132,19 @@ impl AptosDB {
             /* estimated_entry_charge = */ 0,
         );
 
-        let (ledger_db, state_merkle_db, state_kv_db) = Self::open_dbs(
+        let (ledger_db, hot_state_merkle_db, state_merkle_db, state_kv_db) = Self::open_dbs(
             db_paths,
             rocksdb_configs,
             Some(&env),
             Some(&block_cache),
             readonly,
             max_num_nodes_per_lru_cache_shard,
+            reset_hot_state,
         )?;
 
         let mut myself = Self::new_with_dbs(
             ledger_db,
+            hot_state_merkle_db,
             state_merkle_db,
             state_kv_db,
             pruner_config,
@@ -243,6 +249,7 @@ impl AptosDB {
             buffered_state_target_items,
             max_num_nodes_per_lru_cache_shard,
             None,
+            /* reset_hot_state = */ true,
         )
         .expect("Unable to open AptosDB")
     }

storage/aptosdb/src/db/aptosdb_reader.rs

@@ -660,12 +660,17 @@ impl DbReader for AptosDB {
         key_hash: &HashValue,
         version: Version,
         root_depth: usize,
+        use_hot_state: bool,
     ) -> Result<SparseMerkleProofExt> {
         gauged_api("get_state_proof_by_version_ext", || {
             self.error_if_state_merkle_pruned("State merkle", version)?;
 
-            self.state_store
-                .get_state_proof_by_version_ext(key_hash, version, root_depth)
+            self.state_store.get_state_proof_by_version_ext(
+                key_hash,
+                version,
+                root_depth,
+                use_hot_state,
+            )
         })
     }
 
@@ -674,12 +679,17 @@ impl DbReader for AptosDB {
         key_hash: &HashValue,
         version: Version,
         root_depth: usize,
+        use_hot_state: bool,
     ) -> Result<(Option<StateValue>, SparseMerkleProofExt)> {
         gauged_api("get_state_value_with_proof_by_version_ext", || {
             self.error_if_state_merkle_pruned("State merkle", version)?;
 
-            self.state_store
-                .get_state_value_with_proof_by_version_ext(key_hash, version, root_depth)
+            self.state_store.get_state_value_with_proof_by_version_ext(
+                key_hash,
+                version,
+                root_depth,
+                use_hot_state,
+            )
         })
     }