perf(beacon): cache the total active balance per epoch

crates/blockchain/state_transition/src/beacon/block_production.rs

@@ -38,8 +38,8 @@ use super::bls;
 use super::config::Config;
 use super::error::{Error, Result, verify};
 use super::helpers::accessors::{
-    CommitteeCache, get_beacon_proposer_index, get_block_root, get_current_epoch,
-    get_previous_epoch, get_randao_mix,
+    ActiveBalanceCache, CommitteeCache, get_beacon_proposer_index, get_block_root,
+    get_current_epoch, get_previous_epoch, get_randao_mix,
 };
 use super::helpers::electra::{
     get_attesting_indices, get_indexed_attestation, is_valid_indexed_attestation,
@@ -343,6 +343,7 @@ pub fn assemble_block(
         config,
         &ExecutionEngine::valid(),
         &CommitteeCache::default(),
+        &ActiveBalanceCache::default(),
     )?;
     block.state_root = post.hash_tree_root();
     Ok(block)
@@ -435,6 +436,7 @@ mod tests {
             &Config::mainnet(),
             &ExecutionEngine::valid(),
             &CommitteeCache::default(),
+            &ActiveBalanceCache::default(),
         )
         .unwrap();
         assert_eq!(block.state_root, post.hash_tree_root());

crates/blockchain/state_transition/src/beacon/fork_choice.rs

@@ -2370,8 +2370,18 @@ pub fn on_block(
             stf::ExecutionEngine::valid()
         }
     };
-    let transition =
-        stf::state_transition(&mut state, &signed_block, true, config, &engine, committees);
+    // The store's own cache, so every block of this chain shares one entry per
+    // epoch; see `ActiveBalanceCacheExt`.
+    let active_balances = store.active_balance_cache();
+    let transition = stf::state_transition(
+        &mut state,
+        &signed_block,
+        true,
+        config,
+        &engine,
+        committees,
+        &active_balances,
+    );
 
     // `optimistic-sync.md`: a block deemed `INVALIDATED` MUST NOT be included
     // in the canonical chain. That is stated here, on the verdict, rather than

crates/blockchain/state_transition/src/beacon/helpers/accessors.rs

@@ -11,7 +11,9 @@ use std::sync::Arc;
 // and so on), so every existing import of the type this module used to define
 // is unaffected by its move to `ethlambda-storage`; see [`CommitteeCacheExt`]
 // below for what this crate still contributes.
-pub use ethlambda_storage::{CommitteeCache, Lookup, ShufflingKey};
+pub use ethlambda_storage::{
+    ActiveBalanceCache, ActiveBalanceKey, CommitteeCache, Lookup, ShufflingKey,
+};
 
 use crate::beacon::config::Config;
 use crate::beacon::constants;
@@ -301,6 +303,96 @@ impl CommitteeCacheExt for CommitteeCache {
     }
 }
 
