Cache target attester balances for unrealized FFG progression calcula…

…tion (sigp#4362)

## Issue Addressed

sigp#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](sigp#4362 (comment)))
- [x] Add attesting balance metrics



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

beacon_node/beacon_chain/src/beacon_chain.rs

@@ -2898,7 +2898,9 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
                     block_delay,
                     &state,
                     payload_verification_status,
+                    self.config.progressive_balances_mode,
                     &self.spec,
+                    &self.log,
                 )
                 .map_err(|e| BlockError::BeaconChainError(e.into()))?;
         }

beacon_node/beacon_chain/src/builder.rs

@@ -338,7 +338,7 @@ where
         let beacon_block = genesis_block(&mut beacon_state, &self.spec)?;
 
         beacon_state
-            .build_all_caches(&self.spec)
+            .build_caches(&self.spec)
             .map_err(|e| format!("Failed to build genesis state caches: {:?}", e))?;
 
         let beacon_state_root = beacon_block.message().state_root();
@@ -437,7 +437,7 @@ where
         // Prime all caches before storing the state in the database and computing the tree hash
         // root.
         weak_subj_state
-            .build_all_caches(&self.spec)
+            .build_caches(&self.spec)
             .map_err(|e| format!("Error building caches on checkpoint state: {e:?}"))?;
 
         let computed_state_root = weak_subj_state
@@ -687,6 +687,8 @@ where
                 store.clone(),
                 Some(current_slot),
                 &self.spec,
+                self.chain_config.progressive_balances_mode,
+                &log,
             )?;
         }
 
@@ -700,7 +702,7 @@ where
 
         head_snapshot
             .beacon_state
-            .build_all_caches(&self.spec)
+            .build_caches(&self.spec)
             .map_err(|e| format!("Failed to build state caches: {:?}", e))?;
 
         // Perform a check to ensure that the finalization points of the head and fork choice are

beacon_node/beacon_chain/src/chain_config.rs

@@ -1,7 +1,7 @@
 pub use proto_array::{DisallowedReOrgOffsets, ReOrgThreshold};
 use serde_derive::{Deserialize, Serialize};
 use std::time::Duration;
-use types::{Checkpoint, Epoch};
+use types::{Checkpoint, Epoch, ProgressiveBalancesMode};
 
 pub const DEFAULT_RE_ORG_THRESHOLD: ReOrgThreshold = ReOrgThreshold(20);
 pub const DEFAULT_RE_ORG_MAX_EPOCHS_SINCE_FINALIZATION: Epoch = Epoch::new(2);
@@ -81,6 +81,8 @@ pub struct ChainConfig {
     pub always_prepare_payload: bool,
     /// Whether backfill sync processing should be rate-limited.
     pub enable_backfill_rate_limiting: bool,
+    /// Whether to use `ProgressiveBalancesCache` in unrealized FFG progression calculation.
+    pub progressive_balances_mode: ProgressiveBalancesMode,
 }
 
 impl Default for ChainConfig {
@@ -111,6 +113,7 @@ impl Default for ChainConfig {
             genesis_backfill: false,
             always_prepare_payload: false,
             enable_backfill_rate_limiting: true,
+            progressive_balances_mode: ProgressiveBalancesMode::Checked,
         }
     }
 }

beacon_node/beacon_chain/src/fork_revert.rs

