[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.

Author: wqfish

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::{HotStateValue, 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},
@@ -209,10 +212,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,10 +225,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.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()));
                 },
             }
         }
@@ -234,7 +239,7 @@ impl VersionState {
             let shard_id = k.get_shard_id();
             if let Some(slot) = hot_state[shard_id].get_mut(k) {
                 slot.refresh(version);
-                continue;
+                // continue;
             }
             let slot = match state.get(k) {
                 Some((value_version, value)) => StateSlot::HotOccupied {
@@ -248,7 +253,7 @@ impl VersionState {
                     lru_info: LRUEntry::uninitialized(),
                 },
             };
-            hot_smt_updates.push((k.hash(), slot.as_state_value_opt().map(|v| v.hash())));
+            hot_smt_updates.push((k.hash(), Some(HotStateValue::clone_from_slot(&slot).hash())));
             hot_state[shard_id].put(k.clone(), slot);
         }
 
@@ -268,6 +273,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,

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,7 +204,7 @@ 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)
@@ -215,7 +219,7 @@ impl State {
                         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 {
@@ -264,16 +268,18 @@ impl State {
         )
     }
 
+    /// Applies the update the returns the `HotStateValue` that will later go into the hot state
+    /// Merkle tree.
     fn apply_one_update(
         lru: &mut HotStateLRU,
         overlay: &LayeredMap<StateKey, StateSlot>,
         read_cache: &StateCacheShard,
         key: &StateKey,
         update: &StateUpdateRef,
-    ) -> Option<StateValue> {
+    ) -> 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 HotStateValue::new(state_value_opt.cloned(), update.version);
         }
 
         if let Some(mut slot) = lru.get_slot(key) {
@@ -283,14 +289,15 @@ impl State {
             } else {
                 slot.to_hot(update.version)
             };
-            let ret = slot_to_insert.as_state_value_opt().cloned();
+            let ret = HotStateValue::clone_from_slot(&slot_to_insert);
             lru.insert((*key).clone(), slot_to_insert);
             ret
         } 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));
+            let slot = slot.to_hot(update.version);
+            let ret = HotStateValue::clone_from_slot(&slot);
+            lru.insert((*key).clone(), slot);
             ret
         }
     }

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

@@ -125,8 +125,8 @@ impl StateSummary {
                 shard
                     .insertions
                     .iter()
-                    .map(|(k, value_opt)| (k, value_opt.as_ref().map(|v| v.hash())))
-                    .chain(shard.evictions.keys().map(|k| (k, None)))
+                    .map(|(k, item)| (k, Some(item.hash())))
+                    .chain(shard.evictions.iter().map(|k| (k, None)))
                     .sorted_by_key(|(k, _)| k.crypto_hash_ref())
                     .collect_vec()
             })
