Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix: improve performance of getExpectedWithdrawals #7045

Merged
merged 3 commits into from
Sep 4, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
47 changes: 39 additions & 8 deletions packages/state-transition/src/block/processWithdrawals.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -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;
Comment on lines +109 to +112
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
const allPendingPartialWithdrawals =
stateElectra.pendingPartialWithdrawals.length <= MAX_PENDING_PARTIALS_PER_WITHDRAWALS_SWEEP
? stateElectra.pendingPartialWithdrawals.getAllReadonly()
: null;
const numPendingPartialWithdrawals = stateElectra.pendingPartialWithdrawals.length;
const allPendingPartialWithdrawals =
numPendingPartialWithdrawals <= MAX_PENDING_PARTIALS_PER_WITHDRAWALS_SWEEP
? stateElectra.pendingPartialWithdrawals.getAllReadonly()
: Array.from(
{ length: numPendingPartialWithdrawals },
(_, i) => stateElectra.pendingPartialWithdrawals.getReadonly(i)
);

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

See if this makes sense. I think it's cleaner this way cuz we don't need to set allPendingPartialWithdrawals to null, and no more conditional assignment for withdrawal in the for loop

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

that'd actually make it worse than the initial version because getAllReadonly() is way more performant than calling getReadonly(i) n times especially if n is big


// 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;
}
Expand All @@ -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;
Expand All @@ -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++;
}
Expand Down
55 changes: 2 additions & 53 deletions packages/state-transition/src/util/electra.ts
Original file line number Diff line number Diff line change
@@ -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;
}
Expand All @@ -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);

Expand Down
Loading