Use head state for exit verification (#4183)

## Issue Addressed

NA

## Proposed Changes

Similar to #4181 but without the version bump and a more nuanced fix.

Patches the high CPU usage seen after the Capella fork which was caused by processing exits when there are skip slots.

## Additional Info

~~This is an imperfect solution that will cause us to drop some exits at the fork boundary. This is tracked at #4184.~~

beacon_node/beacon_chain/src/beacon_chain.rs

@@ -2206,12 +2206,14 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
         &self,
         exit: SignedVoluntaryExit,
     ) -> Result<ObservationOutcome<SignedVoluntaryExit, T::EthSpec>, Error> {
-        // NOTE: this could be more efficient if it avoided cloning the head state
+        let head_snapshot = self.head().snapshot;
-        let wall_clock_state = self.wall_clock_state()?;
+        let head_state = &head_snapshot.beacon_state;
+        let wall_clock_epoch = self.epoch()?;
+
         Ok(self
             .observed_voluntary_exits
             .lock()
-            .verify_and_observe(exit, &wall_clock_state, &self.spec)
+            .verify_and_observe_at(exit, wall_clock_epoch, head_state, &self.spec)
             .map(|exit| {
                 // this method is called for both API and gossip exits, so this covers all exit events
                 if let Some(event_handler) = self.event_handler.as_ref() {

beacon_node/beacon_chain/src/observed_operations.rs

@@ -1,11 +1,11 @@
 use derivative::Derivative;
 use smallvec::{smallvec, SmallVec};
 use ssz::{Decode, Encode};
-use state_processing::{SigVerifiedOp, VerifyOperation};
+use state_processing::{SigVerifiedOp, VerifyOperation, VerifyOperationAt};
 use std::collections::HashSet;
 use std::marker::PhantomData;
 use types::{
-    AttesterSlashing, BeaconState, ChainSpec, EthSpec, ForkName, ProposerSlashing,
+    AttesterSlashing, BeaconState, ChainSpec, Epoch, EthSpec, ForkName, ProposerSlashing,
     SignedBlsToExecutionChange, SignedVoluntaryExit, Slot,
 };
 
@@ -87,12 +87,16 @@ impl<E: EthSpec> ObservableOperation<E> for SignedBlsToExecutionChange {
 }
 
 impl<T: ObservableOperation<E>, E: EthSpec> ObservedOperations<T, E> {
-    pub fn verify_and_observe(
+    pub fn verify_and_observe_parametric<F>(
         &mut self,
         op: T,
+        validate: F,
         head_state: &BeaconState<E>,
         spec: &ChainSpec,
-    ) -> Result<ObservationOutcome<T, E>, T::Error> {
+    ) -> Result<ObservationOutcome<T, E>, T::Error>
+    where
+        F: Fn(T) -> Result<SigVerifiedOp<T, E>, T::Error>,
+    {
         self.reset_at_fork_boundary(head_state.slot(), spec);
 
         let observed_validator_indices = &mut self.observed_validator_indices;
@@ -112,7 +116,7 @@ impl<T: ObservableOperation<E>, E: EthSpec> ObservedOperations<T, E> {
         }
 
         // Validate the op using operation-specific logic (`verify_attester_slashing`, etc).
-        let verified_op = op.validate(head_state, spec)?;
+        let verified_op = validate(op)?;
 
         // Add the relevant indices to the set of known indices to prevent processing of duplicates
         // in the future.
@@ -121,6 +125,16 @@ impl<T: ObservableOperation<E>, E: EthSpec> ObservedOperations<T, E> {
         Ok(ObservationOutcome::New(verified_op))
     }
 
+    pub fn verify_and_observe(
+        &mut self,
+        op: T,
+        head_state: &BeaconState<E>,
+        spec: &ChainSpec,
+    ) -> Result<ObservationOutcome<T, E>, T::Error> {
+        let validate = |op: T| op.validate(head_state, spec);
+        self.verify_and_observe_parametric(op, validate, head_state, spec)
+    }
+
     /// Reset the cache when crossing a fork boundary.
     ///
     /// This prevents an attacker from crafting a self-slashing which is only valid before the fork
@@ -140,3 +154,16 @@ impl<T: ObservableOperation<E>, E: EthSpec> ObservedOperations<T, E> {
         }
     }
 }
+
+impl<T: ObservableOperation<E> + VerifyOperationAt<E>, E: EthSpec> ObservedOperations<T, E> {
+    pub fn verify_and_observe_at(
+        &mut self,
+        op: T,
+        verify_at_epoch: Epoch,
+        head_state: &BeaconState<E>,
+        spec: &ChainSpec,
+    ) -> Result<ObservationOutcome<T, E>, T::Error> {
+        let validate = |op: T| op.validate_at(head_state, verify_at_epoch, spec);
+        self.verify_and_observe_parametric(op, validate, head_state, spec)
+    }
+}

beacon_node/operation_pool/src/lib.rs

@@ -497,7 +497,8 @@ impl<T: EthSpec> OperationPool<T> {
             |exit| {
                 filter(exit.as_inner())
                     && exit.signature_is_still_valid(&state.fork())
-                    && verify_exit(state, exit.as_inner(), VerifySignatures::False, spec).is_ok()
+                    && verify_exit(state, None, exit.as_inner(), VerifySignatures::False, spec)
+                        .is_ok()
             },
             |exit| exit.as_inner().clone(),
             T::MaxVoluntaryExits::to_usize(),

consensus/state_processing/src/lib.rs

@@ -41,4 +41,4 @@ pub use per_epoch_processing::{
     errors::EpochProcessingError, process_epoch as per_epoch_processing,
 };
 pub use per_slot_processing::{per_slot_processing, Error as SlotProcessingError};
-pub use verify_operation::{SigVerifiedOp, VerifyOperation};
+pub use verify_operation::{SigVerifiedOp, VerifyOperation, VerifyOperationAt};

consensus/state_processing/src/per_block_processing/process_operations.rs

@@ -282,7 +282,8 @@ pub fn process_exits<T: EthSpec>(
     // Verify and apply each exit in series. We iterate in series because higher-index exits may
     // become invalid due to the application of lower-index ones.
     for (i, exit) in voluntary_exits.iter().enumerate() {
-        verify_exit(state, exit, verify_signatures, spec).map_err(|e| e.into_with_index(i))?;
+        verify_exit(state, None, exit, verify_signatures, spec)
+            .map_err(|e| e.into_with_index(i))?;
 
         initiate_validator_exit(state, exit.message.validator_index as usize, spec)?;
     }

consensus/state_processing/src/per_block_processing/tests.rs

@@ -978,8 +978,14 @@ async fn fork_spanning_exit() {
     let head = harness.chain.canonical_head.cached_head();
     let head_state = &head.snapshot.beacon_state;
     assert!(head_state.current_epoch() < spec.altair_fork_epoch.unwrap());
-    verify_exit(head_state, &signed_exit, VerifySignatures::True, &spec)
+    verify_exit(
-        .expect("phase0 exit verifies against phase0 state");
+        head_state,
+        None,
+        &signed_exit,
+        VerifySignatures::True,
+        &spec,
+    )
+    .expect("phase0 exit verifies against phase0 state");
 
     /*
      * Ensure the exit verifies after Altair.
@@ -992,8 +998,14 @@ async fn fork_spanning_exit() {
     let head_state = &head.snapshot.beacon_state;
     assert!(head_state.current_epoch() >= spec.altair_fork_epoch.unwrap());
     assert!(head_state.current_epoch() < spec.bellatrix_fork_epoch.unwrap());
-    verify_exit(head_state, &signed_exit, VerifySignatures::True, &spec)
+    verify_exit(
-        .expect("phase0 exit verifies against altair state");
+        head_state,
+        None,
+        &signed_exit,
+        VerifySignatures::True,
+        &spec,
+    )
+    .expect("phase0 exit verifies against altair state");
 
     /*
      * Ensure the exit no longer verifies after Bellatrix.
@@ -1009,6 +1021,12 @@ async fn fork_spanning_exit() {
     let head = harness.chain.canonical_head.cached_head();
     let head_state = &head.snapshot.beacon_state;
     assert!(head_state.current_epoch() >= spec.bellatrix_fork_epoch.unwrap());
-    verify_exit(head_state, &signed_exit, VerifySignatures::True, &spec)
+    verify_exit(
-        .expect_err("phase0 exit does not verify against bellatrix state");
+        head_state,
+        None,
+        &signed_exit,
+        VerifySignatures::True,
+        &spec,
+    )
+    .expect_err("phase0 exit does not verify against bellatrix state");
 }

consensus/state_processing/src/per_block_processing/verify_exit.rs

@@ -20,10 +20,12 @@ fn error(reason: ExitInvalid) -> BlockOperationError<ExitInvalid> {
 /// Spec v0.12.1
 pub fn verify_exit<T: EthSpec>(
     state: &BeaconState<T>,
+    current_epoch: Option<Epoch>,
     signed_exit: &SignedVoluntaryExit,
     verify_signatures: VerifySignatures,
     spec: &ChainSpec,
 ) -> Result<()> {
+    let current_epoch = current_epoch.unwrap_or(state.current_epoch());
     let exit = &signed_exit.message;
 
     let validator = state
@@ -33,7 +35,7 @@ pub fn verify_exit<T: EthSpec>(
 
     // Verify the validator is active.
     verify!(
-        validator.is_active_at(state.current_epoch()),
+        validator.is_active_at(current_epoch),
         ExitInvalid::NotActive(exit.validator_index)
     );
 
@@ -45,9 +47,9 @@ pub fn verify_exit<T: EthSpec>(
 
     // Exits must specify an epoch when they become valid; they are not valid before then.
     verify!(
-        state.current_epoch() >= exit.epoch,
+        current_epoch >= exit.epoch,
         ExitInvalid::FutureEpoch {
-            state: state.current_epoch(),
+            state: current_epoch,
             exit: exit.epoch
         }
     );
@@ -57,9 +59,9 @@ pub fn verify_exit<T: EthSpec>(
         .activation_epoch
         .safe_add(spec.shard_committee_period)?;
     verify!(
-        state.current_epoch() >= earliest_exit_epoch,
+        current_epoch >= earliest_exit_epoch,
         ExitInvalid::TooYoungToExit {
-            current_epoch: state.current_epoch(),
+            current_epoch,
             earliest_exit_epoch,
         }
     );

consensus/state_processing/src/verify_operation.rs

@@ -134,7 +134,7 @@ impl<E: EthSpec> VerifyOperation<E> for SignedVoluntaryExit {
         state: &BeaconState<E>,
         spec: &ChainSpec,
     ) -> Result<SigVerifiedOp<Self, E>, Self::Error> {
-        verify_exit(state, &self, VerifySignatures::True, spec)?;
+        verify_exit(state, None, &self, VerifySignatures::True, spec)?;
         Ok(SigVerifiedOp::new(self, state))
     }
 
@@ -205,3 +205,35 @@ impl<E: EthSpec> VerifyOperation<E> for SignedBlsToExecutionChange {
         smallvec![]
     }
 }
+
+/// Trait for operations that can be verified and transformed into a
+/// `SigVerifiedOp`.
+///
+/// The `At` suffix indicates that we can specify a particular epoch at which to
+/// verify the operation.
+pub trait VerifyOperationAt<E: EthSpec>: VerifyOperation<E> + Sized {
+    fn validate_at(
+        self,
+        state: &BeaconState<E>,
+        validate_at_epoch: Epoch,
+        spec: &ChainSpec,
+    ) -> Result<SigVerifiedOp<Self, E>, Self::Error>;
+}
+
+impl<E: EthSpec> VerifyOperationAt<E> for SignedVoluntaryExit {
+    fn validate_at(
+        self,
+        state: &BeaconState<E>,
+        validate_at_epoch: Epoch,
+        spec: &ChainSpec,
+    ) -> Result<SigVerifiedOp<Self, E>, Self::Error> {
+        verify_exit(
+            state,
+            Some(validate_at_epoch),
+            &self,
+            VerifySignatures::True,
+            spec,
+        )?;
+        Ok(SigVerifiedOp::new(self, state))
+    }
+}