+/// The key `state`'s total active balance is cached under, or `None` if
+/// `state` cannot name the block that fixes it: the state has not advanced
+/// past that block's slot (the genesis state, asked about its own epoch), or
+/// the slot has fallen out of its `SLOTS_PER_HISTORICAL_ROOT` window. `None`
+/// means the lookup is computed for its caller alone and not cached.
+///
+/// The total for a state in epoch `E` sums the effective balances of the
+/// validators active at `E`, and both inputs are fixed once epoch `E - 1`
+/// ends. Effective balances are written only by `process_effective_balance_updates`,
+/// at the end of `E - 1`; an epoch's update is the last writer before `E`'s
+/// states exist. The active set moves only through `activation_epoch` and
+/// `exit_epoch`, and every assignment goes through `compute_activation_exit_epoch`,
+/// at least `MAX_SEED_LOOKAHEAD` epochs ahead, so none can land in `E` for `E`
+/// itself. Slashing sets `slashed` and queues an exit but leaves effective
+/// balances alone, and the spec's total keeps slashed validators. Deposits
+/// append validators that are not yet active, and the Electra upgrade zeroes
+/// only never-activated validators. So the block root at the last slot of
+/// `E - 1` identifies the history that determines the total: two states
+/// agreeing on it agree on the total, however much they disagree after it.
+///
+/// That is one epoch later than [`shuffling_key`]'s root (the last slot of
+/// `E - 2`), because the shuffle's inputs are fixed `MIN_SEED_LOOKAHEAD` epochs
+/// ahead of use and the total's are not: an effective balance written at the
+/// end of `E - 1` is part of the total for `E`, so a key rooted at `E - 2`
+/// would let two branches that diverge during `E - 1` share one entry.
+///
+/// # Epoch 0
+///
+/// There is no epoch before it, so it takes the genesis block, at slot 0, as
+/// its deciding block, as [`shuffling_key`] does for its first epochs. Nothing
+/// after genesis can reach either input within epoch 0 (the effective-balance
+/// update at its end writes epoch 1's), and a state from a different genesis
+/// has a different genesis block root.
+///
+/// # Callers must be block processing
+///
+/// Between `process_effective_balance_updates` and the slot increment that
+/// follows it, a state is still in epoch `E - 1` but already carries epoch
+/// `E`'s effective balances, so its total would not match a key rooted at
+/// `E - 2`'s end. The only callers routed through the cache are block
+/// processing's (attestations, sync aggregate), which never run in that
+/// window: a block is processed on a state that `process_slots` has already
+/// advanced to the block's slot. Epoch processing keeps calling the
+/// uncached [`get_total_active_balance`].
+fn active_balance_key(state: &BeaconState) -> Option<ActiveBalanceKey> {
+    let epoch = get_current_epoch(state);
+    let decision_slot = compute_start_slot_at_epoch(epoch).saturating_sub(1);
+    let decision_root = get_block_root_at_slot(state, decision_slot).ok()?;
+    Some(ActiveBalanceKey {
+        epoch,
+        decision_root,
+    })
+}
+
+/// Extends `ethlambda-storage`'s [`ActiveBalanceCache`] with the consensus
+/// logic that keys and computes its entries, for the same reason
+/// [`CommitteeCacheExt`] lives here and not in `ethlambda-storage`.
+pub trait ActiveBalanceCacheExt {
+    /// `get_total_active_balance(state)`, computed once per key and then
+    /// served to every later caller naming the same one. Equal to the
+    /// uncached function for every state it is given; see [`active_balance_key`]
+    /// for why, and for who may call it.
+    fn total_active_balance(&self, state: &BeaconState) -> Result<Gwei>;
+}
+
+impl ActiveBalanceCacheExt for ActiveBalanceCache {
+    fn total_active_balance(&self, state: &BeaconState) -> Result<Gwei> {
+        let Some(key) = active_balance_key(state) else {
+            crate::metrics::inc_total_active_balance_lookups("unkeyable");
+            return get_total_active_balance(state);
+        };
+
+        let (total, lookup) = self.get_or_compute(key, || compute_total_active_balance(state));
+        crate::metrics::inc_total_active_balance_lookups(match lookup {
+            Lookup::Hit => "hit",
+            Lookup::Miss => "miss",
+        });
+        // The cross-check the key's soundness argument rests on: a stale hit
+        // would mean some writer moved an input mid-epoch.
+        if lookup == Lookup::Hit {
+            debug_assert_eq!(
+                total,
+                compute_total_active_balance(state),
+                "stale total active balance cache entry"
+            );
+        }
+        Ok(total)
+    }
+}
+
 /// The committee at `slot` with index `index`.
 ///
 /// One epoch's active set is shuffled once and then split across every slot and
@@ -398,9 +490,28 @@ pub fn get_total_balance(state: &BeaconState, indices: &[ValidatorIndex]) -> Res
 }
 
 /// The combined effective balance of the currently active validators.
+///
+/// One in-order pass over the registry, summing as it goes: the same value as
+/// [`get_total_balance`] over [`get_active_validator_indices`] (same
+/// saturating sum, same one-increment floor), without the intermediate index
+/// list or a tree descent per active validator. `state.validators().iter()`
+/// walks the leaves, whereas `state.validator(i)` descends from the root each
+/// time.
 pub fn get_total_active_balance(state: &BeaconState) -> Result<Gwei> {
-    let indices = get_active_validator_indices(state, get_current_epoch(state));
-    get_total_balance(state, &indices)
+    Ok(compute_total_active_balance(state))
+}
+
+/// [`get_total_active_balance`]'s body, without the `Result` it never needed.
+fn compute_total_active_balance(state: &BeaconState) -> Gwei {
+    let epoch = get_current_epoch(state);
+    let total = state
+        .validators()
+        .iter()
+        .filter(|validator| is_active_validator(validator, epoch))
+        .fold(0, |sum: Gwei, validator| {
+            sum.saturating_add(validator.effective_balance)
+        });
+    total.max(preset::EFFECTIVE_BALANCE_INCREMENT)
 }
 
 /// The signing domain for `domain_type` at `epoch`, or at the current epoch when
