feat: sweep computation with electra rules

src/constants.py

@@ -1,39 +1,45 @@
-# https://github.com/ethereum/consensus-specs/blob/dev/specs/phase0/beacon-chain.md#misc
 from src.types import Gwei
 
-FAR_FUTURE_EPOCH = 2 ** 64 - 1
+# https://github.com/ethereum/consensus-specs/blob/dev/specs/phase0/beacon-chain.md#misc
+FAR_FUTURE_EPOCH = 2**64 - 1
 # https://github.com/ethereum/consensus-specs/blob/dev/specs/phase0/beacon-chain.md#time-parameters-1
-MIN_VALIDATOR_WITHDRAWABILITY_DELAY = 2 ** 8
+MIN_VALIDATOR_WITHDRAWABILITY_DELAY = 2**8
 SHARD_COMMITTEE_PERIOD = 256
 MAX_SEED_LOOKAHEAD = 4
 # https://github.com/ethereum/consensus-specs/blob/dev/specs/phase0/beacon-chain.md#state-list-lengths
-EPOCHS_PER_SLASHINGS_VECTOR = 2 ** 13
+EPOCHS_PER_SLASHINGS_VECTOR = 2**13
 # https://github.com/ethereum/consensus-specs/blob/dev/specs/phase0/beacon-chain.md#rewards-and-penalties
 PROPORTIONAL_SLASHING_MULTIPLIER_BELLATRIX = 3
 # https://github.com/ethereum/consensus-specs/blob/dev/specs/phase0/beacon-chain.md#gwei-values
-EFFECTIVE_BALANCE_INCREMENT = 2 ** 0 * 10 ** 9
-MAX_EFFECTIVE_BALANCE = 32 * 10 ** 9
+EFFECTIVE_BALANCE_INCREMENT = 2**0 * 10**9
+MAX_EFFECTIVE_BALANCE = 32 * 10**9
 # https://github.com/ethereum/consensus-specs/blob/dev/specs/capella/beacon-chain.md#execution
-MAX_WITHDRAWALS_PER_PAYLOAD = 2 ** 4
+MAX_WITHDRAWALS_PER_PAYLOAD = 2**4
+# https://github.com/ethereum/consensus-specs/blob/dev/specs/electra/beacon-chain.md#withdrawals-processing
+MAX_VALIDATORS_PER_WITHDRAWALS_SWEEP = 2**14
+MAX_PENDING_PARTIALS_PER_WITHDRAWALS_SWEEP = 2**3
 # https://github.com/ethereum/consensus-specs/blob/dev/specs/phase0/beacon-chain.md#withdrawal-prefixes
 ETH1_ADDRESS_WITHDRAWAL_PREFIX = '0x01'
 # https://github.com/ethereum/consensus-specs/blob/dev/specs/phase0/beacon-chain.md#withdrawal-prefixes
 COMPOUNDING_WITHDRAWAL_PREFIX = '0x02'
 # https://github.com/ethereum/consensus-specs/blob/dev/specs/phase0/beacon-chain.md#validator-cycle
-MIN_PER_EPOCH_CHURN_LIMIT = 2 ** 2
-CHURN_LIMIT_QUOTIENT = 2 ** 16
+MIN_PER_EPOCH_CHURN_LIMIT = 2**2
+CHURN_LIMIT_QUOTIENT = 2**16
+# https://github.com/ethereum/consensus-specs/blob/dev/specs/electra/beacon-chain.md#validator-cycle
+MIN_PER_EPOCH_CHURN_LIMIT_ELECTRA = Gwei(2**7 * 10**9)
+MAX_PER_EPOCH_ACTIVATION_EXIT_CHURN_LIMIT = Gwei(2**8 * 10**9)
 # https://github.com/ethereum/consensus-specs/blob/dev/specs/phase0/beacon-chain.md#time-parameters
 SLOTS_PER_HISTORICAL_ROOT = 8192
 
 # https://github.com/ethereum/consensus-specs/blob/dev/specs/electra/beacon-chain.md#gwei-values
-MIN_ACTIVATION_BALANCE = Gwei(2 ** 5 * 10 ** 9)
-MAX_EFFECTIVE_BALANCE_ELECTRA = Gwei(2 ** 11 * 10 ** 9)
+MIN_ACTIVATION_BALANCE = Gwei(2**5 * 10**9)
+MAX_EFFECTIVE_BALANCE_ELECTRA = Gwei(2**11 * 10**9)
 
 # Local constants
