refactor: update process withdrawals

2026-01-09 15:48:08 -05:00 · 2025-12-16 14:54:34 +01:00
parent 3bf4734ba9
commit 57f37790c1
3 changed files with 76 additions and 70 deletions
--- a/packages/state-transition/src/block/processWithdrawals.ts
+++ b/packages/state-transition/src/block/processWithdrawals.ts
@@ -34,6 +34,7 @@ export function processWithdrawals(
  // processedBuilderWithdrawalsCount is withdrawals coming from builder payment since gloas (EIP-7732)
  const {
    withdrawals: expectedWithdrawals,
+    processedValidatorsSweepCount,
    processedPartialWithdrawalsCount,
    processedBuilderWithdrawalsCount,
  } = getExpectedWithdrawals(fork, state);
@@ -105,15 +106,8 @@ export function processWithdrawals(
  }

  // Update the nextWithdrawalValidatorIndex
-  if (latestWithdrawal && expectedWithdrawals.length === MAX_WITHDRAWALS_PER_PAYLOAD) {
-    // All slots filled, nextWithdrawalValidatorIndex should be validatorIndex having next turn
-    state.nextWithdrawalValidatorIndex = (latestWithdrawal.validatorIndex + 1) % state.validators.length;
-  } else {
-    // expected withdrawals came up short in the bound, so we move nextWithdrawalValidatorIndex to
-    // the next post the bound
-    state.nextWithdrawalValidatorIndex =
-      (state.nextWithdrawalValidatorIndex + MAX_VALIDATORS_PER_WITHDRAWALS_SWEEP) % state.validators.length;
-  }
+  const nextIndex = state.nextWithdrawalValidatorIndex + processedValidatorsSweepCount;
+  state.nextWithdrawalValidatorIndex = nextIndex % state.validators.length;
 }

 export function getExpectedWithdrawals(
@@ -121,7 +115,7 @@ export function getExpectedWithdrawals(
  state: CachedBeaconStateCapella | CachedBeaconStateElectra | CachedBeaconStateGloas
 ): {
  withdrawals: capella.Withdrawal[];
-  sampledValidators: number;
+  processedValidatorsSweepCount: number;
  processedPartialWithdrawalsCount: number;
  processedBuilderWithdrawalsCount: number;
 } {
@@ -217,6 +211,7 @@ export function getExpectedWithdrawals(
      const totalWithdrawn = withdrawnBalances.getOrDefault(withdrawal.validatorIndex);
      const balance = state.balances.get(withdrawal.validatorIndex) - totalWithdrawn;

+      // is_eligible_for_partial_withdrawals
      if (
        validator.exitEpoch === FAR_FUTURE_EPOCH &&
        validator.effectiveBalance >= MIN_ACTIVATION_BALANCE &&
@@ -238,31 +233,36 @@ export function getExpectedWithdrawals(
    }
  }

-  const withdrawalBound = Math.min(validators.length, MAX_VALIDATORS_PER_WITHDRAWALS_SWEEP);
-  let n = 0;
-  // Just run a bounded loop max iterating over all withdrawals
-  // however breaks out once we have MAX_WITHDRAWALS_PER_PAYLOAD
-  for (n = 0; n < withdrawalBound; n++) {
-    // Get next validator in turn
-    const validatorIndex = (nextWithdrawalValidatorIndex + n) % validators.length;
+  // get_validators_sweep_withdrawals
+  const validatorsLimit = Math.min(validators.length, MAX_VALIDATORS_PER_WITHDRAWALS_SWEEP);
+  const withdrawalsLimit = MAX_WITHDRAWALS_PER_PAYLOAD;
+
+  let processedValidatorsSweepCount = 0;
+  let validatorIndex = nextWithdrawalValidatorIndex;
+
+  for (let i = 0; i < validatorsLimit; i++) {
+    if (withdrawals.length === withdrawalsLimit) {
+      break;
+    }

    const validator = validators.getReadonly(validatorIndex);
    const withdrawnBalance = withdrawnBalances.getOrDefault(validatorIndex);
-    const balance = isPostElectra
-      ? // Deduct partially withdrawn balance already queued above
-        balances.get(validatorIndex) - withdrawnBalance
-      : balances.get(validatorIndex);
+    // get_balance_after_withdrawals
+    const balance = balances.get(validatorIndex) - withdrawnBalance;
    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/slashed with
-    // balance zero than not have withdrawal credentials set
+
+    // Early skip for balance = 0 as it's now more likely that validator has exited/slashed with
+    // balance zero than not having withdrawal credentials set
    if (balance === 0 || !hasWithdrawableCredentials) {
+      validatorIndex = (validatorIndex + 1) % validators.length;
+      processedValidatorsSweepCount++;
      continue;
    }

-    // capella full withdrawal
+    // is_fully_withdrawable_validator
    if (withdrawableEpoch <= epoch) {
      withdrawals.push({
        index: withdrawalIndex,
@@ -272,27 +272,31 @@ export function getExpectedWithdrawals(
      });
      withdrawalIndex++;
      withdrawnBalances.set(validatorIndex, withdrawnBalance + balance);
-    } else if (
-      effectiveBalance === (isPostElectra ? getMaxEffectiveBalance(withdrawalCredentials) : MAX_EFFECTIVE_BALANCE) &&
-      balance > effectiveBalance
-    ) {
-      // capella partial withdrawal
-      const partialAmount = balance - effectiveBalance;
-      withdrawals.push({
-        index: withdrawalIndex,
-        validatorIndex,
-        address: validator.withdrawalCredentials.subarray(12),
-        amount: BigInt(partialAmount),
-      });
-      withdrawalIndex++;
-      withdrawnBalances.set(validatorIndex, withdrawnBalance + partialAmount);
+    }
+    // is_partially_withdrawable_validator
+    else {
+      const maxEffectiveBalance = isPostElectra ? getMaxEffectiveBalance(withdrawalCredentials) : MAX_EFFECTIVE_BALANCE;
+      if (effectiveBalance === maxEffectiveBalance && balance > maxEffectiveBalance) {
+        const withdrawableBalance = balance - maxEffectiveBalance;
+        withdrawals.push({
+          index: withdrawalIndex,
+          validatorIndex,
+          address: validator.withdrawalCredentials.subarray(12),
+          amount: BigInt(withdrawableBalance),
+        });
+        withdrawalIndex++;
+        withdrawnBalances.set(validatorIndex, withdrawnBalance + withdrawableBalance);
+      }
    }

-    // Break if we have enough to pack the block
-    if (withdrawals.length >= MAX_WITHDRAWALS_PER_PAYLOAD) {
-      break;
-    }
+    validatorIndex = (validatorIndex + 1) % validators.length;
+    processedValidatorsSweepCount++;
  }

-  return {withdrawals, sampledValidators: n, processedPartialWithdrawalsCount, processedBuilderWithdrawalsCount};
+  return {
+    withdrawals,
+    processedValidatorsSweepCount,
+    processedPartialWithdrawalsCount,
+    processedBuilderWithdrawalsCount,
+  };
 }
--- a/packages/state-transition/test/perf/block/processWithdrawals.test.ts
+++ b/packages/state-transition/test/perf/block/processWithdrawals.test.ts
@@ -21,25 +21,26 @@ describe("getExpectedWithdrawals", () => {
  // blsCredentialRatio represents ratio of validators not eligible for withdrawals which
  // can approximate these two cases in combined manner:
  //  - because of credentials not enabled
-  //  - or they were full withdrawan and zero balance
-  const testCases: (WithdrawalOpts & {cache: boolean; sampled: number})[] = [
+  //  - or they were full withdrawn and zero balance
+  // Note: sweepCount is +1 compared to old "sampled" since we now count validators processed, not the loop index
+  const testCases: (WithdrawalOpts & {cache: boolean; sweepCount: number})[] = [
    // Best case when every probe results into a withdrawal candidate
-    {excessBalance: 1, eth1Credentials: 1, withdrawable: 0, withdrawn: 0, cache: true, sampled: 15},
+    {excessBalance: 1, eth1Credentials: 1, withdrawable: 0, withdrawn: 0, cache: true, sweepCount: 16},
    // Normal case based on mainnet conditions: mainnet network conditions: 95% reward rate
-    {excessBalance: 0.95, eth1Credentials: 0.1, withdrawable: 0.05, withdrawn: 0, cache: true, sampled: 219},
+    {excessBalance: 0.95, eth1Credentials: 0.1, withdrawable: 0.05, withdrawn: 0, cache: true, sweepCount: 220},
    // Intermediate good case
-    {excessBalance: 0.95, eth1Credentials: 0.3, withdrawable: 0.05, withdrawn: 0, cache: true, sampled: 42},
-    {excessBalance: 0.95, eth1Credentials: 0.7, withdrawable: 0.05, withdrawn: 0, cache: true, sampled: 18},
+    {excessBalance: 0.95, eth1Credentials: 0.3, withdrawable: 0.05, withdrawn: 0, cache: true, sweepCount: 43},
+    {excessBalance: 0.95, eth1Credentials: 0.7, withdrawable: 0.05, withdrawn: 0, cache: true, sweepCount: 19},
    // Intermediate bad case
-    {excessBalance: 0.1, eth1Credentials: 0.1, withdrawable: 0, withdrawn: 0, cache: true, sampled: 1_020},
-    {excessBalance: 0.03, eth1Credentials: 0.03, withdrawable: 0, withdrawn: 0, cache: true, sampled: 11_777},
+    {excessBalance: 0.1, eth1Credentials: 0.1, withdrawable: 0, withdrawn: 0, cache: true, sweepCount: 1_021},
+    {excessBalance: 0.03, eth1Credentials: 0.03, withdrawable: 0, withdrawn: 0, cache: true, sweepCount: 11_778},
    // Expected 141_069 but gets bounded at 16_384
-    {excessBalance: 0.01, eth1Credentials: 0.01, withdrawable: 0, withdrawn: 0, cache: true, sampled: 16_384},
+    {excessBalance: 0.01, eth1Credentials: 0.01, withdrawable: 0, withdrawn: 0, cache: true, sweepCount: 16_384},
    // Worst case: All validators 250_000 need to be probed but get bounded at 16_384
-    {excessBalance: 0, eth1Credentials: 0.0, withdrawable: 0, withdrawn: 0, cache: true, sampled: 16_384},
-    {excessBalance: 0, eth1Credentials: 0.0, withdrawable: 0, withdrawn: 0, cache: false, sampled: 16_384},
-    {excessBalance: 0, eth1Credentials: 1, withdrawable: 0, withdrawn: 0, cache: true, sampled: 16_384},
-    {excessBalance: 0, eth1Credentials: 1, withdrawable: 0, withdrawn: 0, cache: false, sampled: 16_384},
+    {excessBalance: 0, eth1Credentials: 0.0, withdrawable: 0, withdrawn: 0, cache: true, sweepCount: 16_384},
+    {excessBalance: 0, eth1Credentials: 0.0, withdrawable: 0, withdrawn: 0, cache: false, sweepCount: 16_384},
+    {excessBalance: 0, eth1Credentials: 1, withdrawable: 0, withdrawn: 0, cache: true, sweepCount: 16_384},
+    {excessBalance: 0, eth1Credentials: 1, withdrawable: 0, withdrawn: 0, cache: false, sweepCount: 16_384},
  ];

  for (const opts of testCases) {
@@ -49,7 +50,7 @@ describe("getExpectedWithdrawals", () => {
      `we:${opts.withdrawable}`,
      `wn:${opts.withdrawn}`,
      opts.cache ? null : "nocache",
-      `smpl:${opts.sampled}`,
+      `swp:${opts.sweepCount}`,
    ]
      .filter((str) => str)
      .join(",");
@@ -70,9 +71,9 @@ describe("getExpectedWithdrawals", () => {
        return opts.cache ? state : state.clone(true);
      },
      fn: (state) => {
-        const {sampledValidators} = getExpectedWithdrawals(ForkSeq.capella, state); // TODO Electra: Do test for electra
-        if (sampledValidators !== opts.sampled) {
-          throw Error(`Wrong sampledValidators ${sampledValidators} != ${opts.sampled}`);
+        const {processedValidatorsSweepCount} = getExpectedWithdrawals(ForkSeq.capella, state); // TODO Electra: Do test for electra
+        if (processedValidatorsSweepCount !== opts.sweepCount) {
+          throw Error(`Wrong processedValidatorsSweepCount ${processedValidatorsSweepCount} != ${opts.sweepCount}`);
        }
      },
    });
--- a/packages/state-transition/test/unit/block/processWithdrawals.test.ts
+++ b/packages/state-transition/test/unit/block/processWithdrawals.test.ts
@@ -8,20 +8,21 @@ import {WithdrawalOpts, getExpectedWithdrawalsTestData} from "../../utils/capell
 describe("getExpectedWithdrawals", () => {
  const vc = numValidators;

-  const testCases: (WithdrawalOpts & {withdrawals: number; sampled: number})[] = [
+  const testCases: (WithdrawalOpts & {withdrawals: number; sweepCount: number})[] = [
    // Best case when every probe results into a withdrawal candidate
-    {excessBalance: 1, eth1Credentials: 1, withdrawable: 0, withdrawn: 0, withdrawals: 16, sampled: 15},
+    // Note: sweepCount is +1 compared to old "sampled" since we now count validators processed, not the loop index
+    {excessBalance: 1, eth1Credentials: 1, withdrawable: 0, withdrawn: 0, withdrawals: 16, sweepCount: 16},
    // Normal case based on mainnet conditions: mainnet network conditions: 95% reward rate
-    {excessBalance: 0.95, eth1Credentials: 0.1, withdrawable: 0.05, withdrawn: 0, withdrawals: 16, sampled: 219},
+    {excessBalance: 0.95, eth1Credentials: 0.1, withdrawable: 0.05, withdrawn: 0, withdrawals: 16, sweepCount: 220},
    // Intermediate good case
-    {excessBalance: 0.95, eth1Credentials: 0.3, withdrawable: 0.05, withdrawn: 0, withdrawals: 16, sampled: 42},
-    {excessBalance: 0.95, eth1Credentials: 0.7, withdrawable: 0.05, withdrawn: 0, withdrawals: 16, sampled: 18},
+    {excessBalance: 0.95, eth1Credentials: 0.3, withdrawable: 0.05, withdrawn: 0, withdrawals: 16, sweepCount: 43},
+    {excessBalance: 0.95, eth1Credentials: 0.7, withdrawable: 0.05, withdrawn: 0, withdrawals: 16, sweepCount: 19},
    // Intermediate bad case
-    {excessBalance: 0.1, eth1Credentials: 0.1, withdrawable: 0, withdrawn: 0, withdrawals: 16, sampled: 1020},
+    {excessBalance: 0.1, eth1Credentials: 0.1, withdrawable: 0, withdrawn: 0, withdrawals: 16, sweepCount: 1021},
    // Expected 141069 but gets bounded by 16384
-    {excessBalance: 0.01, eth1Credentials: 0.01, withdrawable: 0, withdrawn: 0, withdrawals: 2, sampled: 16384},
+    {excessBalance: 0.01, eth1Credentials: 0.01, withdrawable: 0, withdrawn: 0, withdrawals: 2, sweepCount: 16384},
    // Expected 250000 but gets bounded by 16384
-    {excessBalance: 0, eth1Credentials: 0.0, withdrawable: 0, withdrawn: 0, withdrawals: 0, sampled: 16384},
+    {excessBalance: 0, eth1Credentials: 0.0, withdrawable: 0, withdrawn: 0, withdrawals: 0, sweepCount: 16384},
  ];

  for (const opts of testCases) {
@@ -39,8 +40,8 @@ describe("getExpectedWithdrawals", () => {

    // TODO Electra: Add test for electra
    it(`getExpectedWithdrawals ${vc} ${caseID}`, () => {
-      const {sampledValidators, withdrawals} = getExpectedWithdrawals(ForkSeq.capella, state.value);
-      expect(sampledValidators).toBe(opts.sampled);
+      const {processedValidatorsSweepCount, withdrawals} = getExpectedWithdrawals(ForkSeq.capella, state.value);
+      expect(processedValidatorsSweepCount).toBe(opts.sweepCount);
      expect(withdrawals.length).toBe(opts.withdrawals);
    });
  }