[Hot State] Add a separate StateMerkleDb for hot state

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

@@ -25,13 +25,14 @@ impl DBDebuggerInterface {
         Ok(Self(Arc::new(
             AptosDB::open(
                 StorageDirPaths::from_path(db_root_path),
-                true,
+                /* readonly = */ true,
                 NO_OP_STORAGE_PRUNER_CONFIG,
                 RocksdbConfigs::default(),
-                false, /* indexer */
+                /* enable_indexer = */ false,
                 BUFFERED_STATE_TARGET_ITEMS,
                 DEFAULT_MAX_NUM_NODES_PER_LRU_CACHE_SHARD,
-                None,
+                /* internal_indexer_db = */ 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,
         }
     }
 }
@@ -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(())
@@ -740,10 +785,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

@@ -170,13 +170,14 @@ impl GenesisInfo {
         let path = TempPath::new();
         let aptosdb = AptosDB::open(
             StorageDirPaths::from_path(path),
-            false,
+            /* readonly = */ false,
             NO_OP_STORAGE_PRUNER_CONFIG,
             RocksdbConfigs::default(),
-            false, /* indexer */
+            /* enable_indexer = */ false,
             BUFFERED_STATE_TARGET_ITEMS,
             DEFAULT_MAX_NUM_NODES_PER_LRU_CACHE_SHARD,
-            None,
+            /* internal_indexer_db = */ 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

@@ -154,13 +154,14 @@ impl MainnetGenesisInfo {
         let path = TempPath::new();
         let aptosdb = AptosDB::open(
             StorageDirPaths::from_path(path),
-            false,
+            /* readonly = */ false,
             NO_OP_STORAGE_PRUNER_CONFIG,
             RocksdbConfigs::default(),
-            false, /* indexer */
+            /* enable_indexer = */ false,
             BUFFERED_STATE_TARGET_ITEMS,
             DEFAULT_MAX_NUM_NODES_PER_LRU_CACHE_SHARD,
-            None,
+            /* internal_indexer_db = */ 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

@@ -100,10 +100,11 @@ pub(crate) fn bootstrap_with_genesis(
             false, /* readonly */
             NO_OP_STORAGE_PRUNER_CONFIG,
             rocksdb_configs,
-            false, /* indexer */
+            false, /* enable_indexer */
             BUFFERED_STATE_TARGET_ITEMS,
             DEFAULT_MAX_NUM_NODES_PER_LRU_CACHE_SHARD,
-            None,
+            None, /* internal_indexer_db */
+            true, /* reset_hot_state */
         )
         .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,
+            )
         })
     }
 

storage/aptosdb/src/db/aptosdb_test.rs

@@ -256,6 +256,7 @@ pub fn test_state_merkle_pruning_impl(
         BUFFERED_STATE_TARGET_ITEMS_FOR_TEST,
         DEFAULT_MAX_NUM_NODES_PER_LRU_CACHE_SHARD,
         None,
+        /* reset_hot_state = */ true,
     )
     .unwrap();
 

storage/aptosdb/src/db/aptosdb_testonly.rs

@@ -57,6 +57,7 @@ impl AptosDB {
             BUFFERED_STATE_TARGET_ITEMS_FOR_TEST,
             max_node_cache,
             None,
+            /* reset_hot_state = */ true,
         )
         .expect("Unable to open AptosDB")
     }