feat: added sweep module, consensus client service, refactoring

Author: Tarens2

package.json

@@ -30,7 +30,7 @@
     "@fastify/static": "^7.0.4",
     "@lido-nestjs/consensus": "^1.5.0",
     "@lido-nestjs/constants": "^5.2.0",
-    "@lido-nestjs/contracts": "^9.3.0",
+    "@lido-nestjs/contracts": "^9.6.0",
     "@lido-nestjs/decorators": "^1.0.0",
     "@lido-nestjs/execution": "^1.9.3",
     "@lido-nestjs/fetch": "^1.4.0",
@@ -46,6 +46,7 @@
     "@nestjs/terminus": "^9.1.4",
     "@nestjs/throttler": "^6.0.0",
     "@sentry/node": "^7.29.0",
+    "@types/stream-json": "^1.7.8",
     "@willsoto/nestjs-prometheus": "^5.1.0",
     "cache-manager": "^5.7.4",
     "class-transformer": "^0.5.1",
@@ -57,7 +58,7 @@
     "rimraf": "^3.0.2",
     "rxjs": "^7.5.2",
     "stream-chain": "^2.2.5",
-    "stream-json": "^1.8.0"
+    "stream-json": "^1.9.1"
   },
   "devDependencies": {
     "@nestjs/cli": "^10.4.4",

src/common/consensus-provider/consensus-client.service.ts

@@ -0,0 +1,21 @@
+import { ConsensusProviderService } from './index';
+import { Injectable } from '@nestjs/common';
+import { processJsonStreamBeaconState } from './utils/process-json-stream-beacon-state';
+import { BeaconState } from './consensus-provider.types';
+
+@Injectable()
+export class ConsensusClientService {
+  private API_GET_STATE = (stateId: string) => `/eth/v2/debug/beacon/states/${stateId}`;
+  constructor(protected readonly consensusService: ConsensusProviderService) {}
+
+  public async isElectraActivated(epoch: number) {
+    const spec = await this.consensusService.getSpec();
+    return epoch >= +spec.data.ELECTRA_FORK_EPOCH;
+  }
+
+  public async getStateStream(stateId: string): Promise<BeaconState> {
+    const stream = await this.consensusService.fetchStream(this.API_GET_STATE(stateId));
+    const result = await processJsonStreamBeaconState(stream);
+    return result as BeaconState;
+  }
+}

src/common/consensus-provider/consensus-provider.module.ts

@@ -2,6 +2,7 @@ import { ConsensusModule } from '@lido-nestjs/consensus';
 import { Global, Module } from '@nestjs/common';
 import { CONSENSUS_POOL_INTERVAL_MS } from './consensus-provider.constants';
 import { ConsensusFetchModule } from './consensus-fetch.module';
+import { ConsensusClientService } from './consensus-client.service';
 
 @Global()
 @Module({
@@ -11,5 +12,7 @@ import { ConsensusFetchModule } from './consensus-fetch.module';
       pollingInterval: CONSENSUS_POOL_INTERVAL_MS,
     }),
   ],
+  exports: [ConsensusClientService],
+  providers: [ConsensusClientService],
 })
 export class ConsensusProviderModule {}

src/common/consensus-provider/consensus-provider.types.ts

@@ -0,0 +1,27 @@
+import { ConsensusMethodResult } from '@lido-nestjs/consensus/dist/interfaces/consensus.interface';
+
+export type ResponseValidatorsData = Awaited<ConsensusMethodResult<'getStateValidators'>>['data'];
+export type IndexedValidator = ResponseValidatorsData[number];
+export type Validator = ResponseValidatorsData[number]['validator'];
+
+/**
+ * Spec reference:
+ * https://github.com/ethereum/consensus-specs/blob/dev/specs/electra/beacon-chain.md#pendingpartialwithdrawal
+ */
+export interface PendingPartialWithdrawal {
+  validator_index: string;
+  amount: string;
+  withdrawable_epoch: string;
+}
+
+/**
+ * Spec reference:
+ * https://github.com/ethereum/consensus-specs/blob/dev/specs/electra/beacon-chain.md#beaconstate
+ * included only used properties
+ */
+export interface BeaconState {
+  slot: string;
+  pending_partial_withdrawals: PendingPartialWithdrawal[];
+  validators: Validator[];
+  balances: string[];
+}

src/common/consensus-provider/utils/process-json-stream-beacon-state.ts