@@ -297,23 +297,20 @@ impl<'db> ProvableStateSummary<'db> {
         root_depth: usize,
         use_hot_state: bool,
     ) -> Result<SparseMerkleProofExt> {
-        if rand::random::<usize>() % 10000 == 0 {
+        // TODO(HotState): we cannot verify proof yet. In order to verify the proof, we need to
+        // fetch and construct the corresponding `HotStateValue` for `key` at `version`, including
+        // `hot_since_version`. However, the current in-memory hot state does not support this
+        // query, and we might need persist hot state KV to db first.
+        if !use_hot_state && rand::random::<usize>() % 10000 == 0 {
             // 1 out of 10000 times, verify the proof.
             let (val_opt, proof) = self
                 .db
                 // check the full proof
                 .get_state_value_with_proof_by_version_ext(
-                    key,
-                    version,
-                    /* root_depth = */ 0,
-                    use_hot_state,
+                    key, version, /* root_depth = */ 0, /* use_hot_state = */ false,
                 )?;
             proof.verify(
-                if use_hot_state {
-                    self.state_summary.hot_state_summary.root_hash()
-                } else {
-                    self.state_summary.global_state_summary.root_hash()
-                },
+                self.state_summary.global_state_summary.root_hash(),
                 *key,
                 val_opt.as_ref(),
             )?;

types/src/state_store/hot_state.rs

@@ -1,6 +1,17 @@
 // Copyright (c) Aptos Foundation
 // Licensed pursuant to the Innovation-Enabling Source Code License, available at https://github.com/aptos-labs/aptos-core/blob/main/LICENSE
 
+use crate::{
+    state_store::{state_slot::StateSlot, state_value::StateValue},
+    transaction::Version,
+};
+use aptos_crypto::{
+    hash::{CryptoHash, CryptoHasher},
+    HashValue,
+};
+use aptos_crypto_derive::{BCSCryptoHash, CryptoHasher};
+use serde::{Deserialize, Serialize};
+
 #[derive(Clone, Debug, Eq, PartialEq)]
 pub struct LRUEntry<K> {
     /// The key that is slightly newer than the current entry. `None` for the newest entry.
@@ -29,3 +40,105 @@ pub trait THotStateSlot {
     fn set_prev(&mut self, prev: Option<Self::Key>);
     fn set_next(&mut self, next: Option<Self::Key>);
 }
+
+/// `HotStateValue` is what gets hashed into the hot state Merkle tree.
+#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize, BCSCryptoHash, CryptoHasher)]
+pub struct HotStateValue {
+    /// `Some` means occupied and `None` means vacant.
+    value: Option<StateValue>,
+    hot_since_version: Version,
+}
+
+impl HotStateValue {
+    pub fn new(value: Option<StateValue>, hot_since_version: Version) -> Self {
+        Self {
+            value,
+            hot_since_version,
+        }
+    }
+
+    pub fn clone_from_slot(slot: &StateSlot) -> Self {
+        match slot {
+            StateSlot::HotOccupied {
+                value,
+                hot_since_version,
+                ..
+            } => Self::new(Some(value.clone()), *hot_since_version),
+            StateSlot::HotVacant {
+                hot_since_version, ..
+            } => Self::new(None, *hot_since_version),
+            _ => panic!("Must be hot slot"),
+        }
+    }
+}
+
+/// A reference-based version of `HotStateValue` that avoids cloning `StateValue`.
+/// When hashed, it produces the same hash as the equivalent `HotStateValue`.
+#[derive(Serialize)]
+pub struct HotStateValueRef<'a> {
+    value: Option<&'a StateValue>,
+    hot_since_version: Version,
+}
+
+impl<'a> HotStateValueRef<'a> {
+    pub fn new(value: Option<&'a StateValue>, hot_since_version: Version) -> Self {
+        Self {
+            value,
+            hot_since_version,
+        }
+    }
+
+    pub fn from_slot(slot: &'a StateSlot) -> Self {
+        match slot {
+            StateSlot::HotOccupied {
+                value,
+                hot_since_version,
+                ..
+            } => Self::new(Some(value), *hot_since_version),
+            StateSlot::HotVacant {
+                hot_since_version, ..
+            } => Self::new(None, *hot_since_version),
+            _ => panic!("Must be hot slot"),
+        }
+    }
+}
+
+impl CryptoHash for HotStateValueRef<'_> {
+    type Hasher = HotStateValueHasher;
+
+    fn hash(&self) -> HashValue {
+        let mut state = Self::Hasher::default();
+        bcs::serialize_into(&mut state, &self)
+            .expect("BCS serialization of HotStateValueRef should not fail");
+        state.finish()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::{
+        state_store::{
+            hot_state::{HotStateValue, HotStateValueRef},
+            state_value::StateValue,
+        },
+        transaction::Version,
+    };
+    use aptos_crypto::hash::CryptoHash;
+    use proptest::prelude::*;
+
+    proptest! {
+        #[test]
+        fn test_hot_state_value_ref_hash(
+            state_value in any::<StateValue>(),
+            hot_since_version in any::<Version>(),
+        ) {
+            let owned = HotStateValue::new(Some(state_value.clone()), hot_since_version);
+            let borrowed = HotStateValueRef::new(Some(&state_value), hot_since_version);
+            assert_eq!(owned.hash(), borrowed.hash());
+
+            let owned_none = HotStateValue::new(None, hot_since_version);
+            let borrowed_none = HotStateValueRef::new(None, hot_since_version);
+            assert_eq!(owned_none.hash(), borrowed_none.hash());
+        }
+    }
+}