Use state roots iter directly

beacon_node/beacon_chain/src/canonical_head.rs

@@ -8,12 +8,14 @@ use crate::{
 };
 use eth2::types::{EventKind, SseChainReorg, SseFinalizedCheckpoint, SseHead, SseLateHead};
 use fork_choice::{ExecutionStatus, ForkChoiceView, ForkchoiceUpdateParameters, ProtoBlock};
+use itertools::process_results;
 use parking_lot::RwLockWriteGuard;
 use slog::{crit, debug, error, warn, Logger};
 use slot_clock::SlotClock;
 use std::mem;
 use std::sync::Arc;
 use std::time::Duration;
+use store::iter::StateRootsIterator;
 use task_executor::{JoinHandle, ShutdownReason};
 use types::*;
 
@@ -408,7 +410,36 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
 
         // If the finalized checkpoint changed, perform some updates.
         if new_view.finalized_checkpoint != old_view.finalized_checkpoint {
-            if let Err(e) = self.after_finalization(new_head, new_view, finalized_proto_block) {
+            // The store migration task requires the *state at the slot of the finalized epoch*,
+            // rather than the state of the latest finalized block. These two values will only
+            // differ when the first slot of the finalized epoch is a skip slot.
+            //
+            // Use the `StateRootsIterator` directly rather than `BeaconChain::state_root_at_slot`
+            // since the latter will try to take a read-lock on the state.
+            let new_finalized_slot = new_view
+                .finalized_checkpoint
+                .epoch
+                .start_slot(T::EthSpec::slots_per_epoch());
+            let new_finalized_state_root = process_results(
+                StateRootsIterator::new(&self.store, &new_head.beacon_state),
+                |mut iter| {
+                    iter.find_map(|(state_root, slot)| {
+                        if slot == new_finalized_slot {
+                            Some(state_root)
+                        } else {
+                            None
+                        }
+                    })
+                },
+            )?
+            .ok_or(Error::MissingFinalizedStateRoot(new_finalized_slot))?;
+
+            if let Err(e) = self.after_finalization(
+                new_head,
+                new_view,
+                finalized_proto_block,
+                new_finalized_state_root,
+            ) {
                 crit!(
                     self.log,
                     "Error updating finalization";
@@ -597,6 +628,7 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
         new_head: &BeaconSnapshot<T::EthSpec>,
         new_view: ForkChoiceView,
         finalized_proto_block: ProtoBlock,
+        new_finalized_state_root: Hash256,
     ) -> Result<(), Error> {
         self.op_pool
             .prune_all(&new_head.beacon_state, self.epoch()?);
@@ -631,20 +663,6 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
                     );
                 });
 
-            // The store migration task requires the *state at the slot of the finalized epoch*,
-            // rather than the state of the latest finalized block. These two values will only
-            // differ when the first slot of the finalized epoch is a skip slot.
-            //
-            // The `BeaconChain::state_root_at_slot` might take a read-lock on the canonical head.
-            // Another read-lock on the canonical head is being held whilst the function is called,
-            // but that's OK since read-locks alone don't deadlock.
-            let new_finalized_slot = new_view
-                .finalized_checkpoint
-                .epoch
-                .start_slot(T::EthSpec::slots_per_epoch());
-            let new_finalized_state_root = chain
-                .state_root_at_slot(new_finalized_slot)?
-                .ok_or(Error::MissingFinalizedStateRoot(new_finalized_slot))?;
             chain.store_migrator.process_finalization(
                 new_finalized_state_root.into(),
                 new_view.finalized_checkpoint,