Clean up temporary state flags while running (#6422)

* Clean up temporary state flags while running

* Add regression test

* Simplify

beacon_node/beacon_chain/tests/store_tests.rs

@@ -2514,7 +2514,7 @@ async fn pruning_test(
 }
 
 #[tokio::test]
-async fn garbage_collect_temp_states_from_failed_block() {
+async fn garbage_collect_temp_states_from_failed_block_on_startup() {
     let db_path = tempdir().unwrap();
 
     // Wrap these functions to ensure the variables are dropped before we try to open another
@@ -2571,6 +2571,61 @@ async fn garbage_collect_temp_states_from_failed_block() {
     assert_eq!(store.iter_temporary_state_roots().count(), 0);
 }
 
+#[tokio::test]
+async fn garbage_collect_temp_states_from_failed_block_on_finalization() {
+    let db_path = tempdir().unwrap();
+
+    let store = get_store(&db_path);
+    let harness = get_harness(store.clone(), LOW_VALIDATOR_COUNT);
+
+    let slots_per_epoch = E::slots_per_epoch();
+
+    let genesis_state = harness.get_current_state();
+    let block_slot = Slot::new(2 * slots_per_epoch);
+    let ((signed_block, _), state) = harness.make_block(genesis_state, block_slot).await;
+
+    let (mut block, _) = (*signed_block).clone().deconstruct();
+
+    // Mutate the block to make it invalid, and re-sign it.
+    *block.state_root_mut() = Hash256::repeat_byte(0xff);
+    let proposer_index = block.proposer_index() as usize;
+    let block = Arc::new(block.sign(
+        &harness.validator_keypairs[proposer_index].sk,
+        &state.fork(),
+        state.genesis_validators_root(),
+        &harness.spec,
+    ));
+
+    // The block should be rejected, but should store a bunch of temporary states.
+    harness.set_current_slot(block_slot);
+    harness
+        .process_block_result((block, None))
+        .await
+        .unwrap_err();
+
+    assert_eq!(
+        store.iter_temporary_state_roots().count(),
+        block_slot.as_usize() - 1
+    );
+
+    // Finalize the chain without the block, which should result in pruning of all temporary states.
+    let blocks_required_to_finalize = 3 * slots_per_epoch;
+    harness.advance_slot();
+    harness
+        .extend_chain(
+            blocks_required_to_finalize as usize,
+            BlockStrategy::OnCanonicalHead,
+            AttestationStrategy::AllValidators,
+        )
+        .await;
+
+    // Check that the finalization migration ran.
+    assert_ne!(store.get_split_slot(), 0);
+
+    // Check that temporary states have been pruned.
+    assert_eq!(store.iter_temporary_state_roots().count(), 0);
+}
+
 #[tokio::test]
 async fn weak_subjectivity_sync_easy() {
     let num_initial_slots = E::slots_per_epoch() * 11;

beacon_node/store/src/garbage_collection.rs

@@ -21,7 +21,6 @@ where
                 .try_fold(vec![], |mut ops, state_root| {
                     let state_root = state_root?;
                     ops.push(StoreOp::DeleteState(state_root, None));
-                    ops.push(StoreOp::DeleteStateTemporaryFlag(state_root));
                     Result::<_, Error>::Ok(ops)
                 })?;
 
@@ -29,7 +28,7 @@ where
             debug!(
                 self.log,
                 "Garbage collecting {} temporary states",
-                delete_ops.len() / 2
+                delete_ops.len()
             );
             self.do_atomically_with_block_and_blobs_cache(delete_ops)?;
         }

beacon_node/store/src/hot_cold_store.rs

@@ -1160,10 +1160,19 @@ impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> HotColdDB<E, Hot, Cold>
                 }
 
                 StoreOp::DeleteState(state_root, slot) => {
+                    // Delete the hot state summary.
                     let state_summary_key =
                         get_key_for_col(DBColumn::BeaconStateSummary.into(), state_root.as_slice());
                     key_value_batch.push(KeyValueStoreOp::DeleteKey(state_summary_key));
 
+                    // Delete the state temporary flag (if any). Temporary flags are commonly
+                    // created by the state advance routine.
+                    let state_temp_key = get_key_for_col(
+                        DBColumn::BeaconStateTemporary.into(),
+                        state_root.as_slice(),
+                    );
+                    key_value_batch.push(KeyValueStoreOp::DeleteKey(state_temp_key));
+
                     if slot.map_or(true, |slot| slot % E::slots_per_epoch() == 0) {
                         let state_key =
                             get_key_for_col(DBColumn::BeaconState.into(), state_root.as_slice());