Make state cache payload status aware

beacon_node/store/src/hot_cold_store.rs

@@ -1144,10 +1144,13 @@ impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> HotColdDB<E, Hot, Cold>
     pub fn get_advanced_hot_state(
         &self,
         block_root: Hash256,
+        payload_status: StatePayloadStatus,
         max_slot: Slot,
         state_root: Hash256,
     ) -> Result<Option<(Hash256, BeaconState<E>)>, Error> {
-        if let Some(cached) = self.get_advanced_hot_state_from_cache(block_root, max_slot) {
+        if let Some(cached) =
+            self.get_advanced_hot_state_from_cache(block_root, payload_status, max_slot)
+        {
             return Ok(Some(cached));
         }
 
@@ -1169,7 +1172,11 @@ impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> HotColdDB<E, Hot, Cold>
             .into());
         }
 
-        let state_root = if block_root == split.block_root && split.slot <= max_slot {
+        // Split state should always be `Pending`.
+        let state_root = if block_root == split.block_root
+            && let StatePayloadStatus::Pending = payload_status
+            && split.slot <= max_slot
+        {
             split.state_root
         } else {
             state_root
@@ -1216,11 +1223,12 @@ impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> HotColdDB<E, Hot, Cold>
     pub fn get_advanced_hot_state_from_cache(
         &self,
         block_root: Hash256,
+        payload_status: StatePayloadStatus,
         max_slot: Slot,
     ) -> Option<(Hash256, BeaconState<E>)> {
         self.state_cache
             .lock()
-            .get_by_block_root(block_root, max_slot)
+            .get_by_block_root(block_root, payload_status, max_slot)
     }
 
     /// Delete a state, ensuring it is removed from the LRU cache, as well as from on-disk.

beacon_node/store/src/state_cache.rs

@@ -7,7 +7,7 @@ use lru::LruCache;
 use std::collections::{BTreeMap, HashMap, HashSet};
 use std::num::NonZeroUsize;
 use tracing::instrument;
-use types::{BeaconState, ChainSpec, Epoch, EthSpec, Hash256, Slot};
+use types::{BeaconState, ChainSpec, Epoch, EthSpec, Hash256, Slot, execution::StatePayloadStatus};
 
 /// Fraction of the LRU cache to leave intact during culling.
 const CULL_EXEMPT_NUMERATOR: usize = 1;
@@ -23,10 +23,10 @@ pub struct FinalizedState<E: EthSpec> {
     state: BeaconState<E>,
 }
 
-/// Map from block_root -> slot -> state_root.
+/// Map from (block_root, payload_status) -> slot -> state_root.
 #[derive(Debug, Default)]
 pub struct BlockMap {
-    blocks: HashMap<Hash256, SlotMap>,
+    blocks: HashMap<(Hash256, StatePayloadStatus), SlotMap>,
 }
 
 /// Map from slot -> state_root.
@@ -130,8 +130,11 @@ impl<E: EthSpec> StateCache<E> {
             return Err(Error::FinalizedStateDecreasingSlot);
         }
 
+        let payload_status = state.payload_status();
+
         // Add to block map.
-        self.block_map.insert(block_root, state.slot(), state_root);
+        self.block_map
+            .insert(block_root, payload_status, state.slot(), state_root);
 
         // Prune block map.
         let state_roots_to_prune = self.block_map.prune(state.slot());
@@ -254,7 +257,9 @@ impl<E: EthSpec> StateCache<E> {
 
         // Record the connection from block root and slot to this state.
         let slot = state.slot();
-        self.block_map.insert(block_root, slot, state_root);
+        let payload_status = state.payload_status();
+        self.block_map
+            .insert(block_root, payload_status, slot, state_root);
 
         Ok(PutStateOutcome::New(deleted_states))
     }
@@ -303,9 +308,10 @@ impl<E: EthSpec> StateCache<E> {
     pub fn get_by_block_root(
         &mut self,
         block_root: Hash256,
+        payload_status: StatePayloadStatus,
         slot: Slot,
     ) -> Option<(Hash256, BeaconState<E>)> {
-        let slot_map = self.block_map.blocks.get(&block_root)?;
+        let slot_map = self.block_map.blocks.get(&(block_root, payload_status))?;
 
         // Find the state at `slot`, or failing that the most recent ancestor.
         let state_root = slot_map
@@ -399,8 +405,14 @@ impl<E: EthSpec> StateCache<E> {
 }
 
 impl BlockMap {
-    fn insert(&mut self, block_root: Hash256, slot: Slot, state_root: Hash256) {
-        let slot_map = self.blocks.entry(block_root).or_default();
+    fn insert(
+        &mut self,
+        block_root: Hash256,
+        payload_status: StatePayloadStatus,
+        slot: Slot,
+        state_root: Hash256,
+    ) {
+        let slot_map = self.blocks.entry((block_root, payload_status)).or_default();
         slot_map.slots.insert(slot, state_root);
     }
 
@@ -432,7 +444,10 @@ impl BlockMap {
     }
 
     fn delete_block_states(&mut self, block_root: &Hash256) -> Option<SlotMap> {
-        self.blocks.remove(block_root)
+        // TODO(gloas): update return type
+        self.blocks
+            .remove(&(*block_root, StatePayloadStatus::Pending));
+        self.blocks.remove(&(*block_root, StatePayloadStatus::Full))
     }
 }