Prevent writing to state cache when migrating the database (#7067)

* add an update_cache flag to get_state to have more granular control over when we write to the cache

* State cache tweaks

- add state-cache-headroom flag to control pruning
- prune old epoch boundary states ahead of mid-epoch states
- never prune head block's state
- avoid caching ancestor states unless they are on an epoch boundary

---------

Co-authored-by: Michael Sproul <michael@sigmaprime.io>

beacon_node/beacon_chain/src/attestation_rewards.rs

@@ -48,7 +48,7 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
             .ok_or(BeaconChainError::NoStateForSlot(state_slot))?;
 
         let state = self
-            .get_state(&state_root, Some(state_slot))?
+            .get_state(&state_root, Some(state_slot), true)?
             .ok_or(BeaconChainError::MissingBeaconState(state_root))?;
 
         if state.fork_name_unchecked().altair_enabled() {

beacon_node/beacon_chain/src/attester_cache.rs

@@ -287,6 +287,7 @@ impl AttesterCache {
         Ok(())
     }
 
+    // TODO(holesky) check this out
     /// Read the state identified by `state_root` from the database, advance it to the required
     /// slot, use it to prime the cache and return the values for the provided `slot` and
     /// `committee_index`.
@@ -326,7 +327,7 @@ impl AttesterCache {
         }
 
         let mut state: BeaconState<T::EthSpec> = chain
-            .get_state(&state_root, None)?
+            .get_state(&state_root, None, true)?
             .ok_or(Error::MissingBeaconState(state_root))?;
 
         if state.slot() > slot {

beacon_node/beacon_chain/src/beacon_chain.rs

@@ -821,7 +821,7 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
             .get_blinded_block(&block_root)?
             .ok_or(Error::MissingBeaconBlock(block_root))?;
         let state = self
-            .get_state(&block.state_root(), Some(block.slot()))?
+            .get_state(&block.state_root(), Some(block.slot()), true)?
             .ok_or_else(|| Error::MissingBeaconState(block.state_root()))?;
         let iter = BlockRootsIterator::owned(&self.store, state);
         Ok(std::iter::once(Ok((block_root, block.slot())))
@@ -1348,8 +1348,9 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
         &self,
         state_root: &Hash256,
         slot: Option<Slot>,
+        update_cache: bool,
     ) -> Result<Option<BeaconState<T::EthSpec>>, Error> {
-        Ok(self.store.get_state(state_root, slot)?)
+        Ok(self.store.get_state(state_root, slot, update_cache)?)
     }
 
     /// Return the sync committee at `slot + 1` from the canonical chain.
@@ -1524,8 +1525,9 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
                     })?
                     .ok_or(Error::NoStateForSlot(slot))?;
 
+                // Since this slot is less than the current head, don't cache the state
                 Ok(self
-                    .get_state(&state_root, Some(slot))?
+                    .get_state(&state_root, Some(slot), false)?
                     .ok_or(Error::NoStateForSlot(slot))?)
             }
         }
@@ -6943,7 +6945,7 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
 
             let mut beacon_state = self
                 .store
-                .get_state(&beacon_state_root, Some(beacon_block.slot()))?
+                .get_state(&beacon_state_root, Some(beacon_block.slot()), true)?
                 .ok_or_else(|| {
                     Error::DBInconsistent(format!("Missing state {:?}", beacon_state_root))
                 })?;
@@ -7096,7 +7098,7 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
                 if signed_beacon_block.slot() % T::EthSpec::slots_per_epoch() == 0 {
                     let block = self.get_blinded_block(&block_hash).unwrap().unwrap();
                     let state = self
-                        .get_state(&block.state_root(), Some(block.slot()))
+                        .get_state(&block.state_root(), Some(block.slot()), true)
                         .unwrap()
                         .unwrap();
                     finalized_blocks.insert(state.finalized_checkpoint().root);

beacon_node/beacon_chain/src/builder.rs

@@ -299,7 +299,11 @@ where
             .map_err(|e| descriptive_db_error("genesis block", &e))?
             .ok_or("Genesis block not found in store")?;
         let genesis_state = store
-            .get_state(&genesis_block.state_root(), Some(genesis_block.slot()))
+            .get_state(
+                &genesis_block.state_root(),
+                Some(genesis_block.slot()),
+                true,
+            )
             .map_err(|e| descriptive_db_error("genesis state", &e))?
             .ok_or("Genesis state not found in store")?;
 

beacon_node/beacon_chain/src/canonical_head.rs

@@ -655,6 +655,7 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
                     .get_full_block(&new_view.head_block_root)?
                     .ok_or(Error::MissingBeaconBlock(new_view.head_block_root))?;
 
+                // TODO(holesky) when we calculate the new snapshot we might be missing a state
                 let (_, beacon_state) = self
                     .store
                     .get_advanced_hot_state(
@@ -784,6 +785,12 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
             .execution_status
             .is_optimistic_or_invalid();
 
+        // Update the state cache so it doesn't mistakenly prune the new head.
+        self.store
+            .state_cache
+            .lock()
+            .update_head_block_root(new_cached_head.head_block_root());
+
         // Detect and potentially report any re-orgs.
         let reorg_distance = detect_reorg(
             &old_snapshot.beacon_state,

beacon_node/beacon_chain/src/fork_revert.rs

@@ -117,7 +117,7 @@ pub fn reset_fork_choice_to_finalization<E: EthSpec, Hot: ItemStore<E>, Cold: It
     // Advance finalized state to finalized epoch (to handle skipped slots).
     let finalized_state_root = finalized_block.state_root();
     let mut finalized_state = store
-        .get_state(&finalized_state_root, Some(finalized_block.slot()))
+        .get_state(&finalized_state_root, Some(finalized_block.slot()), true)
         .map_err(|e| format!("Error loading finalized state: {:?}", e))?
         .ok_or_else(|| {
             format!(

beacon_node/beacon_chain/src/light_client_server_cache.rs

@@ -320,7 +320,7 @@ impl<T: BeaconChainTypes> LightClientServerCache<T> {
 
         // Compute the value, handling potential errors.
         let mut state = store
-            .get_state(block_state_root, Some(block_slot))?
+            .get_state(block_state_root, Some(block_slot), false)?
             .ok_or_else(|| {
                 BeaconChainError::DBInconsistent(format!("Missing state {:?}", block_state_root))
             })?;

beacon_node/beacon_chain/src/migrate.rs

@@ -356,7 +356,7 @@ impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> BackgroundMigrator<E, Ho
         let finalized_state_root = notif.finalized_state_root;
         let finalized_block_root = notif.finalized_checkpoint.root;
 
-        let finalized_state = match db.get_state(&finalized_state_root.into(), None) {
+        let finalized_state = match db.get_state(&finalized_state_root.into(), None, true) {
             Ok(Some(state)) => state,
             other => {
                 error!(

beacon_node/beacon_chain/src/test_utils.rs

@@ -890,7 +890,7 @@ where
     pub fn get_hot_state(&self, state_hash: BeaconStateHash) -> Option<BeaconState<E>> {
         self.chain
             .store
-            .load_hot_state(&state_hash.into())
+            .load_hot_state(&state_hash.into(), true)
             .unwrap()
             .map(|(state, _)| state)
     }