@@ -10,7 +10,10 @@ use state_processing::{
 use std::sync::Arc;
 use std::time::Duration;
 use store::{iter::ParentRootBlockIterator, HotColdDB, ItemStore};
-use types::{BeaconState, ChainSpec, EthSpec, ForkName, Hash256, SignedBeaconBlock, Slot};
+use types::{
+    BeaconState, ChainSpec, EthSpec, ForkName, Hash256, ProgressiveBalancesMode, SignedBeaconBlock,
+    Slot,
+};
 
 const CORRUPT_DB_MESSAGE: &str = "The database could be corrupt. Check its file permissions or \
                                   consider deleting it by running with the --purge-db flag.";
@@ -100,6 +103,8 @@ pub fn reset_fork_choice_to_finalization<E: EthSpec, Hot: ItemStore<E>, Cold: It
     store: Arc<HotColdDB<E, Hot, Cold>>,
     current_slot: Option<Slot>,
     spec: &ChainSpec,
+    progressive_balances_mode: ProgressiveBalancesMode,
+    log: &Logger,
 ) -> Result<ForkChoice<BeaconForkChoiceStore<E, Hot, Cold>, E>, String> {
     // Fetch finalized block.
     let finalized_checkpoint = head_state.finalized_checkpoint();
@@ -197,7 +202,9 @@ pub fn reset_fork_choice_to_finalization<E: EthSpec, Hot: ItemStore<E>, Cold: It
                 Duration::from_secs(0),
                 &state,
                 payload_verification_status,
+                progressive_balances_mode,
                 spec,
+                log,
             )
             .map_err(|e| format!("Error applying replayed block to fork choice: {:?}", e))?;
     }

beacon_node/beacon_chain/src/test_utils.rs

@@ -754,9 +754,7 @@ where
         complete_state_advance(&mut state, None, slot, &self.spec)
             .expect("should be able to advance state to slot");
 
-        state
-            .build_all_caches(&self.spec)
-            .expect("should build caches");
+        state.build_caches(&self.spec).expect("should build caches");
 
         let proposer_index = state.get_beacon_proposer_index(slot, &self.spec).unwrap();
 
@@ -803,9 +801,7 @@ where
         complete_state_advance(&mut state, None, slot, &self.spec)
             .expect("should be able to advance state to slot");
 
-        state
-            .build_all_caches(&self.spec)
-            .expect("should build caches");
+        state.build_caches(&self.spec).expect("should build caches");
 
         let proposer_index = state.get_beacon_proposer_index(slot, &self.spec).unwrap();
 
@@ -1523,6 +1519,36 @@ where
         .sign(sk, &fork, genesis_validators_root, &self.chain.spec)
     }
 
