Cache target attester balances for unrealized FFG progression calculation (#4362)

## Issue Addressed

#4118 

## Proposed Changes

This PR introduces a "progressive balances" cache on the `BeaconState`, which keeps track of the accumulated target attestation balance for the current & previous epochs. The cached values are utilised by fork choice to calculate unrealized justification and finalization (instead of converting epoch participation arrays to balances for each block we receive).

This optimization will be rolled out gradually to allow for more testing. A new `--progressive-balances disabled|checked|strict|fast` flag is introduced to support this:
- `checked`: enabled with checks against participation cache, and falls back to the existing epoch processing calculation if there is a total target attester balance mismatch. There is no performance gain from this as the participation cache still needs to be computed. **This is the default mode for now.**
- `strict`: enabled with checks against participation cache, returns error if there is a mismatch. **Used for testing only**.
- `fast`: enabled with no comparative checks and without computing the participation cache. This mode gives us the performance gains from the optimization. This is still experimental and not currently recommended for production usage, but will become the default mode in a future release.
- `disabled`: disable the usage of progressive cache, and use the existing method for FFG progression calculation. This mode may be useful if we find a bug and want to stop the frequent error logs.

### Tasks

- [x] Initial cache implementation in `BeaconState`
- [x] Perform checks in fork choice to compare the progressive balances cache against results from `ParticipationCache`
- [x] Add CLI flag, and disable the optimization by default
- [x] Testing on Goerli & Benchmarking
- [x]  Move caching logic from state processing to the `ProgressiveBalancesCache` (see [this comment](https://github.com/sigp/lighthouse/pull/4362#discussion_r1230877001))
- [x] Add attesting balance metrics



Co-authored-by: Jimmy Chen <jimmy@sigmaprime.io>

consensus/types/src/beacon_state.rs

@@ -26,6 +26,8 @@ pub use self::committee_cache::{
     compute_committee_index_in_epoch, compute_committee_range_in_epoch, epoch_committee_count,
     CommitteeCache,
 };
+pub use crate::beacon_state::balance::Balance;
+pub use crate::beacon_state::progressive_balances_cache::*;
 use crate::historical_summary::HistoricalSummary;
 pub use clone_config::CloneConfig;
 pub use eth_spec::*;
@@ -34,9 +36,11 @@ pub use tree_hash_cache::BeaconTreeHashCache;
 
 #[macro_use]
 mod committee_cache;
+mod balance;
 mod clone_config;
 mod exit_cache;
 mod iter;
+mod progressive_balances_cache;
 mod pubkey_cache;
 mod tests;
 mod tree_hash_cache;
@@ -101,6 +105,9 @@ pub enum Error {
     SszTypesError(ssz_types::Error),
     TreeHashCacheNotInitialized,
     NonLinearTreeHashCacheHistory,
+    ParticipationCacheError(String),
+    ProgressiveBalancesCacheNotInitialized,
+    ProgressiveBalancesCacheInconsistent,
     TreeHashCacheSkippedSlot {
         cache: Slot,
         state: Slot,
@@ -317,6 +324,12 @@ where
     #[tree_hash(skip_hashing)]
     #[test_random(default)]
     #[derivative(Clone(clone_with = "clone_default"))]
+    pub progressive_balances_cache: ProgressiveBalancesCache,
+    #[serde(skip_serializing, skip_deserializing)]
+    #[ssz(skip_serializing, skip_deserializing)]
+    #[tree_hash(skip_hashing)]
+    #[test_random(default)]
+    #[derivative(Clone(clone_with = "clone_default"))]
     pub committee_caches: [CommitteeCache; CACHED_EPOCHS],
     #[serde(skip_serializing, skip_deserializing)]
     #[ssz(skip_serializing, skip_deserializing)]
@@ -393,6 +406,7 @@ impl<T: EthSpec> BeaconState<T> {
 
             // Caching (not in spec)
             total_active_balance: None,
+            progressive_balances_cache: <_>::default(),
             committee_caches: [
                 CommitteeCache::default(),
                 CommitteeCache::default(),
@@ -757,7 +771,7 @@ impl<T: EthSpec> BeaconState<T> {
         Ok(signature_hash_int.safe_rem(modulo)? == 0)
     }
 
-    /// Returns the beacon proposer index for the `slot` in the given `relative_epoch`.
+    /// Returns the beacon proposer index for the `slot` in `self.current_epoch()`.
     ///
     /// Spec v0.12.1
     pub fn get_beacon_proposer_index(&self, slot: Slot, spec: &ChainSpec) -> Result<usize, Error> {
@@ -1150,12 +1164,30 @@ impl<T: EthSpec> BeaconState<T> {
     }
 
     /// Convenience accessor for validators and balances simultaneously.
-    pub fn validators_and_balances_mut(&mut self) -> (&mut [Validator], &mut [u64]) {
+    pub fn validators_and_balances_and_progressive_balances_mut(
+        &mut self,
+    ) -> (&mut [Validator], &mut [u64], &mut ProgressiveBalancesCache) {
         match self {
-            BeaconState::Base(state) => (&mut state.validators, &mut state.balances),
-            BeaconState::Altair(state) => (&mut state.validators, &mut state.balances),
-            BeaconState::Merge(state) => (&mut state.validators, &mut state.balances),
-            BeaconState::Capella(state) => (&mut state.validators, &mut state.balances),
+            BeaconState::Base(state) => (
+                &mut state.validators,
+                &mut state.balances,
+                &mut state.progressive_balances_cache,
+            ),
+            BeaconState::Altair(state) => (
+                &mut state.validators,
+                &mut state.balances,
+                &mut state.progressive_balances_cache,
+            ),
+            BeaconState::Merge(state) => (
+                &mut state.validators,
+                &mut state.balances,
+                &mut state.progressive_balances_cache,
+            ),
+            BeaconState::Capella(state) => (
+                &mut state.validators,
+                &mut state.balances,
+                &mut state.progressive_balances_cache,
+            ),
         }
     }
 
@@ -1380,7 +1412,7 @@ impl<T: EthSpec> BeaconState<T> {
     }
 
     /// Build all caches (except the tree hash cache), if they need to be built.
-    pub fn build_all_caches(&mut self, spec: &ChainSpec) -> Result<(), Error> {
+    pub fn build_caches(&mut self, spec: &ChainSpec) -> Result<(), Error> {
         self.build_all_committee_caches(spec)?;
         self.update_pubkey_cache()?;
         self.build_exit_cache(spec)?;
@@ -1412,6 +1444,7 @@ impl<T: EthSpec> BeaconState<T> {
         self.drop_committee_cache(RelativeEpoch::Next)?;
         self.drop_pubkey_cache();
         self.drop_tree_hash_cache();
+        self.drop_progressive_balances_cache();
         *self.exit_cache_mut() = ExitCache::default();
         Ok(())
     }
@@ -1608,6 +1641,11 @@ impl<T: EthSpec> BeaconState<T> {
         *self.pubkey_cache_mut() = PubkeyCache::default()
     }
 
+    /// Completely drops the `progressive_balances_cache` cache, replacing it with a new, empty cache.
+    fn drop_progressive_balances_cache(&mut self) {
+        *self.progressive_balances_cache_mut() = ProgressiveBalancesCache::default();
+    }
+
     /// Initialize but don't fill the tree hash cache, if it isn't already initialized.
     pub fn initialize_tree_hash_cache(&mut self) {
         if !self.tree_hash_cache().is_initialized() {
@@ -1679,6 +1717,9 @@ impl<T: EthSpec> BeaconState<T> {
         if config.tree_hash_cache {
             *res.tree_hash_cache_mut() = self.tree_hash_cache().clone();
         }
+        if config.progressive_balances_cache {
+            *res.progressive_balances_cache_mut() = self.progressive_balances_cache().clone();
+        }
         res
     }
 

consensus/types/src/beacon_state/balance.rs

@@ -0,0 +1,33 @@
+use arbitrary::Arbitrary;
+use safe_arith::{ArithError, SafeArith};
+
+/// A balance which will never be below the specified `minimum`.
+///
+/// This is an effort to ensure the `EFFECTIVE_BALANCE_INCREMENT` minimum is always respected.
+#[derive(PartialEq, Debug, Clone, Copy, Arbitrary)]
+pub struct Balance {
+    raw: u64,
+    minimum: u64,
+}
+
+impl Balance {
+    /// Initialize the balance to `0`, or the given `minimum`.
+    pub fn zero(minimum: u64) -> Self {
+        Self { raw: 0, minimum }
+    }
+
+    /// Returns the balance with respect to the initialization `minimum`.
+    pub fn get(&self) -> u64 {
+        std::cmp::max(self.raw, self.minimum)
+    }
+
+    /// Add-assign to the balance.
+    pub fn safe_add_assign(&mut self, other: u64) -> Result<(), ArithError> {
+        self.raw.safe_add_assign(other)
+    }
+
+    /// Sub-assign to the balance.
+    pub fn safe_sub_assign(&mut self, other: u64) -> Result<(), ArithError> {
+        self.raw.safe_sub_assign(other)
+    }
+}

consensus/types/src/beacon_state/clone_config.rs

@@ -5,6 +5,7 @@ pub struct CloneConfig {
     pub pubkey_cache: bool,
     pub exit_cache: bool,
     pub tree_hash_cache: bool,
+    pub progressive_balances_cache: bool,
 }
 
 impl CloneConfig {
@@ -14,6 +15,7 @@ impl CloneConfig {
             pubkey_cache: true,
             exit_cache: true,
             tree_hash_cache: true,
+            progressive_balances_cache: true,
         }
     }
 

consensus/types/src/beacon_state/progressive_balances_cache.rs

@@ -0,0 +1,184 @@
+use crate::beacon_state::balance::Balance;
+use crate::{BeaconState, BeaconStateError, ChainSpec, Epoch, EthSpec};
+use arbitrary::Arbitrary;
+use safe_arith::SafeArith;
+use serde_derive::{Deserialize, Serialize};
+use strum::{Display, EnumString, EnumVariantNames};
+
+/// This cache keeps track of the accumulated target attestation balance for the current & previous
+/// epochs. The cached values can be utilised by fork choice to calculate unrealized justification
+/// and finalization instead of converting epoch participation arrays to balances for each block we
+/// process.
+#[derive(Default, Debug, PartialEq, Arbitrary, Clone)]
+pub struct ProgressiveBalancesCache {
+    inner: Option<Inner>,
+}
+
+#[derive(Debug, PartialEq, Arbitrary, Clone)]
+struct Inner {
+    pub current_epoch: Epoch,
+    pub previous_epoch_target_attesting_balance: Balance,
+    pub current_epoch_target_attesting_balance: Balance,
+}
+
+impl ProgressiveBalancesCache {
+    pub fn initialize(
+        &mut self,
+        current_epoch: Epoch,
+        previous_epoch_target_attesting_balance: Balance,
+        current_epoch_target_attesting_balance: Balance,
+    ) {
+        self.inner = Some(Inner {
+            current_epoch,
+            previous_epoch_target_attesting_balance,
+            current_epoch_target_attesting_balance,
+        });
+    }
+
+    pub fn is_initialized(&self) -> bool {
+        self.inner.is_some()
+    }
+
+    /// When a new target attestation has been processed, we update the cached
+    /// `current_epoch_target_attesting_balance` to include the validator effective balance.
+    /// If the epoch is neither the current epoch nor the previous epoch, an error is returned.
+    pub fn on_new_target_attestation(
+        &mut self,
+        epoch: Epoch,
+        validator_effective_balance: u64,
+    ) -> Result<(), BeaconStateError> {
+        let cache = self.get_inner_mut()?;
+
+        if epoch == cache.current_epoch {
+            cache
+                .current_epoch_target_attesting_balance
+                .safe_add_assign(validator_effective_balance)?;
+        } else if epoch.safe_add(1)? == cache.current_epoch {
+            cache
+                .previous_epoch_target_attesting_balance
+                .safe_add_assign(validator_effective_balance)?;
+        } else {
+            return Err(BeaconStateError::ProgressiveBalancesCacheInconsistent);
+        }
+
+        Ok(())
+    }
+
+    /// When a validator is slashed, we reduce the `current_epoch_target_attesting_balance` by the
+    /// validator's effective balance to exclude the validator weight.
+    pub fn on_slashing(
+        &mut self,
+        is_previous_epoch_target_attester: bool,
+        is_current_epoch_target_attester: bool,
+        effective_balance: u64,
+    ) -> Result<(), BeaconStateError> {
+        let cache = self.get_inner_mut()?;
+        if is_previous_epoch_target_attester {
+            cache
+                .previous_epoch_target_attesting_balance
+                .safe_sub_assign(effective_balance)?;
+        }
+        if is_current_epoch_target_attester {
+            cache
+                .current_epoch_target_attesting_balance
+                .safe_sub_assign(effective_balance)?;
+        }
+        Ok(())
+    }
+
+    /// When a current epoch target attester has its effective balance changed, we adjust the
+    /// its share of the target attesting balance in the cache.
+    pub fn on_effective_balance_change(
+        &mut self,
+        is_current_epoch_target_attester: bool,
+        old_effective_balance: u64,
+        new_effective_balance: u64,
+    ) -> Result<(), BeaconStateError> {
+        let cache = self.get_inner_mut()?;
+        if is_current_epoch_target_attester {
+            if new_effective_balance > old_effective_balance {
+                cache
+                    .current_epoch_target_attesting_balance
+                    .safe_add_assign(new_effective_balance.safe_sub(old_effective_balance)?)?;
+            } else {
+                cache
+                    .current_epoch_target_attesting_balance
+                    .safe_sub_assign(old_effective_balance.safe_sub(new_effective_balance)?)?;
+            }
+        }
+        Ok(())
+    }
+
+    /// On epoch transition, the balance from current epoch is shifted to previous epoch, and the
+    /// current epoch balance is reset to 0.
+    pub fn on_epoch_transition(&mut self, spec: &ChainSpec) -> Result<(), BeaconStateError> {
+        let cache = self.get_inner_mut()?;
+        cache.current_epoch.safe_add_assign(1)?;
+        cache.previous_epoch_target_attesting_balance =
+            cache.current_epoch_target_attesting_balance;
+        cache.current_epoch_target_attesting_balance =
+            Balance::zero(spec.effective_balance_increment);
+        Ok(())
+    }
+
+    pub fn previous_epoch_target_attesting_balance(&self) -> Result<u64, BeaconStateError> {
+        Ok(self
+            .get_inner()?
+            .previous_epoch_target_attesting_balance
+            .get())
+    }
+
+    pub fn current_epoch_target_attesting_balance(&self) -> Result<u64, BeaconStateError> {
+        Ok(self
+            .get_inner()?
+            .current_epoch_target_attesting_balance
+            .get())
+    }
+
+    fn get_inner_mut(&mut self) -> Result<&mut Inner, BeaconStateError> {
+        self.inner
+            .as_mut()
+            .ok_or(BeaconStateError::ProgressiveBalancesCacheNotInitialized)
+    }
+
+    fn get_inner(&self) -> Result<&Inner, BeaconStateError> {
+        self.inner
+            .as_ref()
+            .ok_or(BeaconStateError::ProgressiveBalancesCacheNotInitialized)
+    }
+}
+
+#[derive(
+    Debug, PartialEq, Eq, Clone, Copy, Deserialize, Serialize, Display, EnumString, EnumVariantNames,
+)]
+#[strum(serialize_all = "lowercase")]
+pub enum ProgressiveBalancesMode {
+    /// Disable the usage of progressive cache, and use the existing `ParticipationCache` calculation.
+    Disabled,
+    /// Enable the usage of progressive cache, with checks against the `ParticipationCache` and falls
+    /// back to the existing calculation if there is a balance mismatch.
+    Checked,
+    /// Enable the usage of progressive cache, with checks against the `ParticipationCache`. Errors
+    /// if there is a balance mismatch. Used in testing only.
+    Strict,
+    /// Enable the usage of progressive cache, with no comparative checks against the
+    /// `ParticipationCache`. This is fast but an experimental mode, use with caution.
+    Fast,
+}
+
+impl ProgressiveBalancesMode {
+    pub fn perform_comparative_checks(&self) -> bool {
+        match self {
+            Self::Disabled | Self::Fast => false,
+            Self::Checked | Self::Strict => true,
+        }
+    }
+}
+
+/// `ProgressiveBalancesCache` is only enabled from `Altair` as it requires `ParticipationCache`.
+pub fn is_progressive_balances_enabled<E: EthSpec>(state: &BeaconState<E>) -> bool {
+    match state {
+        BeaconState::Base(_) => false,
+        BeaconState::Altair(_) | BeaconState::Merge(_) | BeaconState::Capella(_) => true,
+    }
+}

consensus/types/src/beacon_state/tests.rs

@@ -219,17 +219,18 @@ async fn clone_config() {
 
     let mut state = build_state::<MinimalEthSpec>(16).await;
 
-    state.build_all_caches(&spec).unwrap();
+    state.build_caches(&spec).unwrap();
     state
         .update_tree_hash_cache()
         .expect("should update tree hash cache");
 
-    let num_caches = 4;
+    let num_caches = 5;
     let all_configs = (0..2u8.pow(num_caches)).map(|i| CloneConfig {
         committee_caches: (i & 1) != 0,
         pubkey_cache: ((i >> 1) & 1) != 0,
         exit_cache: ((i >> 2) & 1) != 0,
         tree_hash_cache: ((i >> 3) & 1) != 0,
+        progressive_balances_cache: ((i >> 4) & 1) != 0,
     });
 
     for config in all_configs {