@@ -0,0 +1,33 @@
+import { parser } from 'stream-json';
+import { pick } from 'stream-json/filters/Pick';
+import { streamObject } from 'stream-json/streamers/StreamObject';
+import { chain } from 'stream-chain';
+import { BeaconState } from '../consensus-provider.types';
+
+const keys = ['slot', 'pending_partial_withdrawals', 'validators', 'balances'] as const;
+
+export async function processJsonStreamBeaconState(readableStream) {
+  return new Promise((resolve, reject) => {
+    const pipeline = chain([
+      readableStream, // Incoming ReadableStream
+      parser(), // Parses JSON as a stream
+      pick({ filter: 'data' }),
+      streamObject(), // Streams key-value pairs { key, value }
+    ]);
+
+    const result = {} as BeaconState;
+
+    pipeline.on('data', ({ key, value }) => {
+      console.log('key', key);
+      if (keys.includes(key)) {
+        result[key] = value; // Store key-value pairs in an object
+      }
+    });
+
+    pipeline.on('end', () => {
+      resolve(result); // Resolve with the final object
+    });
+
+    pipeline.on('error', reject);
+  });
+}

src/common/sweep/index.ts

@@ -0,0 +1,3 @@
+export * from './sweep.module';
+export * from './sweep.service';
+export * from './sweep.constants';

src/common/sweep/sweep.constants.ts

@@ -0,0 +1,5 @@
+import { MAX_EFFECTIVE_BALANCE } from '../../waiting-time/waiting-time.constants';
+
+export const MAX_WITHDRAWALS_PER_PAYLOAD = 2 ** 4;
+export const MAX_PENDING_PARTIALS_PER_WITHDRAWALS_SWEEP = 2 ** 3;
+export const MIN_ACTIVATION_BALANCE = MAX_EFFECTIVE_BALANCE;

src/common/sweep/sweep.module.ts

@@ -0,0 +1,13 @@
+import { Module } from '@nestjs/common';
+import { ConsensusProviderModule } from 'common/consensus-provider';
+import { LoggerModule } from 'common/logger';
+import { SweepService } from './sweep.service';
+import { ContractConfigStorageModule } from '../../storage';
+import { GenesisTimeModule } from '../genesis-time';
+
+@Module({
+  imports: [LoggerModule, ConsensusProviderModule, ContractConfigStorageModule, GenesisTimeModule],
+  providers: [SweepService],
+  exports: [SweepService],
+})
+export class SweepModule {}

src/common/sweep/sweep.service.ts

