[Hot State] Compute root hash for hot state

Author: wqfish

Cargo.lock

@@ -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-types/src/execution_output.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/block_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/chunk_executor/mod.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_get_execution_output.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/do_state_checkpoint.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(

execution/executor/src/workflow/mod.rs

@@ -20,6 +20,7 @@ aptos-crypto = { workspace = true }
 aptos-db-indexer = { workspace = true }
 aptos-db-indexer-schemas = { workspace = true, features = ["fuzzing"] }
 aptos-executor-types = { workspace = true }
+aptos-experimental-layered-map = { workspace = true }
 aptos-experimental-runtimes = { workspace = true }
 aptos-infallible = { workspace = true }
 aptos-jellyfish-merkle = { workspace = true }

storage/aptosdb/Cargo.toml

@@ -267,7 +267,7 @@ pub(crate) fn save_transactions_impl(
     }
 
     if kv_replay && first_version > 0 && state_store.get_usage(Some(first_version - 1)).is_ok() {
-        let ledger_state = state_store.calculate_state_and_put_updates(
+        let (ledger_state, _hot_state_updates) = state_store.calculate_state_and_put_updates(
             &StateUpdateRefs::index_write_sets(first_version, write_sets, write_sets.len(), vec![]),
             &mut ledger_db_batch.ledger_metadata_db_batches, // used for storing the storage usage
             state_kv_batches,

storage/aptosdb/src/backup/restore_utils.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(),
         )?;