-GWEI_TO_WEI = 10 ** 9
-SHARE_RATE_PRECISION_E27 = 10 ** 27
+GWEI_TO_WEI = 10**9
+SHARE_RATE_PRECISION_E27 = 10**27
 TOTAL_BASIS_POINTS = 10000
 
 MAX_BLOCK_GAS_LIMIT = 30_000_000
 
-UINT64_MAX = 2 ** 64 - 1
+UINT64_MAX = 2**64 - 1

src/modules/ejector/ejector.py

@@ -1,14 +1,15 @@
 import logging
 from functools import reduce
 
-from more_itertools import ilen
 from web3.exceptions import ContractCustomError
 from web3.types import Wei
 
 from src.constants import (
     FAR_FUTURE_EPOCH,
-    MAX_EFFECTIVE_BALANCE,
     MAX_WITHDRAWALS_PER_PAYLOAD,
+    MAX_PENDING_PARTIALS_PER_WITHDRAWALS_SWEEP,
+    MAX_VALIDATORS_PER_WITHDRAWALS_SWEEP,
+    MIN_ACTIVATION_BALANCE,
     MIN_VALIDATOR_WITHDRAWABILITY_DELAY,
 )
 from src.metrics.prometheus.business import CONTRACT_ON_PAUSE
@@ -75,7 +76,7 @@ def __init__(self, w3: Web3):
         self.validators_state_service = LidoValidatorStateService(w3)
 
     def refresh_contracts(self):
-        self.report_contract = self.w3.lido_contracts.validators_exit_bus_oracle
+        self.report_contract = self.w3.lido_contracts.validators_exit_bus_oracle  # type: ignore
 
     def execute_module(self, last_finalized_blockstamp: BlockStamp) -> ModuleExecuteDelay:
         report_blockstamp = self.get_blockstamp_for_report(last_finalized_blockstamp)
@@ -119,7 +120,7 @@ def get_validators_to_eject(self, blockstamp: ReferenceBlockStamp) -> list[tuple
 
         expected_balance = self._get_total_expected_balance(0, blockstamp)
 
-        consensus_version = self.w3.lido_contracts.validators_exit_bus_oracle.get_consensus_version(blockstamp.block_hash)
+        consensus_version = self.consensus_version(blockstamp)
         validators_iterator = iter(self.get_validators_iterator(consensus_version, blockstamp))
 
         validators_to_eject: list[tuple[NodeOperatorGlobalIndex, LidoValidator]] = []
@@ -141,7 +142,7 @@ def get_validators_to_eject(self, blockstamp: ReferenceBlockStamp) -> list[tuple
             'validators_to_eject_count': len(validators_to_eject),
         })
 
-        if consensus_version != 1:
+        if self.consensus_version(blockstamp) != 1:
             forced_validators = validators_iterator.get_remaining_forced_validators()
             if forced_validators:
                 logger.info({'msg': 'Eject forced to exit validators.', 'len': len(forced_validators)})
@@ -203,23 +204,16 @@ def is_reporting_allowed(self, blockstamp: ReferenceBlockStamp) -> bool:
     @lru_cache(maxsize=1)
     def _get_withdrawable_lido_validators_balance(self, on_epoch: EpochNumber, blockstamp: BlockStamp) -> Wei:
         lido_validators = self.w3.lido_validators.get_lido_validators(blockstamp=blockstamp)
