fix: improve performance of getExpectedWithdrawals (#7045)

* fix: improve performance of getExpectedWithdrawals * chore: use isPostElectra variable * chore: check pre-capella
ChainSafe · Sep 4, 2024 · 681bdcd · 681bdcd
1 parent ef99ed7
commit 681bdcd
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 61 deletions.
diff --git a/packages/state-transition/src/block/processWithdrawals.ts b/packages/state-transition/src/block/processWithdrawals.ts
@@ -7,23 +7,26 @@ import {
   MAX_PENDING_PARTIALS_PER_WITHDRAWALS_SWEEP,
   FAR_FUTURE_EPOCH,
   MIN_ACTIVATION_BALANCE,
+  MAX_EFFECTIVE_BALANCE,
 } from "@lodestar/params";
 
 import {toRootHex} from "@lodestar/utils";
 import {CachedBeaconStateCapella, CachedBeaconStateElectra} from "../types.js";
 import {
   decreaseBalance,
   getValidatorMaxEffectiveBalance,
+  hasEth1WithdrawalCredential,
+  hasExecutionWithdrawalCredential,
   isCapellaPayloadHeader,
-  isFullyWithdrawableValidator,
-  isPartiallyWithdrawableValidator,
 } from "../util/index.js";
 
 export function processWithdrawals(
   fork: ForkSeq,
   state: CachedBeaconStateCapella | CachedBeaconStateElectra,
   payload: capella.FullOrBlindedExecutionPayload
 ): void {
+  // partialWithdrawalsCount is withdrawals coming from EL since electra (EIP-7002)
+  // TODO - electra: may switch to executionWithdrawalsCount
   const {withdrawals: expectedWithdrawals, partialWithdrawalsCount} = getExpectedWithdrawals(fork, state);
   const numWithdrawals = expectedWithdrawals.length;
 
@@ -86,16 +89,33 @@ export function getExpectedWithdrawals(
   sampledValidators: number;
   partialWithdrawalsCount: number;
 } {
+  if (fork < ForkSeq.capella) {
+    throw new Error(`getExpectedWithdrawals not supported at forkSeq=${fork} < ForkSeq.capella`);
+  }
+
   const epoch = state.epochCtx.epoch;
   let withdrawalIndex = state.nextWithdrawalIndex;
   const {validators, balances, nextWithdrawalValidatorIndex} = state;
 
   const withdrawals: capella.Withdrawal[] = [];
+  const isPostElectra = fork >= ForkSeq.electra;
 
-  if (fork >= ForkSeq.electra) {
+  if (isPostElectra) {
     const stateElectra = state as CachedBeaconStateElectra;
 
-    for (const withdrawal of stateElectra.pendingPartialWithdrawals.getAllReadonly()) {
+    // MAX_PENDING_PARTIALS_PER_WITHDRAWALS_SWEEP = 8, PENDING_PARTIAL_WITHDRAWALS_LIMIT: 134217728 so we should only call getAllReadonly() if it makes sense
+    // pendingPartialWithdrawals comes from EIP-7002 smart contract where it takes fee so it's more likely than not validator is in correct condition to withdraw
+    // also we may break early if withdrawableEpoch > epoch
+    const allPendingPartialWithdrawals =
+      stateElectra.pendingPartialWithdrawals.length <= MAX_PENDING_PARTIALS_PER_WITHDRAWALS_SWEEP
+        ? stateElectra.pendingPartialWithdrawals.getAllReadonly()
+        : null;
+
+    // EIP-7002: Execution layer triggerable withdrawals
+    for (let i = 0; i < stateElectra.pendingPartialWithdrawals.length; i++) {
+      const withdrawal = allPendingPartialWithdrawals
+        ? allPendingPartialWithdrawals[i]
+        : stateElectra.pendingPartialWithdrawals.getReadonly(i);
       if (withdrawal.withdrawableEpoch > epoch || withdrawals.length === MAX_PENDING_PARTIALS_PER_WITHDRAWALS_SWEEP) {
         break;
       }
@@ -121,6 +141,7 @@ export function getExpectedWithdrawals(
     }
   }
 
+  // partialWithdrawalsCount is withdrawals coming from EL since electra (EIP-7002)
   const partialWithdrawalsCount = withdrawals.length;
   const bound = Math.min(validators.length, MAX_VALIDATORS_PER_WITHDRAWALS_SWEEP);
   let n = 0;
@@ -132,26 +153,36 @@ export function getExpectedWithdrawals(
 
     const validator = validators.getReadonly(validatorIndex);
     const balance = balances.get(validatorIndex);
+    const {withdrawableEpoch, withdrawalCredentials, effectiveBalance} = validator;
+    const hasWithdrawableCredentials = isPostElectra
+      ? hasExecutionWithdrawalCredential(withdrawalCredentials)
+      : hasEth1WithdrawalCredential(withdrawalCredentials);
     // early skip for balance = 0 as its now more likely that validator has exited/slahed with
     // balance zero than not have withdrawal credentials set
-    if (balance === 0) {
+    if (balance === 0 || !hasWithdrawableCredentials) {
       continue;
     }
 
-    if (isFullyWithdrawableValidator(fork, validator, balance, epoch)) {
+    // capella full withdrawal
+    if (withdrawableEpoch <= epoch) {
       withdrawals.push({
         index: withdrawalIndex,
         validatorIndex,
         address: validator.withdrawalCredentials.subarray(12),
         amount: BigInt(balance),
       });
       withdrawalIndex++;
-    } else if (isPartiallyWithdrawableValidator(fork, validator, balance)) {
+    } else if (
+      effectiveBalance ===
+        (isPostElectra ? getValidatorMaxEffectiveBalance(withdrawalCredentials) : MAX_EFFECTIVE_BALANCE) &&
+      balance > effectiveBalance
+    ) {
+      // capella partial withdrawal
       withdrawals.push({
         index: withdrawalIndex,
         validatorIndex,
         address: validator.withdrawalCredentials.subarray(12),
-        amount: BigInt(balance - getValidatorMaxEffectiveBalance(validator.withdrawalCredentials)),
+        amount: BigInt(balance - effectiveBalance),
       });
       withdrawalIndex++;
     }

diff --git a/packages/state-transition/src/util/electra.ts b/packages/state-transition/src/util/electra.ts
@@ -1,17 +1,8 @@
-import {
-  COMPOUNDING_WITHDRAWAL_PREFIX,
-  FAR_FUTURE_EPOCH,
-  ForkSeq,
-  MAX_EFFECTIVE_BALANCE,
-  MIN_ACTIVATION_BALANCE,
-} from "@lodestar/params";
-import {ValidatorIndex, phase0, ssz} from "@lodestar/types";
+import {COMPOUNDING_WITHDRAWAL_PREFIX, FAR_FUTURE_EPOCH, MIN_ACTIVATION_BALANCE} from "@lodestar/params";
+import {ValidatorIndex, ssz} from "@lodestar/types";
 import {CachedBeaconStateElectra} from "../types.js";
-import {getValidatorMaxEffectiveBalance} from "./validator.js";
 import {hasEth1WithdrawalCredential} from "./capella.js";
 
-type ValidatorInfo = Pick<phase0.Validator, "effectiveBalance" | "withdrawableEpoch" | "withdrawalCredentials">;
-
 export function hasCompoundingWithdrawalCredential(withdrawalCredentials: Uint8Array): boolean {
   return withdrawalCredentials[0] === COMPOUNDING_WITHDRAWAL_PREFIX;
 }
@@ -22,48 +13,6 @@ export function hasExecutionWithdrawalCredential(withdrawalCredentials: Uint8Arr
   );
 }
 
-export function isFullyWithdrawableValidator(
-  fork: ForkSeq,
-  validatorCredential: ValidatorInfo,
-  balance: number,
-  epoch: number
-): boolean {
-  const {withdrawableEpoch, withdrawalCredentials} = validatorCredential;
-
-  if (fork < ForkSeq.capella) {
-    throw new Error(`isFullyWithdrawableValidator not supported at forkSeq=${fork} < ForkSeq.capella`);
-  }
-  const hasWithdrawableCredentials =
-    fork >= ForkSeq.electra
-      ? hasExecutionWithdrawalCredential(withdrawalCredentials)
-      : hasEth1WithdrawalCredential(withdrawalCredentials);
-
-  return hasWithdrawableCredentials && withdrawableEpoch <= epoch && balance > 0;
-}
-
-export function isPartiallyWithdrawableValidator(
-  fork: ForkSeq,
-  validatorCredential: ValidatorInfo,
-  balance: number
-): boolean {
-  const {effectiveBalance, withdrawalCredentials} = validatorCredential;
-
-  if (fork < ForkSeq.capella) {
-    throw new Error(`isPartiallyWithdrawableValidator not supported at forkSeq=${fork} < ForkSeq.capella`);
-  }
-  const hasWithdrawableCredentials =
-    fork >= ForkSeq.electra
-      ? hasExecutionWithdrawalCredential(withdrawalCredentials)
-      : hasEth1WithdrawalCredential(withdrawalCredentials);
-
-  const validatorMaxEffectiveBalance =
-    fork >= ForkSeq.electra ? getValidatorMaxEffectiveBalance(withdrawalCredentials) : MAX_EFFECTIVE_BALANCE;
-  const hasMaxEffectiveBalance = effectiveBalance === validatorMaxEffectiveBalance;
-  const hasExcessBalance = balance > validatorMaxEffectiveBalance;
-
-  return hasWithdrawableCredentials && hasMaxEffectiveBalance && hasExcessBalance;
-}
-
 export function switchToCompoundingValidator(state: CachedBeaconStateElectra, index: ValidatorIndex): void {
   const validator = state.validators.get(index);