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/tests/attestation_verification.rs

@@ -1225,7 +1225,7 @@ async fn attestation_that_skips_epochs() {
 
     let mut state = harness
         .chain
-        .get_state(&earlier_block.state_root(), Some(earlier_slot))
+        .get_state(&earlier_block.state_root(), Some(earlier_slot), true)
         .expect("should not error getting state")
         .expect("should find state");
 
@@ -1331,7 +1331,7 @@ async fn attestation_validator_receive_proposer_reward_and_withdrawals() {
     let current_slot = harness.get_current_slot();
     let mut state = harness
         .chain
-        .get_state(&earlier_block.state_root(), Some(earlier_slot))
+        .get_state(&earlier_block.state_root(), Some(earlier_slot), true)
         .expect("should not error getting state")
         .expect("should find state");
 
@@ -1399,7 +1399,7 @@ async fn attestation_to_finalized_block() {
 
     let mut state = harness
         .chain
-        .get_state(&earlier_block.state_root(), Some(earlier_slot))
+        .get_state(&earlier_block.state_root(), Some(earlier_slot), true)
         .expect("should not error getting state")
         .expect("should find state");
 

beacon_node/beacon_chain/tests/rewards.rs

@@ -117,7 +117,7 @@ async fn test_sync_committee_rewards() {
         .unwrap()
         .unwrap();
     let parent_state = chain
-        .get_state(&parent_block.state_root(), Some(parent_block.slot()))
+        .get_state(&parent_block.state_root(), Some(parent_block.slot()), true)
         .unwrap()
         .unwrap();
 

beacon_node/beacon_chain/tests/store_tests.rs

@@ -758,6 +758,7 @@ async fn delete_blocks_and_states() {
         .get_state(
             &faulty_head_block.state_root(),
             Some(faulty_head_block.slot()),
+            true,
         )
         .expect("no db error")
         .expect("faulty head state exists");
@@ -771,7 +772,10 @@ async fn delete_blocks_and_states() {
             break;
         }
         store.delete_state(&state_root, slot).unwrap();
-        assert_eq!(store.get_state(&state_root, Some(slot)).unwrap(), None);
+        assert_eq!(
+            store.get_state(&state_root, Some(slot), true).unwrap(),
+            None
+        );
     }
 
     // Double-deleting should also be OK (deleting non-existent things is fine)
@@ -1055,7 +1059,7 @@ fn get_state_for_block(harness: &TestHarness, block_root: Hash256) -> BeaconStat
         .unwrap();
     harness
         .chain
-        .get_state(&head_block.state_root(), Some(head_block.slot()))
+        .get_state(&head_block.state_root(), Some(head_block.slot()), true)
         .unwrap()
         .unwrap()
 }
@@ -1892,7 +1896,10 @@ fn check_all_states_exist<'a>(
     states: impl Iterator<Item = &'a BeaconStateHash>,
 ) {
     for &state_hash in states {
-        let state = harness.chain.get_state(&state_hash.into(), None).unwrap();
+        let state = harness
+            .chain
+            .get_state(&state_hash.into(), None, true)
+            .unwrap();
         assert!(
             state.is_some(),
             "expected state {:?} to be in DB",
@@ -1910,7 +1917,7 @@ fn check_no_states_exist<'a>(
         assert!(
             harness
                 .chain
-                .get_state(&state_root.into(), None)
+                .get_state(&state_root.into(), None, true)
                 .unwrap()
                 .is_none(),
             "state {:?} should not be in the DB",
@@ -2344,7 +2351,7 @@ async fn weak_subjectivity_sync_test(slots: Vec<Slot>, checkpoint_slot: Slot) {
         .get_or_reconstruct_blobs(&wss_block_root)
         .unwrap();
     let wss_state = full_store
-        .get_state(&wss_state_root, Some(checkpoint_slot))
+        .get_state(&wss_state_root, Some(checkpoint_slot), true)
         .unwrap()
         .unwrap();
 
@@ -2460,7 +2467,7 @@ async fn weak_subjectivity_sync_test(slots: Vec<Slot>, checkpoint_slot: Slot) {
         // Check that the new block's state can be loaded correctly.
         let mut state = beacon_chain
             .store
-            .get_state(&state_root, Some(slot))
+            .get_state(&state_root, Some(slot), true)
             .unwrap()
             .unwrap();
         assert_eq!(state.update_tree_hash_cache().unwrap(), state_root);
@@ -2594,7 +2601,10 @@ async fn weak_subjectivity_sync_test(slots: Vec<Slot>, checkpoint_slot: Slot) {
         .unwrap()
         .map(Result::unwrap)
     {
-        let mut state = store.get_state(&state_root, Some(slot)).unwrap().unwrap();
+        let mut state = store
+            .get_state(&state_root, Some(slot), true)
+            .unwrap()
+            .unwrap();
         assert_eq!(state.slot(), slot);
         assert_eq!(state.canonical_root().unwrap(), state_root);
     }
@@ -3426,7 +3436,7 @@ async fn prune_historic_states() {
     let genesis_state_root = harness.chain.genesis_state_root;
     let genesis_state = harness
         .chain
-        .get_state(&genesis_state_root, None)
+        .get_state(&genesis_state_root, None, true)
         .unwrap()
         .unwrap();
 
@@ -3447,7 +3457,10 @@ async fn prune_historic_states() {
         .map(Result::unwrap)
         .collect::<Vec<_>>();
     for &(state_root, slot) in &first_epoch_state_roots {
-        assert!(store.get_state(&state_root, Some(slot)).unwrap().is_some());
+        assert!(store
+            .get_state(&state_root, Some(slot), true)
+            .unwrap()
+            .is_some());
     }
 
     store
@@ -3462,7 +3475,10 @@ async fn prune_historic_states() {
     // Ensure all epoch 0 states other than the genesis have been pruned.
     for &(state_root, slot) in &first_epoch_state_roots {
         assert_eq!(
-            store.get_state(&state_root, Some(slot)).unwrap().is_some(),
+            store
+                .get_state(&state_root, Some(slot), true)
+                .unwrap()
+                .is_some(),
             slot == 0
         );
     }
@@ -3588,7 +3604,7 @@ fn check_chain_dump(harness: &TestHarness, expected_len: u64) {
             harness
                 .chain
                 .store
-                .get_state(&checkpoint.beacon_state_root(), None)
+                .get_state(&checkpoint.beacon_state_root(), None, true)
                 .expect("no error")
                 .expect("state exists")
                 .slot(),
@@ -3650,7 +3666,7 @@ fn check_iterators(harness: &TestHarness) {
             harness
                 .chain
                 .store
-                .get_state(&state_root, Some(slot))
+                .get_state(&state_root, Some(slot), true)
                 .unwrap()
                 .is_some(),
             "state {:?} from canonical chain should be in DB",

beacon_node/beacon_chain/tests/sync_committee_verification.rs

@@ -755,7 +755,10 @@ async fn unaggregated_gossip_verification() {
 
         // Load the block and state for the given root.
         let block = chain.get_block(&root).await.unwrap().unwrap();
-        let mut state = chain.get_state(&block.state_root(), None).unwrap().unwrap();
+        let mut state = chain
+            .get_state(&block.state_root(), None, true)
+            .unwrap()
+            .unwrap();
 
         // Advance the state to simulate a pre-state for block production.
         let slot = valid_sync_committee_message.slot + 1;