@@ -0,0 +1,147 @@
+import { Inject, Injectable, LoggerService, OnModuleInit } from '@nestjs/common';
+import { LOGGER_PROVIDER } from '@lido-nestjs/logger';
+import { GenesisTimeService, SLOTS_PER_EPOCH } from '../genesis-time';
+import {
+  getMaxEffectiveBalance,
+  isFullyWithdrawableValidator,
+  isPartiallyWithdrawableValidator,
+} from '../../jobs/validators/utils/validator-state-utils';
+import { FAR_FUTURE_EPOCH } from '../../jobs/validators';
+import {
+  MAX_PENDING_PARTIALS_PER_WITHDRAWALS_SWEEP,
+  MAX_WITHDRAWALS_PER_PAYLOAD,
+  MIN_ACTIVATION_BALANCE,
+} from './sweep.constants';
+import { ConsensusClientService } from '../consensus-provider/consensus-client.service';
+import { parseGwei } from '../utils/parse-gwei';
+import { bigNumberMin } from '../utils/big-number-min';
+import { Withdrawal } from './sweep.types';
+import { BeaconState, IndexedValidator, Validator } from '../consensus-provider/consensus-provider.types';
+
+@Injectable()
+export class SweepService implements OnModuleInit {
+  constructor(
+    @Inject(LOGGER_PROVIDER) protected readonly logger: LoggerService,
+    protected readonly consensusClientService: ConsensusClientService,
+    protected readonly genesisService: GenesisTimeService,
+  ) {}
+
+  public async onModuleInit(): Promise<void> {
+    const epoch = this.genesisService.getCurrentEpoch();
+    await this.getSweepDelayInEpochs([], epoch);
+  }
+
+  getConsensusVersion() {
+    // todo add call contract this.vebOracle.getConsensusVersion()
+    return 1;
+  }
+
+  public async getSweepDelayInEpochs(indexedValidators: IndexedValidator[], currentEpoch: number) {
+    const isElectraActivate = await this.consensusClientService.isElectraActivated(currentEpoch);
+    if (this.getConsensusVersion() < 3 || !isElectraActivate) {
+      return this.getSweepDelayInEpochsPreElectra(indexedValidators, currentEpoch);
+    }
+
+    const state = await this.consensusClientService.getStateStream('head');
+    return this.getSweepDelayInEpochsPostElectra(state, indexedValidators);
+  }
+
+  private getSweepDelayInEpochsPreElectra(indexedValidators: IndexedValidator[], epoch: number): number {
+    const totalWithdrawableValidators = this.getWithdrawableValidators(indexedValidators, epoch).length;
+
+    const fullSweepInEpochs = totalWithdrawableValidators / MAX_WITHDRAWALS_PER_PAYLOAD / SLOTS_PER_EPOCH;
+    return Math.floor(fullSweepInEpochs * 0.5);
+  }
+
+  // pre pectra
+  private getWithdrawableValidators(indexedValidators: IndexedValidator[], epoch: number) {
+    return indexedValidators.filter(
+      (v) =>
+        isPartiallyWithdrawableValidator(v.validator, parseGwei(v.balance)) ||
+        isFullyWithdrawableValidator(v.validator, parseGwei(v.balance), epoch),
+    );
+  }
+
+  private getSweepDelayInEpochsPostElectra(state: BeaconState, indexedValidators: IndexedValidator[]): number {
+    const withdrawalsNumberInSweepCycle = this.predictWithdrawalsNumberInSweepCycle(state, indexedValidators);
+    const fullSweepCycleInEpochs = Math.ceil(
+      withdrawalsNumberInSweepCycle / MAX_WITHDRAWALS_PER_PAYLOAD / SLOTS_PER_EPOCH,
+    );
+    return Math.floor(fullSweepCycleInEpochs / 2);
+  }
+
+  private predictWithdrawalsNumberInSweepCycle(state: BeaconState, indexedValidators: IndexedValidator[]): number {
+    const pendingPartialWithdrawals = this.getPendingPartialWithdrawals(state);
+    const validatorsWithdrawals = this.getValidatorsWithdrawals(state, pendingPartialWithdrawals, indexedValidators);
+
+    const pendingPartialWithdrawalsNumber = pendingPartialWithdrawals.length;
+    const validatorsWithdrawalsNumber = validatorsWithdrawals.length;
+
+    const partialWithdrawalsMaxRatio =
+      MAX_PENDING_PARTIALS_PER_WITHDRAWALS_SWEEP /
+      (MAX_WITHDRAWALS_PER_PAYLOAD - MAX_PENDING_PARTIALS_PER_WITHDRAWALS_SWEEP);
+
+    const pendingPartialWithdrawalsMaxNumberInCycle = Math.ceil(
+      validatorsWithdrawalsNumber * partialWithdrawalsMaxRatio,
+    );
+
+    const pendingPartialWithdrawalsNumberInCycle = Math.min(
+      pendingPartialWithdrawalsNumber,
+      pendingPartialWithdrawalsMaxNumberInCycle,
+    );
+
+    return validatorsWithdrawalsNumber + pendingPartialWithdrawalsNumberInCycle;
+  }
+
+  private getPendingPartialWithdrawals(state: BeaconState): Withdrawal[] {
+    const withdrawals: Withdrawal[] = [];
+
+    for (const pendingPartialWithdrawal of state.pending_partial_withdrawals) {
+      const index = pendingPartialWithdrawal.validator_index;
+      const validator: Validator = state.validators[index];
+      const hasSufficientEffectiveBalance = parseGwei(validator.effective_balance).gte(MIN_ACTIVATION_BALANCE);
+      const hasExcessBalance = parseGwei(state.balances[index]).gt(MIN_ACTIVATION_BALANCE);
+
+      if (validator.exit_epoch === FAR_FUTURE_EPOCH.toString() && hasSufficientEffectiveBalance && hasExcessBalance) {
+        const withdrawableBalance = bigNumberMin(
+          parseGwei(state.balances[index]).sub(MIN_ACTIVATION_BALANCE),
+          parseGwei(pendingPartialWithdrawal.amount),
+        );
+        withdrawals.push({ validatorIndex: index, amount: withdrawableBalance });
+      }
+    }
+    return withdrawals;
+  }
+
+  // post pectra
+  getValidatorsWithdrawals(
+    state: BeaconState,
+    partialWithdrawals: Withdrawal[],
+    indexedValidators: IndexedValidator[],
+  ): Withdrawal[] {
+    const epoch = Math.ceil(+state.slot / SLOTS_PER_EPOCH);
+    const withdrawals: Withdrawal[] = [];
+    const partiallyWithdrawnMap: Record<number, number> = {};
+
+    for (const withdrawal of partialWithdrawals) {
+      partiallyWithdrawnMap[withdrawal.validatorIndex] =
+        (partiallyWithdrawnMap[withdrawal.validatorIndex] || 0) + withdrawal.amount;
+    }
+
+    for (const indexedValidator of indexedValidators) {
+      const validatorIndex = indexedValidator.index;
+      const validator = indexedValidator.validator;
+      const partiallyWithdrawnBalance = partiallyWithdrawnMap[validatorIndex] || 0;
+      const balance = parseGwei(state.balances[validatorIndex]).sub(partiallyWithdrawnBalance);
+
+      if (isFullyWithdrawableValidator(validator, balance, epoch)) {
+        withdrawals.push({ validatorIndex, amount: balance });
+      } else if (isPartiallyWithdrawableValidator(validator, balance)) {
+        const maxEffectiveBalance = getMaxEffectiveBalance(validator);
+        withdrawals.push({ validatorIndex, amount: balance.sub(maxEffectiveBalance) });
+      }
+    }
+
+    return withdrawals;
+  }
+}