@@ -649,6 +760,159 @@ mod tests {
         );
     }
 
+    /// SplitMix64: a tiny deterministic generator, so the randomized tests
+    /// below need no dependency.
+    struct SplitMix64(u64);
+
+    impl SplitMix64 {
+        fn next(&mut self) -> u64 {
+            self.0 = self.0.wrapping_add(0x9E37_79B9_7F4A_7C15);
+            let mut z = self.0;
+            z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
+            z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
+            z ^ (z >> 31)
+        }
+    }
+
+    /// A state at `epoch` whose validators get random activation and exit
+    /// epochs around it and random effective balances.
+    fn random_registry_state(rng: &mut SplitMix64, count: usize, epoch: Epoch) -> BeaconState {
+        let mut state = with_validators(count);
+        *state.slot_mut() = compute_start_slot_at_epoch(epoch);
+        for index in 0..count {
+            let validator = state.validator_mut(index as ValidatorIndex).unwrap();
+            validator.activation_epoch = rng.next() % (epoch + 3);
+            validator.exit_epoch = match rng.next() % 3 {
+                0 => constants::FAR_FUTURE_EPOCH,
+                _ => rng.next() % (epoch + 3),
+            };
+            validator.effective_balance = match rng.next() % 4 {
+                0 => 0,
+                1 => u64::MAX - rng.next() % 4,
+                _ => (rng.next() % 64) * preset::EFFECTIVE_BALANCE_INCREMENT,
+            };
+        }
+        state.apply_pending_mutations();
+        state
+    }
+
+    /// The one-pass total is the spec's `get_total_balance` over the active
+    /// indices, including its saturation and its floor.
+    #[test]
+    fn the_one_pass_total_matches_the_spec_formulation() {
+        let mut rng = SplitMix64(0x5EED);
+        for round in 0..64 {
+            let count = (rng.next() % 40) as usize;
+            let epoch = rng.next() % 6;
+            let state = random_registry_state(&mut rng, count, epoch);
+            let indices = get_active_validator_indices(&state, get_current_epoch(&state));
+            assert_eq!(
+                get_total_active_balance(&state).unwrap(),
+                get_total_balance(&state, &indices).unwrap(),
+                "round {round}"
+            );
+        }
+    }
+
+    /// An all-zero registry hits the floor, not zero.
+    #[test]
+    fn the_one_pass_total_is_floored_for_a_zero_registry() {
+        let mut state = with_validators(8);
+        for index in 0..8 {
+            state.validator_mut(index).unwrap().effective_balance = 0;
+        }
+        state.apply_pending_mutations();
+        assert_eq!(
+            get_total_active_balance(&state).unwrap(),
+            preset::EFFECTIVE_BALANCE_INCREMENT
+        );
+    }
+
+    /// A state at the first slot of `epoch` whose every block root is `root`.
+    fn state_at_epoch_with_root(epoch: Epoch, root: u8) -> BeaconState {
+        let mut state = with_validators(16);
+        *state.slot_mut() = compute_start_slot_at_epoch(epoch);
+        for slot in 0..preset::SLOTS_PER_HISTORICAL_ROOT {
+            state.block_roots_mut()[slot] = Root::repeat_byte(root);
+        }
+        state
+    }
+
+    /// A lookup returns the spec value and the second one is a hit.
+    #[test]
+    fn the_active_balance_cache_serves_the_spec_value() {
+        let state = state_at_epoch_with_root(3, 1);
+        let cache = ActiveBalanceCache::default();
+        let expected = get_total_active_balance(&state).unwrap();
+
+        assert_eq!(cache.total_active_balance(&state).unwrap(), expected);
+        let key = active_balance_key(&state).unwrap();
+        assert_eq!(cache.get(key), Some(expected));
+        assert_eq!(cache.total_active_balance(&state).unwrap(), expected);
+    }
+
+    /// Two sibling states in one epoch that disagree on the deciding block
+    /// keep separate entries, each with its own registry's total.
+    #[test]
+    fn sibling_states_with_different_decision_roots_get_different_entries() {
+        let cache = ActiveBalanceCache::default();
+        let a = state_at_epoch_with_root(3, 1);
+        let mut b = state_at_epoch_with_root(3, 2);
+        // Sibling `b` has a validator the fork `a` never saw.
+        b.validator_mut(0).unwrap().effective_balance = 0;
+        b.apply_pending_mutations();
+
+        let total_a = cache.total_active_balance(&a).unwrap();
+        let total_b = cache.total_active_balance(&b).unwrap();
+
+        assert_ne!(total_a, total_b);
+        assert_eq!(total_a, get_total_active_balance(&a).unwrap());
+        assert_eq!(total_b, get_total_active_balance(&b).unwrap());
+        assert_eq!(cache.get(active_balance_key(&a).unwrap()), Some(total_a));
+        assert_eq!(cache.get(active_balance_key(&b).unwrap()), Some(total_b));
+    }
+
+    /// Across an epoch boundary the key changes, so the lookup misses and
+    /// fills a new entry instead of serving the old epoch's total.
+    #[test]
+    fn crossing_an_epoch_boundary_misses_and_refills() {
+        let cache = ActiveBalanceCache::default();
+        let mut state = state_at_epoch_with_root(3, 1);
+        let before = cache.total_active_balance(&state).unwrap();
+        let old_key = active_balance_key(&state).unwrap();
+
+        // The effective-balance update at the boundary, then the new epoch
+        // with its own deciding root.
+        state.validator_mut(0).unwrap().effective_balance = 0;
+        state.apply_pending_mutations();
+        *state.slot_mut() = compute_start_slot_at_epoch(4);
+        let last_slot =
+            (compute_start_slot_at_epoch(4) - 1) as usize % preset::SLOTS_PER_HISTORICAL_ROOT;
+        state.block_roots_mut()[last_slot] = Root::repeat_byte(9);
+
+        let new_key = active_balance_key(&state).unwrap();
+        assert_ne!(old_key, new_key);
+        let after = cache.total_active_balance(&state).unwrap();
+        assert_eq!(after, get_total_active_balance(&state).unwrap());
+        assert_ne!(before, after);
+        assert_eq!(cache.get(old_key), Some(before));
+        assert_eq!(cache.get(new_key), Some(after));
+    }
+
+    /// The genesis state cannot name the deciding block, so it is computed
+    /// for its caller alone and not cached.
+    #[test]
+    fn an_unkeyable_state_is_computed_but_not_cached() {
+        let mut state = with_validators(8);
+        *state.slot_mut() = 0;
+        assert!(active_balance_key(&state).is_none());
+        let cache = ActiveBalanceCache::default();
+        assert_eq!(
+            cache.total_active_balance(&state).unwrap(),
+            get_total_active_balance(&state).unwrap()
+        );
+    }
+
     #[test]
     fn proposer_is_drawn_from_the_active_set() {
         let state = crate::beacon::helpers::test_state::with_validators(32);

crates/blockchain/state_transition/src/beacon/helpers/altair.rs

@@ -182,11 +182,14 @@ pub fn get_next_sync_committee(state: &BeaconState) -> Result<altair::SyncCommit
 /// stake actually earned it rather than assuming every attestation splits the
 /// same fixed fraction.
 pub fn get_base_reward_per_increment(state: &BeaconState) -> Result<Gwei> {
-    let total_active_balance = get_total_active_balance(state)?;
-    Ok(
-        preset::EFFECTIVE_BALANCE_INCREMENT * preset::BASE_REWARD_FACTOR
-            / integer_squareroot(total_active_balance),
-    )
+    Ok(base_reward_per_increment(get_total_active_balance(state)?))
+}
+
+/// [`get_base_reward_per_increment`] for a total active balance the caller
+/// already holds, such as one served by an `ActiveBalanceCache`.
+pub fn base_reward_per_increment(total_active_balance: Gwei) -> Gwei {
+    preset::EFFECTIVE_BALANCE_INCREMENT * preset::BASE_REWARD_FACTOR
+        / integer_squareroot(total_active_balance)
 }
 
 /// The base reward for the validator at `index`.