From b82f7843ffb173826e42ecd46c394b101c8133b4 Mon Sep 17 00:00:00 2001 From: Michael Sproul Date: Thu, 5 Oct 2023 06:03:24 +0000 Subject: [PATCH] Use peeking_take_while in BlockReplayer (#4803) ## Issue Addressed While reviewing #4801 I noticed that our use of `take_while` in the block replayer means that if a state root iterator _with gaps_ is provided, some additonal state roots will be dropped unnecessarily. In practice the impact is small, because once there's _one_ state root miss, the whole tree hash cache needs to be built anyway, and subsequent misses are less costly. However this was still a little inefficient, so I figured it's better to fix it. ## Proposed Changes Use [`peeking_take_while`](https://docs.rs/itertools/latest/itertools/trait.Itertools.html#method.peeking_take_while) to avoid consuming the next element when checking whether it satisfies the slot predicate. ## Additional Info There's a gist here that shows the basic dynamics in isolation: https://play.rust-lang.org/?version=stable&mode=debug&edition=2021&gist=40b623cc0febf9ed51705d476ab140c5. Changing the `peeking_take_while` to a `take_while` causes the assert to fail. Similarly I've added a new test `block_replayer_peeking_state_roots` which fails if the same change is applied inside `get_state_root`. --- .../state_processing/src/block_replayer.rs | 10 ++-- .../src/per_block_processing/tests.rs | 50 ++++++++++++++++++- 2 files changed, 55 insertions(+), 5 deletions(-) diff --git a/consensus/state_processing/src/block_replayer.rs b/consensus/state_processing/src/block_replayer.rs index ed5e64294..f502d7f69 100644 --- a/consensus/state_processing/src/block_replayer.rs +++ b/consensus/state_processing/src/block_replayer.rs @@ -3,6 +3,8 @@ use crate::{ BlockProcessingError, BlockSignatureStrategy, ConsensusContext, SlotProcessingError, VerifyBlockRoot, }; +use itertools::Itertools; +use std::iter::Peekable; use std::marker::PhantomData; use types::{BeaconState, BlindedPayload, ChainSpec, EthSpec, Hash256, SignedBeaconBlock, Slot}; @@ -25,7 +27,7 @@ pub struct BlockReplayer< 'a, Spec: EthSpec, Error = BlockReplayError, - StateRootIter = StateRootIterDefault, + StateRootIter: Iterator> = StateRootIterDefault, > { state: BeaconState, spec: &'a ChainSpec, @@ -36,7 +38,7 @@ pub struct BlockReplayer< post_block_hook: Option>, pre_slot_hook: Option>, post_slot_hook: Option>, - state_root_iter: Option, + pub(crate) state_root_iter: Option>, state_root_miss: bool, _phantom: PhantomData, } @@ -138,7 +140,7 @@ where /// `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.state_root_iter = Some(iter.peekable()); self } @@ -192,7 +194,7 @@ where // 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)) + .peeking_take_while(|res| res.as_ref().map_or(true, |(_, s)| *s <= slot)) .find(|res| res.as_ref().map_or(true, |(_, s)| *s == slot)) .transpose()?; diff --git a/consensus/state_processing/src/per_block_processing/tests.rs b/consensus/state_processing/src/per_block_processing/tests.rs index df5aa9f7a..8f7fb43b1 100644 --- a/consensus/state_processing/src/per_block_processing/tests.rs +++ b/consensus/state_processing/src/per_block_processing/tests.rs @@ -5,7 +5,7 @@ use crate::per_block_processing::errors::{ DepositInvalid, HeaderInvalid, IndexedAttestationInvalid, IntoWithIndex, ProposerSlashingInvalid, }; -use crate::{per_block_processing, StateProcessingStrategy}; +use crate::{per_block_processing, BlockReplayError, BlockReplayer, StateProcessingStrategy}; use crate::{ per_block_processing::{process_operations, verify_exit::verify_exit}, BlockSignatureStrategy, ConsensusContext, VerifyBlockRoot, VerifySignatures, @@ -1035,3 +1035,51 @@ async fn fork_spanning_exit() { ) .expect_err("phase0 exit does not verify against bellatrix state"); } + +/// Check that the block replayer does not consume state roots unnecessarily. +#[tokio::test] +async fn block_replayer_peeking_state_roots() { + let harness = get_harness::(EPOCH_OFFSET, VALIDATOR_COUNT).await; + + let target_state = harness.get_current_state(); + let target_block_root = harness.head_block_root(); + let target_block = harness + .chain + .get_blinded_block(&target_block_root) + .unwrap() + .unwrap(); + + let parent_block_root = target_block.parent_root(); + let parent_block = harness + .chain + .get_blinded_block(&parent_block_root) + .unwrap() + .unwrap(); + let parent_state = harness + .chain + .get_state(&parent_block.state_root(), Some(parent_block.slot())) + .unwrap() + .unwrap(); + + // Omit the state root for `target_state` but provide a dummy state root at the *next* slot. + // If the block replayer is peeking at the state roots rather than consuming them, then the + // dummy state should still be there after block replay completes. + let dummy_state_root = Hash256::repeat_byte(0xff); + let dummy_slot = target_state.slot() + 1; + let state_root_iter = vec![Ok::<_, BlockReplayError>((dummy_state_root, dummy_slot))]; + let block_replayer = BlockReplayer::new(parent_state, &harness.chain.spec) + .state_root_iter(state_root_iter.into_iter()) + .no_signature_verification() + .apply_blocks(vec![target_block], None) + .unwrap(); + + assert_eq!( + block_replayer + .state_root_iter + .unwrap() + .next() + .unwrap() + .unwrap(), + (dummy_state_root, dummy_slot) + ); +}