[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. (We don't have to do this in tools/tests though,
just the nodes.)

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 = */ false,
         )
         .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

@@ -29,25 +29,26 @@ pub const BUFFERED_STATE_TARGET_ITEMS_FOR_TEST: usize = 10;
 
 #[derive(Clone, Debug, Default, Deserialize, PartialEq, Eq, Serialize)]
 #[serde(default, deny_unknown_fields)]
-pub struct DbPathConfig {
-    pub ledger_db_path: Option<PathBuf>,
-    pub state_kv_db_path: Option<ShardedDbPathConfig>,
-    pub state_merkle_db_path: Option<ShardedDbPathConfig>,
-    pub hot_state_kv_db_path: Option<ShardedDbPathConfig>,
+struct DbPathConfig {
+    ledger_db_path: Option<PathBuf>,
+    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)]
 #[serde(deny_unknown_fields)]
-pub struct ShardedDbPathConfig {
-    pub metadata_path: Option<PathBuf>,
-    pub shard_paths: Vec<ShardPathConfig>,
+struct ShardedDbPathConfig {
+    metadata_path: Option<PathBuf>,
+    shard_paths: Vec<ShardPathConfig>,
 }
 
 #[derive(Clone, Debug, Deserialize, PartialEq, Eq, Serialize)]
 #[serde(deny_unknown_fields)]
-pub struct ShardPathConfig {
-    pub shards: String,
-    pub path: PathBuf,
+struct ShardPathConfig {
+    shards: String,
+    path: PathBuf,
 }
 
 impl ShardedDbPathConfig {
@@ -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,
         }
     }
 }
@@ -284,7 +288,7 @@ pub struct StorageConfig {
     /// Fine grained control for db paths of individal databases/shards.
     /// If not specificed, will use `dir` as default.
     /// Only allowed when sharding is enabled.
-    pub db_path_overrides: Option<DbPathConfig>,
+    db_path_overrides: Option<DbPathConfig>,
     /// ensure `ulimit -n`, set to 0 to not ensure.
     pub ensure_rlimit_nofile: u64,
     /// panic if failed to ensure `ulimit -n`
@@ -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,
         }
     }
 }
@@ -741,10 +769,8 @@ impl ConfigSanitizer for StorageConfig {
 
 #[cfg(test)]
 mod test {
-    use crate::config::{
-        config_optimizer::ConfigOptimizer, NodeConfig, NodeType, PrunerConfig, ShardPathConfig,
-        ShardedDbPathConfig, StorageConfig,
-    };
+    use super::{ShardPathConfig, ShardedDbPathConfig, StorageConfig};
+    use crate::config::{config_optimizer::ConfigOptimizer, NodeConfig, NodeType, PrunerConfig};
     use aptos_types::chain_id::ChainId;
 
     #[test]

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 = */ false,
         )?;
         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 = */ false,
         )?;
         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 = */ false,
         )
         .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 = */ false,
         )
         .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 = */ false,
     )
     .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 = */ false,
         )
         .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,
+            )
         })
     }