[Hot State] Compute root hash for hot state

Author: wqfish

execution/executor-types/src/execution_output.rs

@@ -10,11 +10,12 @@ use crate::{
 use aptos_config::config::HotStateConfig;
 use aptos_drop_helper::DropHelper;
 use aptos_storage_interface::state_store::{
-    state::LedgerState, state_view::cached_state_view::ShardedStateCache,
+    state::LedgerState, state_view::cached_state_view::ShardedStateCache, HotStateShardUpdates,
 };
 use aptos_types::{
     contract_event::ContractEvent,
     epoch_state::EpochState,
+    state_store::NUM_STATE_SHARDS,
     transaction::{
         block_epilogue::BlockEndInfo, ExecutionStatus, Transaction, TransactionStatus, Version,
     },
@@ -41,6 +42,7 @@ impl ExecutionOutput {
         block_end_info: Option<BlockEndInfo>,
         next_epoch_state: Option<EpochState>,
         subscribable_events: Planned<Vec<ContractEvent>>,
+        hot_state_updates: [Vec<HotStateShardUpdates>; NUM_STATE_SHARDS],
     ) -> Self {
         let next_version = first_version + to_commit.len() as Version;
         assert_eq!(next_version, result_state.latest().next_version());
@@ -65,6 +67,7 @@ impl ExecutionOutput {
             block_end_info,
             next_epoch_state,
             subscribable_events,
+            hot_state_updates,
         })
     }
 
@@ -81,6 +84,7 @@ impl ExecutionOutput {
             block_end_info: None,
             next_epoch_state: None,
             subscribable_events: Planned::ready(vec![]),
+            hot_state_updates: std::array::from_fn(|_| vec![]),
         })
     }
 
@@ -99,6 +103,7 @@ impl ExecutionOutput {
             block_end_info: None,
             next_epoch_state: None,
             subscribable_events: Planned::ready(vec![]),
+            hot_state_updates: std::array::from_fn(|_| vec![]),
         })
     }
 
@@ -119,6 +124,7 @@ impl ExecutionOutput {
             block_end_info: None,
             next_epoch_state: self.next_epoch_state.clone(),
             subscribable_events: Planned::ready(vec![]),
+            hot_state_updates: std::array::from_fn(|_| vec![]),
         })
     }
 
