Add configurable block replayer (#2863)

## Issue Addressed

Successor to #2431

## Proposed Changes

* Add a `BlockReplayer` struct to abstract over the intricacies of calling `per_slot_processing` and `per_block_processing` while avoiding unnecessary tree hashing.
* Add a variant of the forwards state root iterator that does not require an `end_state`.
* Use the `BlockReplayer` when reconstructing states in the database. Use the efficient forwards iterator for frozen states.
* Refactor the iterators to remove `Arc<HotColdDB>` (this seems to be neater than making _everything_ an `Arc<HotColdDB>` as I did in #2431).

Supplying the state roots allow us to avoid building a tree hash cache at all when reconstructing historic states, which saves around 1 second flat (regardless of `slots-per-restore-point`). This is a small percentage of worst-case state load times with 200K validators and SPRP=2048 (~15s vs ~16s) but a significant speed-up for more frequent restore points: state loads with SPRP=32 should be now consistently <500ms instead of 1.5s (a ~3x speedup).

## Additional Info

Required by https://github.com/sigp/lighthouse/pull/2628

consensus/state_processing/src/block_replayer.rs

@@ -0,0 +1,313 @@
+use crate::{
+    per_block_processing, per_epoch_processing::EpochProcessingSummary, per_slot_processing,
+    BlockProcessingError, BlockSignatureStrategy, SlotProcessingError, VerifyBlockRoot,
+};
+use std::marker::PhantomData;
+use types::{BeaconState, ChainSpec, EthSpec, Hash256, SignedBeaconBlock, Slot};
+
+type PreBlockHook<'a, E, Error> =
+    Box<dyn FnMut(&mut BeaconState<E>, &SignedBeaconBlock<E>) -> Result<(), Error> + 'a>;
+type PostBlockHook<'a, E, Error> = PreBlockHook<'a, E, Error>;
+type PreSlotHook<'a, E, Error> = Box<dyn FnMut(&mut BeaconState<E>) -> Result<(), Error> + 'a>;
+type PostSlotHook<'a, E, Error> = Box<
+    dyn FnMut(&mut BeaconState<E>, Option<EpochProcessingSummary<E>>, bool) -> Result<(), Error>
+        + 'a,
+>;
+type StateRootIterDefault<Error> = std::iter::Empty<Result<(Hash256, Slot), Error>>;
+
+/// Efficiently apply blocks to a state while configuring various parameters.
+///
+/// Usage follows a builder pattern.
+pub struct BlockReplayer<
+    'a,
+    Spec: EthSpec,
+    Error = BlockReplayError,
+    StateRootIter = StateRootIterDefault<Error>,
+> {
+    state: BeaconState<Spec>,
+    spec: &'a ChainSpec,
+    state_root_strategy: StateRootStrategy,
+    block_sig_strategy: BlockSignatureStrategy,
+    verify_block_root: Option<VerifyBlockRoot>,
+    pre_block_hook: Option<PreBlockHook<'a, Spec, Error>>,
+    post_block_hook: Option<PostBlockHook<'a, Spec, Error>>,
+    pre_slot_hook: Option<PreSlotHook<'a, Spec, Error>>,
+    post_slot_hook: Option<PostSlotHook<'a, Spec, Error>>,
+    state_root_iter: Option<StateRootIter>,
+    state_root_miss: bool,
+    _phantom: PhantomData<Error>,
+}
+
+#[derive(Debug)]
+pub enum BlockReplayError {
+    NoBlocks,
+    SlotProcessing(SlotProcessingError),
+    BlockProcessing(BlockProcessingError),
+}
+
+impl From<SlotProcessingError> for BlockReplayError {
+    fn from(e: SlotProcessingError) -> Self {
+        Self::SlotProcessing(e)
+    }
+}
+
+impl From<BlockProcessingError> for BlockReplayError {
+    fn from(e: BlockProcessingError) -> Self {
+        Self::BlockProcessing(e)
+    }
+}
+
+/// Defines how state roots should be computed during block replay.
+#[derive(PartialEq)]
+pub enum StateRootStrategy {
+    /// Perform all transitions faithfully to the specification.
+    Accurate,
+    /// Don't compute state roots, eventually computing an invalid beacon state that can only be
+    /// used for obtaining shuffling.
+    Inconsistent,
+}
+
+impl<'a, E, Error, StateRootIter> BlockReplayer<'a, E, Error, StateRootIter>
+where
+    E: EthSpec,
+    StateRootIter: Iterator<Item = Result<(Hash256, Slot), Error>>,
+    Error: From<BlockReplayError>,
+{
+    /// Create a new replayer that will apply blocks upon `state`.
+    ///
+    /// Defaults:
+    ///
+    /// - Full (bulk) signature verification
+    /// - Accurate state roots
+    /// - Full block root verification
+    pub fn new(state: BeaconState<E>, spec: &'a ChainSpec) -> Self {
+        Self {
+            state,
+            spec,
+            state_root_strategy: StateRootStrategy::Accurate,
+            block_sig_strategy: BlockSignatureStrategy::VerifyBulk,
+            verify_block_root: Some(VerifyBlockRoot::True),
+            pre_block_hook: None,
+            post_block_hook: None,
+            pre_slot_hook: None,
+            post_slot_hook: None,
+            state_root_iter: None,
+            state_root_miss: false,
+            _phantom: PhantomData,
+        }
+    }
+
+    /// Set the replayer's state root strategy different from the default.
+    pub fn state_root_strategy(mut self, state_root_strategy: StateRootStrategy) -> Self {
+        if state_root_strategy == StateRootStrategy::Inconsistent {
+            self.verify_block_root = None;
+        }
+        self.state_root_strategy = state_root_strategy;
+        self
+    }
+
+    /// Set the replayer's block signature verification strategy.
+    pub fn block_signature_strategy(mut self, block_sig_strategy: BlockSignatureStrategy) -> Self {
+        self.block_sig_strategy = block_sig_strategy;
+        self
+    }
+
+    /// Disable signature verification during replay.
+    ///
+    /// If you are truly _replaying_ blocks then you will almost certainly want to disable
+    /// signature checks for performance.
+    pub fn no_signature_verification(self) -> Self {
+        self.block_signature_strategy(BlockSignatureStrategy::NoVerification)
+    }
+
+    /// Verify only the block roots of the initial few blocks, and trust the rest.
+    pub fn minimal_block_root_verification(mut self) -> Self {
+        self.verify_block_root = None;
+        self
+    }
+
+    /// Supply a state root iterator to accelerate slot processing.
+    ///
+    /// If possible the state root iterator should return a state root for every slot from
+    /// `self.state.slot` to the `target_slot` supplied to `apply_blocks` (inclusive of both
+    /// endpoints).
+    pub fn state_root_iter(mut self, iter: StateRootIter) -> Self {
+        self.state_root_iter = Some(iter);
+        self
+    }
+
+    /// Run a function immediately before each block that is applied during `apply_blocks`.
+    ///
+    /// This can be used to inspect the state as blocks are applied.
+    pub fn pre_block_hook(mut self, hook: PreBlockHook<'a, E, Error>) -> Self {
+        self.pre_block_hook = Some(hook);
+        self
+    }
+
+    /// Run a function immediately after each block that is applied during `apply_blocks`.
+    ///
+    /// This can be used to inspect the state as blocks are applied.
+    pub fn post_block_hook(mut self, hook: PostBlockHook<'a, E, Error>) -> Self {
+        self.post_block_hook = Some(hook);
+        self
+    }
+
+    /// Run a function immediately before slot processing advances the state to the next slot.
+    pub fn pre_slot_hook(mut self, hook: PreSlotHook<'a, E, Error>) -> Self {
+        self.pre_slot_hook = Some(hook);
+        self
+    }
+
+    /// Run a function immediately after slot processing has advanced the state to the next slot.
+    ///
+    /// The hook receives the state and a bool indicating if this state corresponds to a skipped
+    /// slot (i.e. it will not have a block applied).
+    pub fn post_slot_hook(mut self, hook: PostSlotHook<'a, E, Error>) -> Self {
+        self.post_slot_hook = Some(hook);
+        self
+    }
+
+    /// Compute the state root for `slot` as efficiently as possible.
+    ///
+    /// The `blocks` should be the full list of blocks being applied and `i` should be the index of
+    /// the next block that will be applied, or `blocks.len()` if all blocks have already been
+    /// applied.
+    fn get_state_root(
+        &mut self,
+        slot: Slot,
+        blocks: &[SignedBeaconBlock<E>],
+        i: usize,
+    ) -> Result<Option<Hash256>, Error> {
+        // If we don't care about state roots then return immediately.
+        if self.state_root_strategy == StateRootStrategy::Inconsistent {
+            return Ok(Some(Hash256::zero()));
+        }
+
+        // If a state root iterator is configured, use it to find the root.
+        if let Some(ref mut state_root_iter) = self.state_root_iter {
+            let opt_root = state_root_iter
+                .take_while(|res| res.as_ref().map_or(true, |(_, s)| *s <= slot))
+                .find(|res| res.as_ref().map_or(true, |(_, s)| *s == slot))
+                .transpose()?;
+
+            if let Some((root, _)) = opt_root {
+                return Ok(Some(root));
+            }
+        }
+
+        // Otherwise try to source a root from the previous block.
+        if let Some(prev_i) = i.checked_sub(1) {
+            if let Some(prev_block) = blocks.get(prev_i) {
+                if prev_block.slot() == slot {
+                    return Ok(Some(prev_block.state_root()));
+                }
+            }
+        }
+
+        self.state_root_miss = true;
+        Ok(None)
+    }
+
+    /// Apply `blocks` atop `self.state`, taking care of slot processing.
+    ///
+    /// If `target_slot` is provided then the state will be advanced through to `target_slot`
+    /// after the blocks have been applied.
+    pub fn apply_blocks(
+        mut self,
+        blocks: Vec<SignedBeaconBlock<E>>,
+        target_slot: Option<Slot>,
+    ) -> Result<Self, Error> {
+        for (i, block) in blocks.iter().enumerate() {
+            // Allow one additional block at the start which is only used for its state root.
+            if i == 0 && block.slot() <= self.state.slot() {
+                continue;
+            }
+
+            while self.state.slot() < block.slot() {
+                if let Some(ref mut pre_slot_hook) = self.pre_slot_hook {
+                    pre_slot_hook(&mut self.state)?;
+                }
+
+                let state_root = self.get_state_root(self.state.slot(), &blocks, i)?;
+                let summary = per_slot_processing(&mut self.state, state_root, self.spec)
+                    .map_err(BlockReplayError::from)?;
+
+                if let Some(ref mut post_slot_hook) = self.post_slot_hook {
+                    let is_skipped_slot = self.state.slot() < block.slot();
+                    post_slot_hook(&mut self.state, summary, is_skipped_slot)?;
+                }
+            }
+
+            if let Some(ref mut pre_block_hook) = self.pre_block_hook {
+                pre_block_hook(&mut self.state, block)?;
+            }
+
+            let verify_block_root = self.verify_block_root.unwrap_or_else(|| {
+                // If no explicit policy is set, verify only the first 1 or 2 block roots if using
+                // accurate state roots. Inaccurate state roots require block root verification to
+                // be off.
+                if i <= 1 && self.state_root_strategy == StateRootStrategy::Accurate {
+                    VerifyBlockRoot::True
+                } else {
+                    VerifyBlockRoot::False
+                }
+            });
+            per_block_processing(
+                &mut self.state,
+                block,
+                None,
+                self.block_sig_strategy,
+                verify_block_root,
+                self.spec,
+            )
+            .map_err(BlockReplayError::from)?;
+
+            if let Some(ref mut post_block_hook) = self.post_block_hook {
+                post_block_hook(&mut self.state, block)?;
+            }
+        }
+
+        if let Some(target_slot) = target_slot {
+            while self.state.slot() < target_slot {
+                if let Some(ref mut pre_slot_hook) = self.pre_slot_hook {
+                    pre_slot_hook(&mut self.state)?;
+                }
+
+                let state_root = self.get_state_root(self.state.slot(), &blocks, blocks.len())?;
+                let summary = per_slot_processing(&mut self.state, state_root, self.spec)
+                    .map_err(BlockReplayError::from)?;
+
+                if let Some(ref mut post_slot_hook) = self.post_slot_hook {
+                    // No more blocks to apply (from our perspective) so we consider these slots
+                    // skipped.
+                    let is_skipped_slot = true;
+                    post_slot_hook(&mut self.state, summary, is_skipped_slot)?;
+                }
+            }
+        }
+
+        Ok(self)
+    }
+
+    /// After block application, check if a state root miss occurred.
+    pub fn state_root_miss(&self) -> bool {
+        self.state_root_miss
+    }
+
+    /// Convert the replayer into the state that was built.
+    pub fn into_state(self) -> BeaconState<E> {
+        self.state
+    }
+}
+
+impl<'a, E, Error> BlockReplayer<'a, E, Error, StateRootIterDefault<Error>>
+where
+    E: EthSpec,
+    Error: From<BlockReplayError>,
+{
+    /// If type inference fails to infer the state root iterator type you can use this method
+    /// to hint that no state root iterator is desired.
+    pub fn no_state_root_iter(self) -> Self {
+        self
+    }
+}

consensus/state_processing/src/lib.rs

@@ -16,6 +16,7 @@
 mod macros;
 mod metrics;
 
+pub mod block_replayer;
 pub mod common;
 pub mod genesis;
 pub mod per_block_processing;
@@ -25,13 +26,14 @@ pub mod state_advance;
 pub mod upgrade;
 pub mod verify_operation;
 
+pub use block_replayer::{BlockReplayError, BlockReplayer, StateRootStrategy};
 pub use genesis::{
     eth2_genesis_time, initialize_beacon_state_from_eth1, is_valid_genesis_state,
     process_activations,
 };
 pub use per_block_processing::{
     block_signature_verifier, errors::BlockProcessingError, per_block_processing, signature_sets,
-    BlockSignatureStrategy, BlockSignatureVerifier, VerifySignatures,
+    BlockSignatureStrategy, BlockSignatureVerifier, VerifyBlockRoot, VerifySignatures,
 };
 pub use per_epoch_processing::{
     errors::EpochProcessingError, process_epoch as per_epoch_processing,

consensus/state_processing/src/per_block_processing.rs

@@ -68,6 +68,14 @@ impl VerifySignatures {
     }
 }
 
+/// Control verification of the latest block header.
+#[cfg_attr(feature = "arbitrary-fuzz", derive(Arbitrary))]
+#[derive(PartialEq, Clone, Copy)]
+pub enum VerifyBlockRoot {
+    True,
+    False,
+}
+
 /// Updates the state for a new block, whilst validating that the block is valid, optionally
 /// checking the block proposer signature.
 ///
@@ -84,6 +92,7 @@ pub fn per_block_processing<T: EthSpec>(
     signed_block: &SignedBeaconBlock<T>,
     block_root: Option<Hash256>,
     block_signature_strategy: BlockSignatureStrategy,
+    verify_block_root: VerifyBlockRoot,
     spec: &ChainSpec,
 ) -> Result<(), BlockProcessingError> {
     let block = signed_block.message();
@@ -120,7 +129,7 @@ pub fn per_block_processing<T: EthSpec>(
         BlockSignatureStrategy::VerifyRandao => VerifySignatures::False,
     };
 
-    let proposer_index = process_block_header(state, block, spec)?;
+    let proposer_index = process_block_header(state, block, verify_block_root, spec)?;
 
     if verify_signatures.is_true() {
         verify_block_signature(state, signed_block, block_root, spec)?;
@@ -167,6 +176,7 @@ pub fn per_block_processing<T: EthSpec>(
 pub fn process_block_header<T: EthSpec>(
     state: &mut BeaconState<T>,
     block: BeaconBlockRef<'_, T>,
+    verify_block_root: VerifyBlockRoot,
     spec: &ChainSpec,
 ) -> Result<u64, BlockOperationError<HeaderInvalid>> {
     // Verify that the slots match
@@ -195,14 +205,16 @@ pub fn process_block_header<T: EthSpec>(
         }
     );
 
-    let expected_previous_block_root = state.latest_block_header().tree_hash_root();
-    verify!(
-        block.parent_root() == expected_previous_block_root,
-        HeaderInvalid::ParentBlockRootMismatch {
-            state: expected_previous_block_root,
-            block: block.parent_root(),
-        }
-    );
+    if verify_block_root == VerifyBlockRoot::True {
+        let expected_previous_block_root = state.latest_block_header().tree_hash_root();
+        verify!(
+            block.parent_root() == expected_previous_block_root,
+            HeaderInvalid::ParentBlockRootMismatch {
+                state: expected_previous_block_root,
+                block: block.parent_root(),
+            }
+        );
+    }
 
     *state.latest_block_header_mut() = block.temporary_block_header();
 

consensus/state_processing/src/per_block_processing/tests.rs

@@ -6,7 +6,10 @@ use crate::per_block_processing::errors::{
     DepositInvalid, HeaderInvalid, IndexedAttestationInvalid, IntoWithIndex,
     ProposerSlashingInvalid,
 };
-use crate::{per_block_processing::process_operations, BlockSignatureStrategy, VerifySignatures};
+use crate::{
+    per_block_processing::process_operations, BlockSignatureStrategy, VerifyBlockRoot,
+    VerifySignatures,
+};
 use beacon_chain::test_utils::{BeaconChainHarness, EphemeralHarnessType};
 use lazy_static::lazy_static;
 use ssz_types::Bitfield;
@@ -65,6 +68,7 @@ fn valid_block_ok() {
         &block,
         None,
         BlockSignatureStrategy::VerifyIndividual,
+        VerifyBlockRoot::True,
         &spec,
     );
 
@@ -88,6 +92,7 @@ fn invalid_block_header_state_slot() {
         &SignedBeaconBlock::from_block(block, signature),
         None,
         BlockSignatureStrategy::VerifyIndividual,
+        VerifyBlockRoot::True,
         &spec,
     );
 
@@ -116,6 +121,7 @@ fn invalid_parent_block_root() {
         &SignedBeaconBlock::from_block(block, signature),
         None,
         BlockSignatureStrategy::VerifyIndividual,
+        VerifyBlockRoot::True,
         &spec,
     );
 
@@ -145,6 +151,7 @@ fn invalid_block_signature() {
         &SignedBeaconBlock::from_block(block, Signature::empty()),
         None,
         BlockSignatureStrategy::VerifyIndividual,
+        VerifyBlockRoot::True,
         &spec,
     );
 
@@ -174,6 +181,7 @@ fn invalid_randao_reveal_signature() {
         &signed_block,
         None,
         BlockSignatureStrategy::VerifyIndividual,
+        VerifyBlockRoot::True,
         &spec,
     );