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

bin/ethlambda/src/era/replay.rs

@@ -18,7 +18,7 @@
 use std::path::Path;
 use std::time::Instant;
 
-use ethlambda_state_transition::beacon::helpers::accessors::CommitteeCache;
+use ethlambda_state_transition::beacon::helpers::accessors::{ActiveBalanceCache, CommitteeCache};
 use ethlambda_state_transition::beacon::stf::{self, ExecutionEngine};
 use ethlambda_types::beacon::config::Config;
 use ethlambda_types::beacon::containers::{BeaconState, SignedBeaconBlock};
@@ -53,6 +53,7 @@ pub(super) struct Replay<'a> {
     engine: ExecutionEngine,
     /// One for the whole run, as the node holds one across its imports.
     committees: CommitteeCache,
+    active_balances: ActiveBalanceCache,
     max_blocks: Option<u64>,
     carried: Option<Carried>,
 }
@@ -64,6 +65,7 @@ impl<'a> Replay<'a> {
             config,
             engine: ExecutionEngine::valid(),
             committees: CommitteeCache::default(),
+            active_balances: ActiveBalanceCache::default(),
             max_blocks,
             carried: None,
         }
@@ -166,6 +168,7 @@ impl<'a> Replay<'a> {
                 self.config,
                 &self.engine,
                 &self.committees,
+                &self.active_balances,
             )
             .map_err(|err| {
                 eyre!(

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_block_root_at_slot,
-    get_current_epoch, get_domain, get_previous_epoch, get_randao_mix,
+    ActiveBalanceCache, CommitteeCache, get_beacon_proposer_index, get_block_root,
+    get_block_root_at_slot, get_current_epoch, get_domain, get_previous_epoch, get_randao_mix,
 };
 use super::helpers::electra::{
     get_attesting_indices, get_indexed_attestation, is_valid_indexed_attestation,
@@ -591,6 +591,7 @@ pub fn assemble_block(
         config,
         &ExecutionEngine::valid(),
         &CommitteeCache::default(),
+        &ActiveBalanceCache::default(),
         options,
     )?;
     block.state_root = post.hash_tree_root();
@@ -696,6 +697,7 @@ mod tests {
             &Config::mainnet(),
             &ExecutionEngine::valid(),
             &CommitteeCache::default(),
+            &ActiveBalanceCache::default(),
             BlockProcessingOptions::default(),
         )
         .unwrap();
@@ -759,6 +761,7 @@ mod tests {
             &config,
             &ExecutionEngine::valid(),
             &CommitteeCache::default(),
+            &ActiveBalanceCache::default(),
             skip_randao(),
         )
         .unwrap();

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

@@ -4604,8 +4604,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
@@ -5278,7 +5288,7 @@ mod tests {
     use super::*;
     use crate::beacon::bls;
     use crate::beacon::containers::BeaconBlockHeader;
-    use crate::beacon::helpers::accessors;
+    use crate::beacon::helpers::accessors::{self, ActiveBalanceCache};
     use crate::beacon::helpers::misc::compute_signing_root;
     use crate::beacon::helpers::test_state;
     use crate::beacon::primitives::BlsSignature;
@@ -6440,6 +6450,7 @@ mod tests {
             &config,
             &stf::ExecutionEngine::valid(),
             &trial_committees,
+            &ActiveBalanceCache::default(),
         )
         .expect("the drafted block must transition cleanly");
         let state_root = trial_state.hash_tree_root();

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

@@ -26,7 +26,8 @@ use super::bls;
 use super::config::Config;
 use super::error::{Error, Result, verify};
 use super::helpers::accessors::{
-    CommitteeCache, get_beacon_proposer_index, get_current_epoch, get_randao_mix,
+    ActiveBalanceCache, CommitteeCache, get_beacon_proposer_index, get_current_epoch,
+    get_randao_mix,
 };
 use super::helpers::gloas::{
     can_builder_cover_bid, get_indexed_payload_attestation, get_ptc, gloas_state_ref,
@@ -407,6 +408,7 @@ pub fn assemble_gloas_block_on_bid(
         &block,
         config,
         &CommitteeCache::default(),
+        &ActiveBalanceCache::default(),
         options,
     )?;
     block.state_root = post.hash_tree_root();
@@ -532,6 +534,7 @@ pub fn assemble_gloas_block(
         &block,
         config,
         &CommitteeCache::default(),
+        &ActiveBalanceCache::default(),
         options,
     )?;
     block.state_root = post.hash_tree_root();
@@ -747,6 +750,7 @@ pub mod test_support {
             block,
             &config(),
             &CommitteeCache::default(),
+            &ActiveBalanceCache::default(),
             BlockProcessingOptions::default(),
         )
         .unwrap();
@@ -910,6 +914,7 @@ mod gloas_block_production_tests {
             &config(),
             &ExecutionEngine::valid(),
             &CommitteeCache::default(),
+            &ActiveBalanceCache::default(),
             BlockProcessingOptions::default(),
         )
         .unwrap();
@@ -1042,6 +1047,7 @@ mod gloas_block_production_tests {
             &produced.block,
             &config(),
             &CommitteeCache::default(),
+            &ActiveBalanceCache::default(),
             skip,
         )
         .unwrap();