+    pub fn add_proposer_slashing(&self, validator_index: u64) -> Result<(), String> {
+        let propposer_slashing = self.make_proposer_slashing(validator_index);
+        if let ObservationOutcome::New(verified_proposer_slashing) = self
+            .chain
+            .verify_proposer_slashing_for_gossip(propposer_slashing)
+            .expect("should verify proposer slashing for gossip")
+        {
+            self.chain
+                .import_proposer_slashing(verified_proposer_slashing);
+            Ok(())
+        } else {
+            Err("should observe new proposer slashing".to_string())
+        }
+    }
+
+    pub fn add_attester_slashing(&self, validator_indices: Vec<u64>) -> Result<(), String> {
+        let attester_slashing = self.make_attester_slashing(validator_indices);
+        if let ObservationOutcome::New(verified_attester_slashing) = self
+            .chain
+            .verify_attester_slashing_for_gossip(attester_slashing)
+            .expect("should verify attester slashing for gossip")
+        {
+            self.chain
+                .import_attester_slashing(verified_attester_slashing);
+            Ok(())
+        } else {
+            Err("should observe new attester slashing".to_string())
+        }
+    }
+
     pub fn add_bls_to_execution_change(
         &self,
         validator_index: u64,

beacon_node/beacon_chain/tests/capella.rs

@@ -133,13 +133,8 @@ async fn base_altair_merge_capella() {
     for _ in (merge_fork_slot.as_u64() + 3)..capella_fork_slot.as_u64() {
         harness.extend_slots(1).await;
         let block = &harness.chain.head_snapshot().beacon_block;
-        let full_payload: FullPayload<E> = block
-            .message()
-            .body()
-            .execution_payload()
-            .unwrap()
-            .clone()
-            .into();
+        let full_payload: FullPayload<E> =
+            block.message().body().execution_payload().unwrap().into();
         // pre-capella shouldn't have withdrawals
         assert!(full_payload.withdrawals_root().is_err());
         execution_payloads.push(full_payload);
@@ -151,13 +146,8 @@ async fn base_altair_merge_capella() {
     for _ in 0..16 {
         harness.extend_slots(1).await;
         let block = &harness.chain.head_snapshot().beacon_block;
-        let full_payload: FullPayload<E> = block
-            .message()
-            .body()
-            .execution_payload()
-            .unwrap()
-            .clone()
-            .into();
+        let full_payload: FullPayload<E> =
+            block.message().body().execution_payload().unwrap().into();
         // post-capella should have withdrawals
         assert!(full_payload.withdrawals_root().is_ok());
         execution_payloads.push(full_payload);

beacon_node/beacon_chain/tests/payload_invalidation.rs

@@ -1064,8 +1064,9 @@ async fn invalid_parent() {
             Duration::from_secs(0),
             &state,
             PayloadVerificationStatus::Optimistic,
+            rig.harness.chain.config.progressive_balances_mode,
             &rig.harness.chain.spec,
-
+            rig.harness.logger()
         ),
         Err(ForkChoiceError::ProtoArrayStringError(message))
         if message.contains(&format!(

beacon_node/http_api/src/block_rewards.rs

@@ -49,7 +49,7 @@ pub fn get_block_rewards<T: BeaconChainTypes>(
         .map_err(beacon_chain_error)?;
 
     state
-        .build_all_caches(&chain.spec)
+        .build_caches(&chain.spec)
         .map_err(beacon_state_error)?;
 
     let mut reward_cache = Default::default();

beacon_node/src/cli.rs

@@ -1,5 +1,6 @@
 use clap::{App, Arg};
 use strum::VariantNames;
+use types::ProgressiveBalancesMode;
 
 pub fn cli_app<'a, 'b>() -> App<'a, 'b> {
     App::new("beacon_node")
@@ -1117,4 +1118,17 @@ pub fn cli_app<'a, 'b>() -> App<'a, 'b> {
                     developers. This directory is not pruned, users should be careful to avoid \
                     filling up their disks.")
         )
+        .arg(
+            Arg::with_name("progressive-balances")
+                .long("progressive-balances")
+                .value_name("MODE")
+                .help("Options to enable or disable the progressive balances cache for \
+                        unrealized FFG progression calculation. The default `checked` mode compares \
+                        the progressive balances from the cache against results from the existing \
+                        method. If there is a mismatch, it falls back to the existing method. The \
+                        optimized mode (`fast`) is faster but is still experimental, and is \
+                        not recommended for mainnet usage at this time.")
+                .takes_value(true)
+                .possible_values(ProgressiveBalancesMode::VARIANTS)
+        )
 }

beacon_node/src/config.rs

@@ -800,6 +800,12 @@ pub fn get_config<E: EthSpec>(
         client_config.network.invalid_block_storage = Some(path);
     }
 
+    if let Some(progressive_balances_mode) =
+        clap_utils::parse_optional(cli_args, "progressive-balances")?
+    {
+        client_config.chain.progressive_balances_mode = progressive_balances_mode;
+    }
+
     Ok(client_config)
 }
 

beacon_node/store/src/partial_beacon_state.rs

@@ -373,6 +373,7 @@ macro_rules! impl_try_into_beacon_state {
 
             // Caching
             total_active_balance: <_>::default(),
+            progressive_balances_cache: <_>::default(),
             committee_caches: <_>::default(),
             pubkey_cache: <_>::default(),
             exit_cache: <_>::default(),

beacon_node/store/src/reconstruct.rs

@@ -63,7 +63,7 @@ where
             .load_cold_state_by_slot(lower_limit_slot)?
             .ok_or(HotColdDBError::MissingLowerLimitState(lower_limit_slot))?;
 
-        state.build_all_caches(&self.spec)?;
+        state.build_caches(&self.spec)?;
 
         process_results(block_root_iter, |iter| -> Result<(), Error> {
             let mut io_batch = vec![];