[Hot State] Include HotVacant slots in root hash

In #18390, we simply set the format of the hot state Merkle tree to
`Map<StateKey, StateValue>`. This means that for `HotVacant` entries, for
instance, entries that are read recently and determined to not exist in storage,
we do not hash them into the tree. (And if they existed before, we remove them
from the Merkle tree.) This is inaccurate.

This commit introduces the `HotStateValue` struct and changes the format to
`Map<StateKey, HotStateValue>`. This way, both `HotOccupied` and `HotVacant`
entries are summarized into the root hash.

In addition, we now check refresh interval and avoid refreshing read-only keys
every block.

Author: wqfish

aptos-move/block-executor/src/hot_state_op_accumulator.rs

@@ -18,9 +18,6 @@ pub struct BlockHotStateOpAccumulator<Key> {
     writes: hashbrown::HashSet<Key>,
     /// To prevent the block epilogue from being too heavy.
     max_promotions_per_block: usize,
-    /// Every now and then refresh `hot_since_version` for hot items to prevent them from being
-    /// evicted.
-    _refresh_interval_versions: usize,
 }
 
 impl<Key> BlockHotStateOpAccumulator<Key>
@@ -29,25 +26,16 @@ where
 {
     /// TODO(HotState): make on-chain config
     const MAX_PROMOTIONS_PER_BLOCK: usize = 1024 * 10;
-    /// TODO(HotState): make on-chain config
-    const REFRESH_INTERVAL_VERSIONS: usize = 1_000_000;
 
     pub fn new() -> Self {
-        Self::new_with_config(
-            Self::MAX_PROMOTIONS_PER_BLOCK,
-            Self::REFRESH_INTERVAL_VERSIONS,
-        )
+        Self::new_with_config(Self::MAX_PROMOTIONS_PER_BLOCK)
     }
 
-    pub fn new_with_config(
-        max_promotions_per_block: usize,
-        refresh_interval_versions: usize,
-    ) -> Self {
+    pub fn new_with_config(max_promotions_per_block: usize) -> Self {
         Self {
             to_make_hot: BTreeSet::new(),
             writes: hashbrown::HashSet::new(),
             max_promotions_per_block,
-            _refresh_interval_versions: refresh_interval_versions,
         }
     }
 

config/src/config/storage_config.rs

@@ -243,6 +243,9 @@ impl Default for RocksdbConfigs {
 pub struct HotStateConfig {
     /// Max number of items in each shard.
     pub max_items_per_shard: usize,
+    /// Every now and then refresh `hot_since_version` for hot items to prevent them from being
+    /// evicted.
+    pub refresh_interval_versions: u64,
     /// Whether to delete persisted data on disk on restart. Used during development.
     pub delete_on_restart: bool,
     /// Whether we compute root hashes for hot state in executor and commit the resulting JMT to
@@ -254,6 +257,7 @@ impl Default for HotStateConfig {
     fn default() -> Self {
         Self {
             max_items_per_shard: 250_000,
+            refresh_interval_versions: 100_000,
             delete_on_restart: true,
             compute_root_hash: true,
         }

storage/aptosdb/src/state_store/state_snapshot_committer.rs

@@ -25,7 +25,7 @@ use aptos_storage_interface::{
     jmt_update_refs, state_store::state_with_summary::StateWithSummary, Result,
 };
 use aptos_types::{
-    state_store::{state_key::StateKey, NUM_STATE_SHARDS},
+    state_store::{hot_state::HotStateValueRef, state_key::StateKey, NUM_STATE_SHARDS},
     transaction::Version,
 };
 use rayon::prelude::*;
@@ -99,6 +99,7 @@ impl StateSnapshotCommitter {
                         .map(|(v, _e)| v);
                     let min_version = self.last_snapshot.next_version();
 
+                    // Element format: (key_hash, Option<(value_hash, key)>)
                     let (hot_updates, all_updates): (Vec<_>, Vec<_>) = snapshot
                         .make_delta(&self.last_snapshot)
                         .shards
@@ -111,8 +112,10 @@ impl StateSnapshotCommitter {
                                 if slot.is_hot() {
                                     hot_updates.push((
                                         CryptoHash::hash(&key),
-                                        slot.as_state_value_opt()
-                                            .map(|value| (CryptoHash::hash(value), key.clone())),
+                                        Some((
+                                            HotStateValueRef::from_slot(&slot).hash(),
+                                            key.clone(),
+                                        )),
                                     ));
                                 } else {
                                     hot_updates.push((CryptoHash::hash(&key), None));

storage/aptosdb/src/state_store/tests/speculative_state_workflow.rs

@@ -21,9 +21,12 @@ use aptos_storage_interface::{
 use aptos_types::{
     proof::SparseMerkleProofExt,
     state_store::{
-        hot_state::LRUEntry, state_key::StateKey, state_slot::StateSlot,
-        state_storage_usage::StateStorageUsage, state_value::StateValue, StateViewId,
-        StateViewResult, TStateView, NUM_STATE_SHARDS,
+        hot_state::{HotStateValueRef, LRUEntry},
+        state_key::StateKey,
+        state_slot::StateSlot,
+        state_storage_usage::StateStorageUsage,
+        state_value::StateValue,
+        StateViewId, StateViewResult, TStateView, NUM_STATE_SHARDS,
     },
     transaction::Version,
     write_set::{BaseStateOp, HotStateOp, WriteOp},
@@ -52,10 +55,11 @@ const HOT_STATE_MAX_ITEMS_PER_SHARD: usize = NUM_KEYS / NUM_STATE_SHARDS / 2;
 
 // TODO(HotState): these are not used much at the moment.
 const MAX_PROMOTIONS_PER_BLOCK: usize = NUM_KEYS;
-const REFRESH_INTERVAL_VERSIONS: usize = 50;
+const REFRESH_INTERVAL_VERSIONS: Version = 50;
 
 const TEST_CONFIG: HotStateConfig = HotStateConfig {
     max_items_per_shard: HOT_STATE_MAX_ITEMS_PER_SHARD,
+    refresh_interval_versions: REFRESH_INTERVAL_VERSIONS,
     delete_on_restart: false,
     compute_root_hash: true,
 };
@@ -209,10 +213,11 @@ impl VersionState {
                         hot_since_version: version,
                         lru_info: LRUEntry::uninitialized(),
                     };
+                    hot_smt_updates
+                        .push((k.hash(), Some(HotStateValueRef::from_slot(&slot).hash())));
                     hot_state[shard_id].put(k.clone(), slot);
-                    state.remove(k);
-                    hot_smt_updates.push((k.hash(), None));
                     smt_updates.push((k.hash(), None));
+                    state.remove(k);
                 },
                 Some(v) => {
                     let slot = StateSlot::HotOccupied {
@@ -221,35 +226,47 @@ impl VersionState {
                         hot_since_version: version,
                         lru_info: LRUEntry::uninitialized(),
                     };
+                    hot_smt_updates
+                        .push((k.hash(), Some(HotStateValueRef::from_slot(&slot).hash())));
                     hot_state[shard_id].put(k.clone(), slot);
-                    state.insert(k.clone(), (version, v.clone()));
-                    hot_smt_updates.push((k.hash(), Some(v.hash())));
                     smt_updates.push((k.hash(), Some(v.hash())));
+                    state.insert(k.clone(), (version, v.clone()));
                 },
             }
         }
 
         for k in promotions {
             assert!(is_checkpoint, "No promotions unless in checkpoints");
             let shard_id = k.get_shard_id();
-            if let Some(slot) = hot_state[shard_id].get_mut(k) {
-                slot.refresh(version);
-                continue;
+
+            let hot_value_hash;
+            if let Some(slot) = hot_state[shard_id].peek_mut(k) {
+                if slot.expect_hot_since_version() + REFRESH_INTERVAL_VERSIONS <= version {
+                    slot.refresh(version);
+                    hot_value_hash = Some(HotStateValueRef::from_slot(slot).hash());
+                } else {
+                    hot_value_hash = None;
+                }
+            } else {
+                let slot = match state.get(k) {
+                    Some((value_version, value)) => StateSlot::HotOccupied {
+                        value_version: *value_version,
+                        value: value.clone(),
+                        hot_since_version: version,
+                        lru_info: LRUEntry::uninitialized(),
+                    },
+                    None => StateSlot::HotVacant {
+                        hot_since_version: version,
+                        lru_info: LRUEntry::uninitialized(),
+                    },
+                };
+                hot_value_hash = Some(HotStateValueRef::from_slot(&slot).hash());
+                hot_state[shard_id].put(k.clone(), slot);
+            }
+            if hot_value_hash.is_some() {
+                assert!(hot_state[shard_id].promote(k));
+                hot_smt_updates.push((k.hash(), hot_value_hash));
             }
-            let slot = match state.get(k) {
-                Some((value_version, value)) => StateSlot::HotOccupied {
-                    value_version: *value_version,
-                    value: value.clone(),
-                    hot_since_version: version,
-                    lru_info: LRUEntry::uninitialized(),
-                },
-                None => StateSlot::HotVacant {
-                    hot_since_version: version,
-                    lru_info: LRUEntry::uninitialized(),
-                },
-            };
-            hot_smt_updates.push((k.hash(), slot.as_state_value_opt().map(|v| v.hash())));
-            hot_state[shard_id].put(k.clone(), slot);
         }
 
         if is_checkpoint {
@@ -268,6 +285,11 @@ impl VersionState {
         let bytes = state.iter().map(|(k, v)| k.size() + v.1.size()).sum();
         let usage = StateStorageUsage::new(items, bytes);
 
+        assert_eq!(
+            hot_state.iter().map(|x| x.len()).sum::<usize>(),
+            hot_summary.leaves.len()
+        );
+
         Self {
             usage,
             hot_state,
@@ -660,10 +682,8 @@ fn naive_run_blocks(blocks: Vec<(Vec<UserTxn>, bool)>) -> (Vec<Txn>, StateByVers
     let mut all_txns = vec![];
     let mut state_by_version = StateByVersion::new_empty();
     for (block_txns, append_epilogue) in blocks {
-        let mut op_accu = BlockHotStateOpAccumulator::<StateKey>::new_with_config(
-            MAX_PROMOTIONS_PER_BLOCK,
-            REFRESH_INTERVAL_VERSIONS,
-        );
+        let mut op_accu =
+            BlockHotStateOpAccumulator::<StateKey>::new_with_config(MAX_PROMOTIONS_PER_BLOCK);
         let num_txns = block_txns.len();
         for (idx, txn) in block_txns.into_iter().enumerate() {
             // No promotions except for block epilogue. Also note that in case of reconfig, there's

storage/storage-interface/src/state_store/mod.rs

@@ -10,22 +10,18 @@ pub mod state_view;
 pub mod state_with_summary;
 pub mod versioned_state_value;
 
-use aptos_types::state_store::{state_key::StateKey, state_value::StateValue, NUM_STATE_SHARDS};
-use std::collections::HashMap;
+use aptos_types::state_store::{hot_state::HotStateValue, state_key::StateKey, NUM_STATE_SHARDS};
+use std::collections::{HashMap, HashSet};
 
-// TODO(HotState): this isn't the final form yet. For instance, `HotVacant` entries probably should
-// be included, but right now they are `None` in these and will be removed from the Merkle trees.
 #[derive(Debug)]
 pub(crate) struct HotStateShardUpdates {
-    insertions: HashMap<StateKey, Option<StateValue>>,
-    evictions: HashMap<StateKey, Option<StateValue>>,
+    insertions: HashMap<StateKey, HotStateValue>,
+    // TODO(HotState): only keys are needed for now, since evictions do not affect cold state.
+    evictions: HashSet<StateKey>,
 }
 
 impl HotStateShardUpdates {
-    pub fn new(
-        insertions: HashMap<StateKey, Option<StateValue>>,
-        evictions: HashMap<StateKey, Option<StateValue>>,
-    ) -> Self {
+    pub fn new(insertions: HashMap<StateKey, HotStateValue>, evictions: HashSet<StateKey>) -> Self {
         Self {
             insertions,
             evictions,

storage/storage-interface/src/state_store/state.rs

@@ -24,16 +24,20 @@ use aptos_experimental_layered_map::{LayeredMap, MapLayer};
 use aptos_metrics_core::TimerHelper;
 use aptos_types::{
     state_store::{
-        state_key::StateKey, state_slot::StateSlot, state_storage_usage::StateStorageUsage,
-        state_value::StateValue, StateViewId, NUM_STATE_SHARDS,
+        hot_state::HotStateValue, state_key::StateKey, state_slot::StateSlot,
+        state_storage_usage::StateStorageUsage, StateViewId, NUM_STATE_SHARDS,
     },
     transaction::Version,
 };
 use arr_macro::arr;
 use derive_more::Deref;
 use itertools::Itertools;
 use rayon::prelude::*;
-use std::{collections::HashMap, num::NonZeroUsize, sync::Arc};
+use std::{
+    collections::{HashMap, HashSet},
+    num::NonZeroUsize,
+    sync::Arc,
+};
 
 #[derive(Clone, Debug)]
 pub struct HotStateMetadata {
@@ -200,30 +204,42 @@ impl State {
                     );
                     let mut all_updates = per_version.iter();
                     let mut insertions = HashMap::new();
-                    let mut evictions = HashMap::new();
+                    let mut evictions = HashSet::new();
                     for ckpt_version in all_checkpoint_versions {
                         for (key, update) in
                             all_updates.take_while_ref(|(_k, u)| u.version <= *ckpt_version)
                         {
                             evictions.remove(*key);
-                            insertions.insert(
-                                (*key).clone(),
-                                Self::apply_one_update(&mut lru, overlay, cache, key, update),
-                            );
+                            if let Some(hot_state_value) = Self::apply_one_update(
+                                &mut lru,
+                                overlay,
+                                cache,
+                                key,
+                                update,
+                                self.hot_state_config.refresh_interval_versions,
+                            ) {
+                                insertions.insert((*key).clone(), hot_state_value);
+                            }
                         }
                         // Only evict at the checkpoints.
                         evictions.extend(lru.maybe_evict().into_iter().map(|(key, slot)| {
                             insertions.remove(&key);
                             assert!(slot.is_hot());
-                            (key, slot.into_state_value_opt())
+                            key
                         }));
                     }
                     for (key, update) in all_updates {
                         evictions.remove(*key);
-                        insertions.insert(
-                            (*key).clone(),
-                            Self::apply_one_update(&mut lru, overlay, cache, key, update),
-                        );
+                        if let Some(hot_state_value) = Self::apply_one_update(
+                            &mut lru,
+                            overlay,
+                            cache,
+                            key,
+                            update,
+                            self.hot_state_config.refresh_interval_versions,
+                        ) {
+                            insertions.insert((*key).clone(), hot_state_value);
+                        }
                     }
 
                     let (new_items, new_head, new_tail, new_num_items) = lru.into_updates();
@@ -264,34 +280,48 @@ impl State {
         )
     }
 
+    /// Applies the update the returns the `HotStateValue` that will later go into the hot state
+    /// Merkle tree. `None` if the op is `MakeHot` and it's determined that refresh is not
+    /// necessary.
     fn apply_one_update(
         lru: &mut HotStateLRU,
         overlay: &LayeredMap<StateKey, StateSlot>,
         read_cache: &StateCacheShard,
         key: &StateKey,
         update: &StateUpdateRef,
-    ) -> Option<StateValue> {
+        refresh_interval: Version,
+    ) -> Option<HotStateValue> {
         if let Some(state_value_opt) = update.state_op.as_state_value_opt() {
             lru.insert((*key).clone(), update.to_result_slot().unwrap());
-            return state_value_opt.cloned();
+            return Some(HotStateValue::new(state_value_opt.cloned(), update.version));
         }
 
         if let Some(mut slot) = lru.get_slot(key) {
+            let mut refreshed = true;
             let slot_to_insert = if slot.is_hot() {
-                slot.refresh(update.version);
+                if slot.expect_hot_since_version() + refresh_interval <= update.version {
+                    slot.refresh(update.version);
+                } else {
+                    refreshed = false;
+                }
                 slot
             } else {
                 slot.to_hot(update.version)
             };
-            let ret = slot_to_insert.as_state_value_opt().cloned();
-            lru.insert((*key).clone(), slot_to_insert);
-            ret
+            if refreshed {
+                let ret = HotStateValue::clone_from_slot(&slot_to_insert);
+                lru.insert((*key).clone(), slot_to_insert);
+                Some(ret)
+            } else {
+                None
+            }
         } else {
             let slot = Self::expect_old_slot(overlay, read_cache, key);
             assert!(slot.is_cold());
-            let ret = slot.as_state_value_opt().cloned();
-            lru.insert((*key).clone(), slot.to_hot(update.version));
-            ret
+            let slot = slot.to_hot(update.version);
+            let ret = HotStateValue::clone_from_slot(&slot);
+            lru.insert((*key).clone(), slot);
+            Some(ret)
         }
     }