src/common/sweep/sweep.types.ts

@@ -0,0 +1,6 @@
+import { BigNumber } from '@ethersproject/bignumber';
+
+export interface Withdrawal {
+  validatorIndex: string;
+  amount: BigNumber;
+}

src/common/utils/big-number-min.ts

@@ -0,0 +1,5 @@
+import { BigNumber } from '@ethersproject/bignumber';
+
+export const bigNumberMin = (a: BigNumber, b: BigNumber) => {
+  return a.lt(b) ? a : b;
+};

src/common/utils/parse-gwei-to-big-number.ts renamed to src/common/utils/parse-gwei.ts

@@ -1,6 +1,6 @@
 import { BigNumber } from '@ethersproject/bignumber';
 
-export const parseGweiToWei = (gweiValue: string) => {
+export const parseGwei = (gweiValue: string) => {
   const toWei = BigNumber.from('1000000000');
   return BigNumber.from(gweiValue).mul(toWei);
 };

src/jobs/validators/lido-keys/lido-keys.service.ts

@@ -1,9 +1,9 @@
 import { Inject } from '@nestjs/common';
 import { LOGGER_PROVIDER, LoggerService } from 'common/logger';
+import { IndexedValidator } from 'common/consensus-provider/consensus-provider.types';
 import { ConfigService } from 'common/config';
 import { LidoKey, LidoKeysData } from './lido-keys.types';
 import { KEYS_API_ADDRESS } from './lido-keys.constants';
-import { Validator } from '../validators.types';
 
 export class LidoKeysService {
   constructor(
@@ -19,7 +19,7 @@ export class LidoKeysService {
     return lidoKeys;
   }
 
-  public async getLidoValidatorsByKeys(keys: LidoKey[], validators: Validator[]) {
+  public async getLidoValidatorsByKeys(keys: LidoKey[], validators: IndexedValidator[]) {
     const keysDict = keys.reduce((acc, lidoKey) => {
       acc[lidoKey.key] = true;
       return acc;

src/jobs/validators/utils/get-validator-withdrawal-timestamp.ts

@@ -26,15 +26,15 @@ Examples:
 
 */
 export function getValidatorWithdrawalTimestamp(
-  index: BigNumber,
+  validatorIndex: BigNumber,
   lastWithdrawalValidatorIndex: BigNumber,
   activeValidatorCount: number,
   totalValidatorsCount: number,
 ) {
-  const diff = index.sub(lastWithdrawalValidatorIndex);
+  const diff = validatorIndex.sub(lastWithdrawalValidatorIndex);
   const percentOfActiveValidators = activeValidatorCount / totalValidatorsCount;
   const lengthQueueValidators = diff.lt(0)
-    ? BigNumber.from(totalValidatorsCount).sub(lastWithdrawalValidatorIndex).add(index)
+    ? BigNumber.from(totalValidatorsCount).sub(lastWithdrawalValidatorIndex).add(validatorIndex)
     : diff;
 
   const slots = lengthQueueValidators.div(BigNumber.from(WITHDRAWALS_VALIDATORS_PER_SLOT));