-
-        def get_total_withdrawable_balance(balance: Wei, validator: Validator) -> Wei:
-            if is_fully_withdrawable_validator(validator, on_epoch):
-                return Wei(balance + self._get_predicted_withdrawable_balance(validator))
-
-            return balance
-
-        result = reduce(
-            get_total_withdrawable_balance,
-            lido_validators,
-            Wei(0),
+        return Wei(
+            sum(
+                self._get_predicted_withdrawable_balance(v)
+                for v in lido_validators
+                if is_fully_withdrawable_validator(v, on_epoch)
+            )
         )
 
-        return result
-
     def _get_predicted_withdrawable_balance(self, validator: Validator) -> Wei:
-        return self.w3.to_wei(min(int(validator.balance), MAX_EFFECTIVE_BALANCE), 'gwei')
+        return self.w3.to_wei(min(int(validator.balance), MIN_ACTIVATION_BALANCE), 'gwei')
 
     @lru_cache(maxsize=1)
     def _get_total_el_balance(self, blockstamp: BlockStamp) -> Wei:
@@ -286,19 +280,60 @@ def _get_latest_exit_epoch(self, blockstamp: ReferenceBlockStamp) -> tuple[Epoch
     def _get_sweep_delay_in_epochs(self, blockstamp: ReferenceBlockStamp) -> int:
         """Returns amount of epochs that will take to sweep all validators in chain."""
         chain_config = self.get_chain_config(blockstamp)
-        total_withdrawable_validators = self._get_total_withdrawable_validators(blockstamp)
 
-        full_sweep_in_epochs = total_withdrawable_validators / MAX_WITHDRAWALS_PER_PAYLOAD / chain_config.slots_per_epoch
-        return int(full_sweep_in_epochs * self.AVG_EXPECTING_WITHDRAWALS_SWEEP_DURATION_MULTIPLIER)
+        if self.consensus_version(blockstamp) in (1, 2):
+            total_withdrawable_validators = len(self._get_withdrawable_validators(blockstamp))
+            logger.info({'msg': 'Calculate total withdrawable validators.', 'value': total_withdrawable_validators})
+
+            full_sweep_in_epochs = total_withdrawable_validators / MAX_WITHDRAWALS_PER_PAYLOAD / chain_config.slots_per_epoch
+            return int(full_sweep_in_epochs * self.AVG_EXPECTING_WITHDRAWALS_SWEEP_DURATION_MULTIPLIER)
+
+        state_view = self.w3.cc.get_state_view(blockstamp.state_root)
+        partial_withdrawals = list(
+            range(
+                len(state_view.pending_partial_withdrawals) + self.w3.withdrawal_requests.get_queue_len(blockstamp)
+            )
+        )
+        next_sweep_index = int(state_view.next_withdrawal_validator_index)
+
+        withdrawable_indices = set(int(v.index) for v in self._get_withdrawable_validators(blockstamp))
+        total_validators = len(self.w3.cc.get_validators(blockstamp))
+
+        # TODO: Maybe to raise an exception.
+        if not total_validators:
+            return 0
+
+        slots_to_sweep = 0
+
+        while partial_withdrawals or withdrawable_indices:
+            slots_to_sweep += 1
+
+            capacity = MAX_WITHDRAWALS_PER_PAYLOAD
+            capacity -= len(partial_withdrawals[:MAX_PENDING_PARTIALS_PER_WITHDRAWALS_SWEEP])
+            partial_withdrawals = partial_withdrawals[MAX_PENDING_PARTIALS_PER_WITHDRAWALS_SWEEP:]
+
+            if not capacity:
+                continue
+
+            for _ in range(MAX_VALIDATORS_PER_WITHDRAWALS_SWEEP):
+                if next_sweep_index in withdrawable_indices:
+                    withdrawable_indices.remove(next_sweep_index)
+                    capacity -= 1
+
+                next_sweep_index = (next_sweep_index + 1) % total_validators
+
+                if not capacity:
+                    break
 
-    def _get_total_withdrawable_validators(self, blockstamp: ReferenceBlockStamp) -> int:
-        total_withdrawable_validators = ilen(filter(lambda validator: (
-            is_partially_withdrawable_validator(validator) or
-            is_fully_withdrawable_validator(validator, blockstamp.ref_epoch)
-        ), self.w3.cc.get_validators(blockstamp)))
+        # TODO: Do we need to introduce some new multiplier?
+        return (slots_to_sweep - 1) // chain_config.slots_per_epoch + 1
 
-        logger.info({'msg': 'Calculate total withdrawable validators.', 'value': total_withdrawable_validators})
-        return total_withdrawable_validators
+    def _get_withdrawable_validators(self, blockstamp: ReferenceBlockStamp) -> list[Validator]:
+        return [
+            v
+            for v in self.w3.cc.get_validators(blockstamp)
+            if is_partially_withdrawable_validator(v) or is_fully_withdrawable_validator(v, blockstamp.ref_epoch)
+        ]
 
     @lru_cache(maxsize=1)
     def _get_churn_limit(self, blockstamp: ReferenceBlockStamp) -> int:

src/modules/submodules/consensus.py

@@ -87,6 +87,10 @@ def _get_consensus_contract_members(self, blockstamp: BlockStamp):
         consensus_contract = self._get_consensus_contract(blockstamp)
         return consensus_contract.get_members(blockstamp.block_hash)
 
+    @lru_cache(maxsize=1)
+    def consensus_version(self, blockstamp: BlockStamp):
+        return self.report_contract.get_consensus_version(blockstamp.block_hash)
+
     @lru_cache(maxsize=1)
     def get_chain_config(self, blockstamp: BlockStamp) -> ChainConfig:
         consensus_contract = self._get_consensus_contract(blockstamp)

src/providers/consensus/client.py

@@ -6,6 +6,7 @@
 from src.metrics.logging import logging
 from src.metrics.prometheus.basic import CL_REQUESTS_DURATION
 from src.providers.consensus.types import (
+    BeaconStateView,
     BlockDetailsResponse,
     BlockHeaderFullResponse,
     BlockHeaderResponseData,
@@ -16,7 +17,7 @@
     SlotAttestationCommittee, BlockAttestation,
 )
 from src.providers.http_provider import HTTPProvider, NotOkResponse
-from src.types import BlockRoot, BlockStamp, SlotNumber, EpochNumber
+from src.types import BlockRoot, BlockStamp, SlotNumber, EpochNumber, StateRoot
 from src.utils.dataclass import list_of_dataclasses
 from src.utils.cache import global_lru_cache as lru_cache
 
@@ -151,6 +152,28 @@ def get_state_block_roots(self, state_id: SlotNumber) -> list[BlockRoot]:
             ))
         return list(streamed_json['data']['block_roots'])
 