@@ -166,6 +172,8 @@ pub struct Inner {
     /// state cache.
     pub next_epoch_state: Option<EpochState>,
     pub subscribable_events: Planned<Vec<ContractEvent>>,
+    /// Per checkpoint and optionally a partial block.
+    pub hot_state_updates: [Vec<HotStateShardUpdates>; NUM_STATE_SHARDS],
 }
 
 impl Inner {

execution/executor/src/block_executor/mod.rs

@@ -315,7 +315,14 @@ where
                 output.set_state_checkpoint_output(DoStateCheckpoint::run(
                     &output.execution_output,
                     parent_block.output.ensure_result_state_summary()?,
-                    &ProvableStateSummary::new_persisted(self.db.reader.as_ref())?,
+                    &ProvableStateSummary::new_persisted(
+                        self.db.reader.as_ref(),
+                        /* is_hot = */ true,
+                    )?,
+                    &ProvableStateSummary::new_persisted(
+                        self.db.reader.as_ref(),
+                        /* is_hot = */ false,
+                    )?,
                     None,
                 )?);
                 output.set_ledger_update_output(DoLedgerUpdate::run(

execution/executor/src/chunk_executor/mod.rs

@@ -346,7 +346,14 @@ impl<V: VMBlockExecutor> ChunkExecutorInner<V> {
         let state_checkpoint_output = DoStateCheckpoint::run(
             &output.execution_output,
             &parent_state_summary,
-            &ProvableStateSummary::new_persisted(self.db.reader.as_ref())?,
+            &ProvableStateSummary::new_persisted(
+                self.db.reader.as_ref(),
+                /* is_hot = */ true,
+            )?,
+            &ProvableStateSummary::new_persisted(
+                self.db.reader.as_ref(),
+                /* is_hot = */ false,
+            )?,
             Some(
                 chunk_verifier
                     .transaction_infos()

execution/executor/src/workflow/do_get_execution_output.rs

@@ -417,7 +417,7 @@ impl Parser {
             },
         )?;
 
-        let result_state = parent_state.update_with_memorized_reads(
+        let (result_state, hot_state_updates) = parent_state.update_with_memorized_reads(
             base_state_view.persisted_hot_state(),
             base_state_view.persisted_state(),
             to_commit.state_update_refs(),
@@ -437,6 +437,7 @@ impl Parser {
             block_end_info,
             next_epoch_state,
             Planned::place_holder(),
+            hot_state_updates,
         );
         let ret = out.clone();
         ret.subscribable_events

execution/executor/src/workflow/do_state_checkpoint.rs

@@ -18,13 +18,16 @@ impl DoStateCheckpoint {
     pub fn run(
         execution_output: &ExecutionOutput,
         parent_state_summary: &LedgerStateSummary,
-        persisted_state_summary: &ProvableStateSummary,
+        hot_persisted_state_summary: &ProvableStateSummary,
+        cold_persisted_state_summary: &ProvableStateSummary,
         known_state_checkpoints: Option<Vec<Option<HashValue>>>,
     ) -> Result<StateCheckpointOutput> {
         let _timer = OTHER_TIMERS.timer_with(&["do_state_checkpoint"]);
 
         let state_summary = parent_state_summary.update(
-            persisted_state_summary,
+            hot_persisted_state_summary,
+            &execution_output.hot_state_updates,
+            cold_persisted_state_summary,
             execution_output.to_commit.state_update_refs(),
         )?;
 

execution/executor/src/workflow/mod.rs

@@ -27,7 +27,8 @@ impl ApplyExecutionOutput {
         let state_checkpoint_output = DoStateCheckpoint::run(
             &execution_output,
             &base_view.state_summary,
-            &ProvableStateSummary::new_persisted(reader)?,
+            &ProvableStateSummary::new_persisted(reader, /* is_hot = */ true)?,
+            &ProvableStateSummary::new_persisted(reader, /* is_hot = */ false)?,
             None,
         )?;
         let ledger_update_output = DoLedgerUpdate::run(

storage/aptosdb/src/db/aptosdb_testonly.rs

@@ -166,15 +166,17 @@ impl AptosDB {
 
         let current = self.state_store.current_state_locked().clone();
         let (hot_state, persisted_state) = self.state_store.get_persisted_state()?;
-        let (new_state, reads) = current.ledger_state().update_with_db_reader(
+        let (new_state, reads, hot_state_updates) = current.ledger_state().update_with_db_reader(
             &persisted_state,
             hot_state,
             transactions_to_keep.state_update_refs(),
             self.state_store.clone(),
         )?;
         let persisted_summary = self.state_store.get_persisted_state_summary()?;
         let new_state_summary = current.ledger_state_summary().update(
-            &ProvableStateSummary::new(persisted_summary, self),
+            &ProvableStateSummary::new(persisted_summary.clone(), self, /* is_hot = */ true),
+            &hot_state_updates,
+            &ProvableStateSummary::new(persisted_summary, self, /* is_hot = */ false),
             transactions_to_keep.state_update_refs(),
         )?;
 

storage/aptosdb/src/state_store/mod.rs

@@ -41,7 +41,11 @@ use aptos_db_indexer_schemas::{
     schema::indexer_metadata::InternalIndexerMetadataSchema,
 };
 use aptos_infallible::Mutex;
-use aptos_jellyfish_merkle::iterator::JellyfishMerkleIterator;
+use aptos_jellyfish_merkle::{
+    iterator::JellyfishMerkleIterator,
+    node_type::{Node, NodeKey},
+    TreeWriter,
+};
 use aptos_logger::info;
 use aptos_metrics_core::TimerHelper;
 use aptos_schemadb::batch::{NativeBatch, SchemaBatch, WriteBatch};
@@ -58,6 +62,7 @@ use aptos_storage_interface::{
         },
         state_with_summary::{LedgerStateWithSummary, StateWithSummary},
         versioned_state_value::StateUpdateRef,
+        HotStateShardUpdates,
     },
     AptosDbError, DbReader, Result, StateSnapshotReceiver,
 };
@@ -79,6 +84,7 @@ use claims::{assert_ge, assert_le};
 use itertools::Itertools;
 use rayon::prelude::*;
 use std::{
+    collections::HashMap,
     ops::Deref,
     sync::{Arc, MutexGuard},
 };
@@ -554,6 +560,7 @@ impl StateStore {
             latest_snapshot_version = latest_snapshot_version,
             "Initializing BufferedState."
         );
+        // TODO(HotState): read hot root hash from DB.
         let latest_snapshot_root_hash = if let Some(version) = latest_snapshot_version {
             state_db
                 .state_merkle_db
@@ -595,6 +602,18 @@ impl StateStore {
             );
         }
 
+        if snapshot_next_version > 0 {
+            let prev_version = snapshot_next_version - 1;
+            let node_key = NodeKey::new_empty_path(prev_version);
+            let node = Node::Null;
+            let mut node_batch = HashMap::new();
+            node_batch.insert(node_key, node);
+            if let Some(db) = &state_db.hot_state_merkle_db {
+                db.write_node_batch(&node_batch)?;
+            }
+            info!("Wrote null node for hot state at version {prev_version}");
+        }
+
         // Replaying the committed write sets after the latest snapshot.
         if snapshot_next_version < num_transactions {
             if check_max_versions_after_snapshot {
@@ -605,6 +624,8 @@ impl StateStore {
                     num_transactions,
                 );
             }
+            info!("Replaying writesets from {snapshot_next_version} to {num_transactions} to let state Merkle DB catch up.");
+
             let write_sets = state_db
                 .ledger_db
                 .write_set_db()
@@ -628,17 +649,28 @@ impl StateStore {
             );
             let current_state = out_current_state.lock().clone();
             let (hot_state, state) = out_persisted_state.get_state();
-            let (new_state, _state_reads) = current_state.ledger_state().update_with_db_reader(
-                &state,
-                hot_state,
-                &state_update_refs,
-                state_db.clone(),
-            )?;
+            let (new_state, _state_reads, hot_state_updates) = current_state
+                .ledger_state()
+                .update_with_db_reader(&state, hot_state, &state_update_refs, state_db.clone())?;
             let state_summary = out_persisted_state.get_state_summary();
             let new_state_summary = current_state.ledger_state_summary().update(
-                &ProvableStateSummary::new(state_summary, state_db.as_ref()),
+                &ProvableStateSummary::new(
+                    state_summary.clone(),
+                    state_db.as_ref(),
+                    /* is_hot = */ true,
+                ),
+                &hot_state_updates,
+                &ProvableStateSummary::new(
+                    state_summary,
+                    state_db.as_ref(),
+                    /* is_hot = */ false,
+                ),
                 &state_update_refs,
             )?;
+            info!(
+                "new hot state summary root hash: {}",
+                new_state_summary.hot_root_hash()
+            );
             let updated =
                 LedgerStateWithSummary::from_state_and_summary(new_state, new_state_summary);
 
@@ -720,7 +752,7 @@ impl StateStore {
     ) -> Result<LedgerState> {
         let current = self.current_state_locked().ledger_state();
         let (hot_state, persisted) = self.get_persisted_state()?;
-        let (new_state, reads) = current.update_with_db_reader(
+        let (new_state, reads, _hot_state_updates) = current.update_with_db_reader(
             &persisted,
             hot_state,
             state_update_refs,
@@ -1403,27 +1435,29 @@ mod test_only {
             let current = self.current_state_locked().ledger_state_summary();
             let persisted = self.persisted_state.get_state_summary();
 
-            let new_state_summary = current
-                .update(
-                    &ProvableStateSummary::new(persisted, self.state_db.as_ref()),
-                    &state_update_refs,
-                )
-                .unwrap();
-            let root_hash = new_state_summary.root_hash();
-
-            self.buffered_state
-                .lock()
-                .update(
-                    LedgerStateWithSummary::from_state_and_summary(
-                        new_ledger_state,
-                        new_state_summary,
-                    ),
-                    0,    /* estimated_items, doesn't matter since we sync-commit */
-                    true, /* sync_commit */
-                )
-                .unwrap();
-
-            root_hash
+            unimplemented!();
+
+            // let new_state_summary = current
+            //     .update(
+            //         &ProvableStateSummary::new(persisted, self.state_db.as_ref()),
+            //         &state_update_refs,
+            //     )
+            //     .unwrap();
+            // let root_hash = new_state_summary.root_hash();
+
+            // self.buffered_state
+            //     .lock()
+            //     .update(
+            //         LedgerStateWithSummary::from_state_and_summary(
+            //             new_ledger_state,
+            //             new_state_summary,
+            //         ),
+            //         0,    /* estimated_items, doesn't matter since we sync-commit */
+            //         true, /* sync_commit */
+            //     )
+            //     .unwrap();
+
+            // root_hash
         }
     }
 }

storage/aptosdb/src/state_store/state_merkle_batch_committer.rs

@@ -25,14 +25,14 @@ pub struct StateMerkleBatch {
 
 pub(crate) struct StateMerkleBatchCommitter {
     state_db: Arc<StateDb>,
-    state_merkle_batch_receiver: Receiver<CommitMessage<StateMerkleBatch>>,
+    state_merkle_batch_receiver: Receiver<CommitMessage<(StateMerkleBatch, StateMerkleBatch)>>,
     persisted_state: PersistedState,
 }
 
 impl StateMerkleBatchCommitter {
     pub fn new(
         state_db: Arc<StateDb>,
-        state_merkle_batch_receiver: Receiver<CommitMessage<StateMerkleBatch>>,
+        state_merkle_batch_receiver: Receiver<CommitMessage<(StateMerkleBatch, StateMerkleBatch)>>,
         persisted_state: PersistedState,
     ) -> Self {
         Self {
@@ -46,7 +46,12 @@ impl StateMerkleBatchCommitter {
         while let Ok(msg) = self.state_merkle_batch_receiver.recv() {
             let _timer = OTHER_TIMERS_SECONDS.timer_with(&["batch_committer_work"]);
             match msg {
-                CommitMessage::Data(state_merkle_batch) => {
+                CommitMessage::Data((hot_state_merkle_batch, state_merkle_batch)) => {
+                    let StateMerkleBatch {
+                        top_levels_batch: hot_top_levels_batch,
+                        batches_for_shards: hot_batches_for_shards,
+                        snapshot,
+                    } = hot_state_merkle_batch;
                     let StateMerkleBatch {
                         top_levels_batch,
                         batches_for_shards,
@@ -61,6 +66,31 @@ impl StateMerkleBatchCommitter {
                     // commit jellyfish merkle nodes
                     let _timer =
                         OTHER_TIMERS_SECONDS.timer_with(&["commit_jellyfish_merkle_nodes"]);
+                    self.state_db
+                        .hot_state_merkle_db
+                        .as_ref()
+                        .unwrap()
+                        .commit(
+                            current_version,
+                            hot_top_levels_batch,
+                            hot_batches_for_shards,
+                        )
+                        .expect("Hot state merkle nodes commit failed.");
+                    if let Some(lru_cache) = self
+                        .state_db
+                        .hot_state_merkle_db
+                        .as_ref()
+                        .unwrap()
+                        .lru_cache()
+                    {
+                        self.state_db
+                            .hot_state_merkle_db
+                            .as_ref()
+                            .unwrap()
+                            .version_caches()
+                            .iter()
+                            .for_each(|(_, cache)| cache.maybe_evict_version(lru_cache));
+                    }
                     self.state_db
                         .state_merkle_db
                         .commit(current_version, top_levels_batch, batches_for_shards)