feat: electra churn

src/modules/ejector/ejector.py

@@ -1,14 +1,13 @@
 import logging
 import math
-from functools import reduce
 
 from web3.exceptions import ContractCustomError
 from web3.types import Wei
 
 from src.constants import (
+    EFFECTIVE_BALANCE_INCREMENT,
     FAR_FUTURE_EPOCH,
     MAX_WITHDRAWALS_PER_PAYLOAD,
-    MIN_ACTIVATION_BALANCE,
     MIN_VALIDATOR_WITHDRAWABILITY_DELAY,
 )
 from src.metrics.prometheus.business import CONTRACT_ON_PAUSE
@@ -29,11 +28,13 @@
 from src.services.exit_order_v2.iterator import ValidatorExitIteratorV2
 from src.services.prediction import RewardsPredictionService
 from src.services.validator_state import LidoValidatorStateService
-from src.types import BlockStamp, EpochNumber, NodeOperatorGlobalIndex, ReferenceBlockStamp
+from src.types import BlockStamp, EpochNumber, Gwei, NodeOperatorGlobalIndex, ReferenceBlockStamp
 from src.utils.cache import global_lru_cache as lru_cache
 from src.utils.validator_state import (
     compute_activation_exit_epoch,
-    compute_exit_churn_limit,
+    get_activation_exit_churn_limit,
+    get_validator_churn_limit,
+    get_max_effective_balance,
     is_active_validator,
     is_fully_withdrawable_validator,
     is_partially_withdrawable_validator,
@@ -117,7 +118,7 @@ def get_validators_to_eject(self, blockstamp: ReferenceBlockStamp) -> list[tuple
         EJECTOR_TO_WITHDRAW_WEI_AMOUNT.set(to_withdraw_amount)
         logger.info({'msg': 'Calculate to withdraw amount.', 'value': to_withdraw_amount})
 
-        expected_balance = self._get_total_expected_balance(0, blockstamp)
+        expected_balance = self._get_total_expected_balance([], blockstamp)
 
         consensus_version = self.consensus_version(blockstamp)
         validators_iterator = iter(self.get_validators_iterator(consensus_version, blockstamp))
@@ -129,8 +130,11 @@ def get_validators_to_eject(self, blockstamp: ReferenceBlockStamp) -> list[tuple
             while expected_balance < to_withdraw_amount:
                 gid, next_validator = next(validators_iterator)
                 validators_to_eject.append((gid, next_validator))
-                validator_to_eject_balance_sum += self._get_predicted_withdrawable_balance(next_validator)
-                expected_balance = self._get_total_expected_balance(len(validators_to_eject), blockstamp) + validator_to_eject_balance_sum
+                validator_to_eject_balance_sum += self.w3.to_wei(self._get_predicted_withdrawable_balance(next_validator), "gwei")
+                expected_balance = (
+                    self._get_total_expected_balance([v for (_, v) in validators_to_eject], blockstamp)
+                    + validator_to_eject_balance_sum
+                )
         except StopIteration:
             pass
 
@@ -149,7 +153,7 @@ def get_validators_to_eject(self, blockstamp: ReferenceBlockStamp) -> list[tuple
 
         return validators_to_eject
 
-    def _get_total_expected_balance(self, vals_to_exit: int, blockstamp: ReferenceBlockStamp):
+    def _get_total_expected_balance(self, vals_to_exit: list[Validator], blockstamp: ReferenceBlockStamp):
         chain_config = self.get_chain_config(blockstamp)
 
         validators_going_to_exit = self.validators_state_service.get_recently_requested_but_not_exited_validators(blockstamp, chain_config)
@@ -165,7 +169,7 @@ def _get_total_expected_balance(self, vals_to_exit: int, blockstamp: ReferenceBl
         rewards_speed_per_epoch = self.prediction_service.get_rewards_per_epoch(blockstamp, chain_config)
         logger.info({'msg': 'Calculate average rewards speed per epoch.', 'value': rewards_speed_per_epoch})
 
-        withdrawal_epoch = self._get_predicted_withdrawable_epoch(blockstamp, len(validators_going_to_exit) + vals_to_exit + 1)
+        withdrawal_epoch = self._get_predicted_withdrawable_epoch(blockstamp, validators_going_to_exit + vals_to_exit)
         logger.info({'msg': 'Withdrawal epoch', 'value': withdrawal_epoch})
         EJECTOR_MAX_WITHDRAWAL_EPOCH.set(withdrawal_epoch)
 
@@ -211,8 +215,8 @@ def _get_withdrawable_lido_validators_balance(self, on_epoch: EpochNumber, block
             )
         )
 
-    def _get_predicted_withdrawable_balance(self, validator: Validator) -> Wei:
-        return self.w3.to_wei(min(int(validator.balance), MIN_ACTIVATION_BALANCE), 'gwei')
+    def _get_predicted_withdrawable_balance(self, validator: Validator) -> Gwei:
+        return Gwei(min(int(validator.balance), get_max_effective_balance(validator)))
 
     @lru_cache(maxsize=1)
     def _get_total_el_balance(self, blockstamp: BlockStamp) -> Wei:
@@ -225,11 +229,23 @@ def _get_total_el_balance(self, blockstamp: BlockStamp) -> Wei:
     def _get_predicted_withdrawable_epoch(
         self,
         blockstamp: ReferenceBlockStamp,
-        validators_to_eject_count: int,
+        validators_to_eject: list[Validator],
     ) -> EpochNumber:
         """
         Returns epoch when all validators in queue and validators_to_eject will be withdrawn.
         """
+        spec = self.w3.cc.get_config_spec()
+
+        if blockstamp.ref_epoch < int(spec.ELECTRA_FORK_EPOCH):
+            return self._get_predicted_withdrawable_epoch_pre_electra(blockstamp, validators_to_eject)
+
+        return self._get_predicted_withdrawable_epoch_post_electra(blockstamp, validators_to_eject)
+
+    def _get_predicted_withdrawable_epoch_pre_electra(
+        self,
+        blockstamp: ReferenceBlockStamp,
+        validators_to_eject: list[Validator],
+    ) -> EpochNumber:
         max_exit_epoch_number, latest_to_exit_validators_count = self._get_latest_exit_epoch(blockstamp)
 
         activation_exit_epoch = compute_activation_exit_epoch(blockstamp.ref_epoch)
@@ -240,10 +256,32 @@ def _get_predicted_withdrawable_epoch(
 
         churn_limit = self._get_churn_limit(blockstamp)
 
-        epochs_required_to_exit_validators = (validators_to_eject_count + latest_to_exit_validators_count) // churn_limit
+        epochs_required_to_exit_validators = (len(validators_to_eject) + 1 + latest_to_exit_validators_count) // churn_limit
 
         return EpochNumber(max_exit_epoch_number + epochs_required_to_exit_validators + MIN_VALIDATOR_WITHDRAWABILITY_DELAY)
 
+    def _get_predicted_withdrawable_epoch_post_electra(
+        self,
+        blockstamp: ReferenceBlockStamp,
+        validators_to_eject: list[Validator],
+    ) -> EpochNumber:
+        per_epoch_churn = get_activation_exit_churn_limit(self._get_total_active_balance(blockstamp))
+        activation_exit_epoch = compute_activation_exit_epoch(blockstamp.ref_epoch)
+        state_view = self.w3.cc.get_state_view(blockstamp.state_root)
+
+        if state_view.earliest_exit_epoch < activation_exit_epoch:
+            earliest_exit_epoch = activation_exit_epoch
+            exit_balance_to_consume = per_epoch_churn
+        else:
+            earliest_exit_epoch = state_view.earliest_exit_epoch
+            exit_balance_to_consume = state_view.exit_balance_to_consume
+
+        exit_balance = sum(self._get_predicted_withdrawable_balance(v) for v in validators_to_eject)
+        balance_to_process = max(0, exit_balance - exit_balance_to_consume)
+        additional_epochs = math.ceil(balance_to_process / per_epoch_churn)
+
+        return EpochNumber(earliest_exit_epoch + additional_epochs + MIN_VALIDATOR_WITHDRAWABILITY_DELAY)
+
     @lru_cache(maxsize=1)
     def _get_latest_exit_epoch(self, blockstamp: ReferenceBlockStamp) -> tuple[EpochNumber, int]:
         """
@@ -307,19 +345,20 @@ def _get_withdrawable_validators(self, blockstamp: ReferenceBlockStamp) -> list[
 
     @lru_cache(maxsize=1)
     def _get_churn_limit(self, blockstamp: ReferenceBlockStamp) -> int:
-        total_active_validators = self._get_total_active_validators(blockstamp)
-        churn_limit = compute_exit_churn_limit(total_active_validators)
+        total_active_validators = len(self._get_active_validators(blockstamp))
+        logger.info({'msg': 'Calculate total active validators.', 'value': total_active_validators})
+        churn_limit = get_validator_churn_limit(total_active_validators)
         logger.info({'msg': 'Calculate churn limit.', 'value': churn_limit})
         return churn_limit
 
-    def _get_total_active_validators(self, blockstamp: ReferenceBlockStamp) -> int:
-        total_active_validators = reduce(
-            lambda total, validator: total + int(is_active_validator(validator, blockstamp.ref_epoch)),
-            self.w3.cc.get_validators(blockstamp),
-            0,
-        )
-        logger.info({'msg': 'Calculate total active validators.', 'value': total_active_validators})
-        return total_active_validators
+    # https://github.com/ethereum/consensus-specs/blob/dev/specs/phase0/beacon-chain.md#get_total_active_balance
+    def _get_total_active_balance(self, blockstamp: ReferenceBlockStamp) -> Gwei:
+        active_validators = self._get_active_validators(blockstamp)
+        return Gwei(max(EFFECTIVE_BALANCE_INCREMENT, sum(int(v.validator.effective_balance) for v in active_validators)))
+
+    @lru_cache(maxsize=1)
+    def _get_active_validators(self, blockstamp: ReferenceBlockStamp) -> list[Validator]:
+        return [v for v in self.w3.cc.get_validators(blockstamp) if is_active_validator(v, blockstamp.ref_epoch)]
 
     def is_main_data_submitted(self, blockstamp: BlockStamp) -> bool:
         processing_state = self._get_processing_state(blockstamp)

src/providers/consensus/client.py

@@ -1,11 +1,12 @@
 from http import HTTPStatus
 from typing import Literal, cast
 
-from json_stream.base import TransientStreamingJSONObject  # type: ignore
+from json_stream.base import TransientAccessException, TransientStreamingJSONObject  # type: ignore
 
 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
 
@@ -51,7 +52,7 @@ class ConsensusClient(HTTPProvider):
     API_GET_SPEC = 'eth/v1/config/spec'
     API_GET_GENESIS = 'eth/v1/beacon/genesis'
 
-    def get_config_spec(self):
+    def get_config_spec(self) -> BeaconSpecResponse:
         """Spec: https://ethereum.github.io/beacon-APIs/#/Config/getSpec"""
         data, _ = self._get(self.API_GET_SPEC)
         if not isinstance(data, dict):
@@ -151,6 +152,25 @@ 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: 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, see TransientStreamingJSONObject.
+            view['slot'] = int(data['slot'])
+            view['exit_balance_to_consume'] = int(data['exit_balance_to_consume'])
+            view['earliest_exit_epoch'] = int(data['earliest_exit_epoch'])
+        except TransientAccessException:
+            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,8 +1,9 @@
 from dataclasses import dataclass
 from enum import Enum
 
-from src.types import BlockHash, BlockRoot, StateRoot
+from src.types import BlockHash, BlockRoot, Gwei, SlotNumber, StateRoot
 from src.utils.dataclass import Nested, FromResponse
+from src.constants import FAR_FUTURE_EPOCH
 
 
 @dataclass
@@ -12,6 +13,7 @@ class BeaconSpecResponse(FromResponse):
     SECONDS_PER_SLOT: str
     DEPOSIT_CONTRACT_ADDRESS: str
     SLOTS_PER_HISTORICAL_ROOT: str
+    ELECTRA_FORK_EPOCH: str = str(FAR_FUTURE_EPOCH)
 
 
 @dataclass
@@ -150,3 +152,13 @@ class SlotAttestationCommittee(FromResponse):
     index: str
     slot: str
     validators: list[str]
+
+
+@dataclass
+class BeaconStateView(Nested, FromResponse):
+    """A view to BeaconState with only the required keys presented"""
+
+    slot: SlotNumber
+    # This fields are new in Electra, so here are default values for backward compatibility.
+    exit_balance_to_consume: Gwei = Gwei(0)
+    earliest_exit_epoch: int = 0

src/utils/blockstamp.py

@@ -31,5 +31,5 @@ def _build_blockstamp_data(
         "state_root": slot_details.message.state_root,
         "block_number": BlockNumber(int(execution_payload.block_number)),
         "block_hash": execution_payload.block_hash,
-        "block_timestamp": Timestamp(int(execution_payload.timestamp))
+        "block_timestamp": Timestamp(int(execution_payload.timestamp)),
     }

src/utils/types.py

@@ -7,3 +7,13 @@ def bytes_to_hex_str(b: bytes) -> HexStr:
 
 def hex_str_to_bytes(hex_str: HexStr) -> bytes:
     return bytes.fromhex(hex_str[2:])
+
+
+def is_4bytes_hex(s: str) -> bool:
+    if not s.startswith("0x"):
+        return False
+
+    try:
+        return len(bytes.fromhex(s[2:])) == 4
+    except ValueError:
+        return False

src/utils/validator_state.py

@@ -1,16 +1,18 @@
 from typing import Sequence
 
 from src.constants import (
-    ETH1_ADDRESS_WITHDRAWAL_PREFIX,
-    SHARD_COMMITTEE_PERIOD,
-    FAR_FUTURE_EPOCH,
-    EFFECTIVE_BALANCE_INCREMENT,
-    MAX_SEED_LOOKAHEAD,
-    MIN_PER_EPOCH_CHURN_LIMIT,
     CHURN_LIMIT_QUOTIENT,
     COMPOUNDING_WITHDRAWAL_PREFIX,
+    EFFECTIVE_BALANCE_INCREMENT,
+    ETH1_ADDRESS_WITHDRAWAL_PREFIX,
+    FAR_FUTURE_EPOCH,
     MAX_EFFECTIVE_BALANCE_ELECTRA,
+    MAX_PER_EPOCH_ACTIVATION_EXIT_CHURN_LIMIT,
+    MAX_SEED_LOOKAHEAD,
     MIN_ACTIVATION_BALANCE,
+    MIN_PER_EPOCH_CHURN_LIMIT,
+    MIN_PER_EPOCH_CHURN_LIMIT_ELECTRA,
+    SHARD_COMMITTEE_PERIOD,
 )
 from src.providers.consensus.types import Validator
 from src.types import EpochNumber, Gwei
@@ -131,10 +133,21 @@ def compute_activation_exit_epoch(ref_epoch: EpochNumber):
     return ref_epoch + 1 + MAX_SEED_LOOKAHEAD
 
 
-def compute_exit_churn_limit(active_validators_count: int):
+def get_validator_churn_limit(active_validators_count: int):
     return max(MIN_PER_EPOCH_CHURN_LIMIT, active_validators_count // CHURN_LIMIT_QUOTIENT)
 
 
+# @see https://github.com/ethereum/consensus-specs/blob/dev/specs/electra/beacon-chain.md#new-get_activation_exit_churn_limit
+def get_activation_exit_churn_limit(total_active_balance: Gwei) -> Gwei:
+    return min(MAX_PER_EPOCH_ACTIVATION_EXIT_CHURN_LIMIT, get_balance_churn_limit(total_active_balance))
+
+
+# @see https://github.com/ethereum/consensus-specs/blob/dev/specs/electra/beacon-chain.md#new-get_balance_churn_limit
+def get_balance_churn_limit(total_active_balance: Gwei) -> Gwei:
+    churn = max(MIN_PER_EPOCH_CHURN_LIMIT_ELECTRA, total_active_balance // CHURN_LIMIT_QUOTIENT)
+    return Gwei(churn - churn % EFFECTIVE_BALANCE_INCREMENT)
+
+
 def get_max_effective_balance(validator: Validator) -> Gwei:
     """
     Get max effective balance for ``validator``.

tests/factory/no_registry.py

@@ -5,7 +5,7 @@
 from faker import Faker
 from pydantic_factories import Use
 
-from src.constants import EFFECTIVE_BALANCE_INCREMENT, FAR_FUTURE_EPOCH, MIN_ACTIVATION_BALANCE
+from src.constants import EFFECTIVE_BALANCE_INCREMENT, FAR_FUTURE_EPOCH, MAX_EFFECTIVE_BALANCE, MIN_ACTIVATION_BALANCE
 from src.providers.consensus.types import Validator, ValidatorState
 from src.providers.keys.types import LidoKey
 from src.types import Gwei
@@ -86,11 +86,12 @@ def build_exit_vals(cls, epoch, **kwargs: Any):
         )
 
     @classmethod
-    def build_with_balance(cls, balance: Gwei, **kwargs: Any):
+    def build_with_balance(cls, balance: float, meb: int = MAX_EFFECTIVE_BALANCE, **kwargs: Any):
         return cls.build(
             balance=balance,
             validator=ValidatorStateFactory.build(
-                effective_balance=min(balance - balance % EFFECTIVE_BALANCE_INCREMENT, MIN_ACTIVATION_BALANCE),
+                effective_balance=min(balance - balance % EFFECTIVE_BALANCE_INCREMENT, meb),
+                withdrawal_credentials="0x01" if meb == MAX_EFFECTIVE_BALANCE else "0x02",
             ),
             **kwargs,
         )