+    @lru_cache(maxsize=1)
+    def get_state_view(self, state_id: LiteralState| SlotNumber | StateRoot) -> BeaconStateView:
+        """Spec: https://ethereum.github.io/beacon-APIs/#/Debug/getStateV2"""
+        streamed_json = cast(TransientStreamingJSONObject, self._get(
+                self.API_GET_STATE,
+                path_params=(state_id,),
+                stream=True,
+            ))
+        view = {}
+        data = streamed_json['data']
+        try:
+            # NOTE: Keep in mind: the order is important.
+            view['slot'] = data['slot']
+            view['next_withdrawal_validator_index'] = data['next_withdrawal_validator_index']
+            view['deposit_balance_to_consume'] = data['deposit_balance_to_consume']
+            view['exit_balance_to_consume'] = data['exit_balance_to_consume']
+            view['earliest_exit_epoch'] = data['earliest_exit_epoch']
+            view['pending_partial_withdrawals'] = data['pending_partial_withdrawals'].persistent()
+        except KeyError:
+            pass
+        return BeaconStateView.from_response(**view)
+
     @lru_cache(maxsize=1)
     def get_validators(self, blockstamp: BlockStamp) -> list[Validator]:
         """Spec: https://ethereum.github.io/beacon-APIs/#/Beacon/getStateValidators"""

src/providers/consensus/types.py

@@ -1,6 +1,7 @@
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from enum import Enum
 
+from src.constants import FAR_FUTURE_EPOCH
 from src.types import BlockHash, BlockRoot, StateRoot
 from src.utils.dataclass import Nested, FromResponse
 
@@ -150,3 +151,23 @@ class SlotAttestationCommittee(FromResponse):
     index: str
     slot: str
     validators: list[str]
+
+
+@dataclass
+class PendingPartialWithdrawal:
+    index: str # ValidatorIndex
+    amount: str
+    withdrawable_epoch: str
+
+
+@dataclass
+class BeaconStateView(Nested, FromResponse):
+    """A view to BeaconState with only the required keys presented"""
+
+    slot: str
+    next_withdrawal_validator_index: str
+    # This fields are new in Electra, so here are default values for backward compatibility.
+    deposit_balance_to_consume: str = "0"
+    exit_balance_to_consume: str = "0"
+    earliest_exit_epoch: str = str(FAR_FUTURE_EPOCH) # XXX: Maybe we need something else here.
+    pending_partial_withdrawals: list[PendingPartialWithdrawal] = field(default_factory=list)