fix: improve processEffectiveBalanceUpdates

Author: twoeths

packages/beacon-node/src/chain/historicalState/worker.ts

@@ -82,6 +82,10 @@ if (metricsRegister) {
       buckets: [0.05, 0.1, 0.2, 0.5, 1, 1.5],
       labelNames: ["source"],
     }),
+    numEffectiveBalanceUpdates: metricsRegister.gauge({
+      name: "lodestar_historical_state_stfn_num_effective_balance_updates_count",
+      help: "Count of effective balance updates in epoch transition",
+    }),
     preStateBalancesNodesPopulatedMiss: metricsRegister.gauge<{source: StateCloneSource}>({
       name: "lodestar_historical_state_stfn_balances_nodes_populated_miss_total",
       help: "Total count state.balances nodesPopulated is false on stfn",

packages/beacon-node/src/metrics/metrics/lodestar.ts

@@ -332,6 +332,10 @@ export function createLodestarMetrics(
       buckets: [0.05, 0.1, 0.2, 0.5, 1, 1.5],
       labelNames: ["source"],
     }),
+    numEffectiveBalanceUpdates: register.gauge({
+      name: "lodestar_stfn_effective_balance_updates_count",
+      help: "Total count of effective balance updates",
+    }),
     preStateBalancesNodesPopulatedMiss: register.gauge<{source: StateCloneSource}>({
       name: "lodestar_stfn_balances_nodes_populated_miss_total",
       help: "Total count state.balances nodesPopulated is false on stfn",

packages/state-transition/src/cache/epochTransitionCache.ts

@@ -1,4 +1,4 @@
-import {Epoch, ValidatorIndex} from "@lodestar/types";
+import {Epoch, ValidatorIndex, phase0} from "@lodestar/types";
 import {intDiv} from "@lodestar/utils";
 import {EPOCHS_PER_SLASHINGS_VECTOR, FAR_FUTURE_EPOCH, ForkSeq, MIN_ACTIVATION_BALANCE} from "@lodestar/params";
 
@@ -127,6 +127,18 @@ export interface EpochTransitionCache {
 
   flags: number[];
 
+  /**
+   * Validators in the current epoch, should use it for read-only value instead of accessing state.validators directly.
+   * Note that during epoch processing, validators could be updated so need to use it with care.
+   */
+  validators: phase0.Validator[];
+
+  /**
+   * This is for electra only
+   * Validators that're switched to compounding during processPendingConsolidations(), not available in beforeProcessEpoch()
+   */
+  newCompoundingValidators?: Set<ValidatorIndex>;
+
   /**
    * balances array will be populated by processRewardsAndPenalties() and consumed by processEffectiveBalanceUpdates().
    * processRewardsAndPenalties() already has a regular Javascript array of balances.
@@ -481,7 +493,9 @@ export function beforeProcessEpoch(
     proposerIndices,
     inclusionDelays,
     flags,
-
+    validators,
+    // will be assigned in processPendingConsolidations()
+    newCompoundingValidators: undefined,
     // Will be assigned in processRewardsAndPenalties()
     balances: undefined,
   };

packages/state-transition/src/epoch/index.ts

@@ -150,8 +150,9 @@ export function processEpoch(
     const timer = metrics?.epochTransitionStepTime.startTimer({
       step: EpochTransitionStep.processEffectiveBalanceUpdates,
     });
-    processEffectiveBalanceUpdates(fork, state, cache);
+    const numUpdate = processEffectiveBalanceUpdates(fork, state, cache);
     timer?.();
+    metrics?.numEffectiveBalanceUpdates.set(numUpdate);
   }
 
   processSlashingsReset(state, cache);

packages/state-transition/src/epoch/processEffectiveBalanceUpdates.ts

@@ -23,12 +23,14 @@ const TIMELY_TARGET = 1 << TIMELY_TARGET_FLAG_INDEX;
  *
  * - On normal mainnet conditions 0 validators change their effective balance
  * - In case of big innactivity event a medium portion of validators may have their effectiveBalance updated
+ *
+ * Return number of validators updated
  */
 export function processEffectiveBalanceUpdates(
   fork: ForkSeq,
   state: CachedBeaconStateAllForks,
   cache: EpochTransitionCache
-): void {
+): number {
   const HYSTERESIS_INCREMENT = EFFECTIVE_BALANCE_INCREMENT / HYSTERESIS_QUOTIENT;
   const DOWNWARD_THRESHOLD = HYSTERESIS_INCREMENT * HYSTERESIS_DOWNWARD_MULTIPLIER;
   const UPWARD_THRESHOLD = HYSTERESIS_INCREMENT * HYSTERESIS_UPWARD_MULTIPLIER;
@@ -43,34 +45,38 @@ export function processEffectiveBalanceUpdates(
   // and updated in processPendingBalanceDeposits() and processPendingConsolidations()
   // so it's recycled here for performance.
   const balances = cache.balances ?? state.balances.getAll();
+  const currentEpochValidators = cache.validators;
+  const newCompoundingValidators = cache.newCompoundingValidators ?? new Set();
 
+  let numUpdate = 0;
   for (let i = 0, len = balances.length; i < len; i++) {
     const balance = balances[i];
 
     // PERF: It's faster to access to get() every single element (4ms) than to convert to regular array then loop (9ms)
     let effectiveBalanceIncrement = effectiveBalanceIncrements[i];
     let effectiveBalance = effectiveBalanceIncrement * EFFECTIVE_BALANCE_INCREMENT;
-    let effectiveBalanceLimit;
+
+    let effectiveBalanceLimit: number;
+    if (fork < ForkSeq.electra) {
+      effectiveBalanceLimit = MAX_EFFECTIVE_BALANCE;
+    } else {
+      // from electra, effectiveBalanceLimit is per validator
+      const isCompoundingValidator =
+        hasCompoundingWithdrawalCredential(currentEpochValidators[i].withdrawalCredentials) ||
+        newCompoundingValidators.has(i);
+      effectiveBalanceLimit = isCompoundingValidator ? MAX_EFFECTIVE_BALANCE_ELECTRA : MIN_ACTIVATION_BALANCE;
+    }
 
     if (
       // Too big
       effectiveBalance > balance + DOWNWARD_THRESHOLD ||
       // Too small. Check effectiveBalance < MAX_EFFECTIVE_BALANCE to prevent unnecessary updates
-      effectiveBalance + UPWARD_THRESHOLD < balance
+      (effectiveBalance < effectiveBalanceLimit && effectiveBalance + UPWARD_THRESHOLD < balance)
     ) {
       // Update the state tree
       // Should happen rarely, so it's fine to update the tree
       const validator = validators.get(i);
 
-      if (fork < ForkSeq.electra) {
-        effectiveBalanceLimit = MAX_EFFECTIVE_BALANCE;
-      } else {
-        // Electra or after
-        effectiveBalanceLimit = hasCompoundingWithdrawalCredential(validator.withdrawalCredentials)
-          ? MAX_EFFECTIVE_BALANCE_ELECTRA
-          : MIN_ACTIVATION_BALANCE;
-      }
-
       effectiveBalance = Math.min(balance - (balance % EFFECTIVE_BALANCE_INCREMENT), effectiveBalanceLimit);
       validator.effectiveBalance = effectiveBalance;
       // Also update the fast cached version
@@ -95,6 +101,7 @@ export function processEffectiveBalanceUpdates(
 
       effectiveBalanceIncrement = newEffectiveBalanceIncrement;
       effectiveBalanceIncrements[i] = effectiveBalanceIncrement;
+      numUpdate++;
     }
 
     // TODO: Do this in afterEpochTransitionCache, looping a Uint8Array should be very cheap
@@ -105,4 +112,5 @@ export function processEffectiveBalanceUpdates(
   }
 
   cache.nextEpochTotalActiveBalanceByIncrement = nextEpochTotalActiveBalanceByIncrement;
+  return numUpdate;
 }

packages/state-transition/src/epoch/processPendingConsolidations.ts

@@ -1,3 +1,4 @@
+import {ValidatorIndex} from "@lodestar/types";
 import {CachedBeaconStateElectra, EpochTransitionCache} from "../types.js";
 import {decreaseBalance, increaseBalance} from "../util/balance.js";
 import {getActiveBalance} from "../util/validator.js";
@@ -20,6 +21,7 @@ export function processPendingConsolidations(state: CachedBeaconStateElectra, ca
   let nextPendingConsolidation = 0;
   const validators = state.validators;
   const cachedBalances = cache.balances;
+  const newCompoundingValidators = new Set<ValidatorIndex>();
 
   for (const pendingConsolidation of state.pendingConsolidations.getAllReadonly()) {
     const {sourceIndex, targetIndex} = pendingConsolidation;
@@ -35,6 +37,7 @@ export function processPendingConsolidations(state: CachedBeaconStateElectra, ca
     }
     // Churn any target excess active balance of target and raise its max
     switchToCompoundingValidator(state, targetIndex);
+    newCompoundingValidators.add(targetIndex);
     // Move active balance to target. Excess balance is withdrawable.
     const activeBalance = getActiveBalance(state, sourceIndex);
     decreaseBalance(state, sourceIndex, activeBalance);
@@ -47,5 +50,6 @@ export function processPendingConsolidations(state: CachedBeaconStateElectra, ca
     nextPendingConsolidation++;
   }
 
+  cache.newCompoundingValidators = newCompoundingValidators;
   state.pendingConsolidations = state.pendingConsolidations.sliceFrom(nextPendingConsolidation);
 }

packages/state-transition/src/metrics.ts

@@ -11,6 +11,7 @@ export type BeaconStateTransitionMetrics = {
   processBlockTime: Histogram;
   processBlockCommitTime: Histogram;
   stateHashTreeRootTime: Histogram<{source: StateHashTreeRootSource}>;
+  numEffectiveBalanceUpdates: Gauge;
   preStateBalancesNodesPopulatedMiss: Gauge<{source: StateCloneSource}>;
   preStateBalancesNodesPopulatedHit: Gauge<{source: StateCloneSource}>;
   preStateValidatorsNodesPopulatedMiss: Gauge<{source: StateCloneSource}>;

packages/state-transition/test/perf/epoch/epochAltair.test.ts

@@ -141,7 +141,9 @@ function benchmarkAltairEpochSteps(stateOg: LazyValue<CachedBeaconStateAllForks>
   itBench({
     id: `${stateId} - altair processEffectiveBalanceUpdates`,
     beforeEach: () => stateOg.value.clone(),
-    fn: (state) => processEffectiveBalanceUpdates(ForkSeq.altair, state, cache.value),
+    fn: (state) => {
+      processEffectiveBalanceUpdates(ForkSeq.altair, state, cache.value);
+    },
   });
 
   itBench({

packages/state-transition/test/perf/epoch/epochCapella.test.ts

@@ -120,7 +120,9 @@ function benchmarkAltairEpochSteps(stateOg: LazyValue<CachedBeaconStateAllForks>
   itBench({
     id: `${stateId} - capella processEffectiveBalanceUpdates`,
     beforeEach: () => stateOg.value.clone(),
-    fn: (state) => processEffectiveBalanceUpdates(ForkSeq.capella, state, cache.value),
+    fn: (state) => {
+      processEffectiveBalanceUpdates(ForkSeq.capella, state, cache.value);
+    },
   });
 
   itBench({

packages/state-transition/test/perf/epoch/epochPhase0.test.ts

@@ -123,7 +123,9 @@ function benchmarkPhase0EpochSteps(stateOg: LazyValue<CachedBeaconStateAllForks>
   itBench({
     id: `${stateId} - phase0 processEffectiveBalanceUpdates`,
     beforeEach: () => stateOg.value.clone(),
-    fn: (state) => processEffectiveBalanceUpdates(ForkSeq.phase0, state, cache.value),
+    fn: (state) => {
+      processEffectiveBalanceUpdates(ForkSeq.phase0, state, cache.value);
+    },
   });
 
   itBench({