feat: modified sweep computation

src/modules/ejector/ejector.py

@@ -1,4 +1,5 @@
 import logging
+import math
 from functools import reduce
 
 from web3.exceptions import ContractCustomError
@@ -11,14 +12,14 @@
     MIN_VALIDATOR_WITHDRAWABILITY_DELAY,
 )
 from src.metrics.prometheus.business import CONTRACT_ON_PAUSE
+from src.metrics.prometheus.duration_meter import duration_meter
 from src.metrics.prometheus.ejector import (
-    EJECTOR_VALIDATORS_COUNT_TO_EJECT,
-    EJECTOR_TO_WITHDRAW_WEI_AMOUNT,
     EJECTOR_MAX_WITHDRAWAL_EPOCH,
+    EJECTOR_TO_WITHDRAW_WEI_AMOUNT,
+    EJECTOR_VALIDATORS_COUNT_TO_EJECT,
 )
-from src.metrics.prometheus.duration_meter import duration_meter
 from src.modules.ejector.data_encode import encode_data
-from src.modules.ejector.types import ReportData, EjectorProcessingState
+from src.modules.ejector.types import EjectorProcessingState, ReportData
 from src.modules.submodules.consensus import ConsensusModule, InitialEpochIsYetToArriveRevert
 from src.modules.submodules.oracle_module import BaseModule, ModuleExecuteDelay
 from src.modules.submodules.types import ZERO_HASH
@@ -28,19 +29,20 @@
 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, ReferenceBlockStamp, NodeOperatorGlobalIndex
+from src.types import BlockStamp, EpochNumber, 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_max_effective_balance,
+    has_execution_withdrawal_credential,
     is_active_validator,
     is_fully_withdrawable_validator,
     is_partially_withdrawable_validator,
-    compute_activation_exit_epoch,
-    compute_exit_churn_limit,
 )
 from src.web3py.extensions.lido_validators import LidoValidator
 from src.web3py.types import Web3
 
-
 logger = logging.getLogger(__name__)
 
 
@@ -61,6 +63,7 @@ class Ejector(BaseModule, ConsensusModule):
 
     3. Decode lido validators into bytes and send report transaction
     """
+
     COMPATIBLE_ONCHAIN_VERSIONS = [(1, 1), (1, 2)]
 
     AVG_EXPECTING_WITHDRAWALS_SWEEP_DURATION_MULTIPLIER = 0.5
@@ -286,7 +289,16 @@ def _get_sweep_delay_in_epochs(self, blockstamp: ReferenceBlockStamp) -> int:
             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)
 
-        raise NotImplementedError
+        # This version is intended for use with Pectra, but we do not currently take into account pending withdrawal
+        # requests. It would require a large amount of pending withdrawal requests to significantly impact sweep
+        # duration. Roughly every 512 requests adds one more epoch to sweep duration in the current state.
+        # On the other side, to consider pending withdrawals it is necessary to fetch the beacon state and query the
+        # EIP-7002 predeployed contract, which adds complexity with limited improvement for predictions.
+        total_withdrawable_validators = len(self._get_expected_withdrawable_validators(blockstamp))
+        logger.info({'msg': 'Calculate total withdrawable validators.', 'value': total_withdrawable_validators})
+        slots_to_sweep = math.ceil(total_withdrawable_validators / MAX_WITHDRAWALS_PER_PAYLOAD)
+        full_sweep_in_epochs = math.ceil(slots_to_sweep / chain_config.slots_per_epoch)
+        return math.ceil(full_sweep_in_epochs * self.AVG_EXPECTING_WITHDRAWALS_SWEEP_DURATION_MULTIPLIER)
 
     def _get_withdrawable_validators(self, blockstamp: ReferenceBlockStamp) -> list[Validator]:
         return [
@@ -295,6 +307,15 @@ def _get_withdrawable_validators(self, blockstamp: ReferenceBlockStamp) -> list[
             if is_partially_withdrawable_validator(v) or is_fully_withdrawable_validator(v, blockstamp.ref_epoch)
         ]
 
+    def _get_expected_withdrawable_validators(self, blockstamp: ReferenceBlockStamp) -> list[Validator]:
+        def is_withdrawing_validator(v: Validator) -> bool:
+            # We assume a recently swept or appeared validator will get enough balance waiting for the next sweep cycle.
+            has_enough_balance_to_sweep = int(v.balance) >= get_max_effective_balance(v)
+            may_be_swept = has_execution_withdrawal_credential(v) and has_enough_balance_to_sweep
+            return may_be_swept or is_fully_withdrawable_validator(v, blockstamp.ref_epoch)
+
+        return [v for v in self.w3.cc.get_validators(blockstamp) if is_withdrawing_validator(v)]
+
     @lru_cache(maxsize=1)
     def _get_churn_limit(self, blockstamp: ReferenceBlockStamp) -> int:
         total_active_validators = self._get_total_active_validators(blockstamp)

tests/factory/no_registry.py

@@ -5,18 +5,20 @@
 from faker import Faker
 from pydantic_factories import Use
 
-from src.constants import FAR_FUTURE_EPOCH
+from src.constants import EFFECTIVE_BALANCE_INCREMENT, FAR_FUTURE_EPOCH, MIN_ACTIVATION_BALANCE
 from src.providers.consensus.types import Validator, ValidatorState
 from src.providers.keys.types import LidoKey
+from src.types import Gwei
+from src.web3py.extensions.lido_validators import LidoValidator, NodeOperator, StakingModule
 from tests.factory.web3_factory import Web3Factory
-from src.web3py.extensions.lido_validators import StakingModule, LidoValidator, NodeOperator
 
 faker = Faker()
 
 
 class ValidatorStateFactory(Web3Factory):
     __model__ = ValidatorState
 
+    withdrawal_credentials = "0x01"
     exit_epoch = FAR_FUTURE_EPOCH
 
 
@@ -60,7 +62,7 @@ def build_not_active_vals(cls, epoch, **kwargs: Any):
                 activation_epoch=str(faker.pyint(min_value=epoch + 1, max_value=FAR_FUTURE_EPOCH)),
                 exit_epoch=str(FAR_FUTURE_EPOCH),
             ),
-            **kwargs
+            **kwargs,
         )
 
     @classmethod
@@ -70,7 +72,7 @@ def build_active_vals(cls, epoch, **kwargs: Any):
                 activation_epoch=str(faker.pyint(min_value=0, max_value=epoch - 1)),
                 exit_epoch=str(faker.pyint(min_value=epoch + 1, max_value=FAR_FUTURE_EPOCH)),
             ),
-            **kwargs
+            **kwargs,
         )
 
     @classmethod
@@ -80,7 +82,17 @@ def build_exit_vals(cls, epoch, **kwargs: Any):
                 activation_epoch='0',
                 exit_epoch=str(faker.pyint(min_value=1, max_value=epoch)),
             ),
-            **kwargs
+            **kwargs,
+        )
+
+    @classmethod
+    def build_with_balance(cls, balance: Gwei, **kwargs: Any):
+        return cls.build(
+            balance=balance,
+            validator=ValidatorStateFactory.build(
+                effective_balance=min(balance - balance % EFFECTIVE_BALANCE_INCREMENT, MIN_ACTIVATION_BALANCE),
+            ),
+            **kwargs,
         )
 
 

tests/modules/ejector/test_ejector.py

@@ -12,7 +12,7 @@
 from src.modules.ejector.types import EjectorProcessingState
 from src.modules.submodules.oracle_module import ModuleExecuteDelay
 from src.modules.submodules.types import ChainConfig, CurrentFrame
-from src.types import BlockStamp, ReferenceBlockStamp
+from src.types import BlockStamp, Gwei, ReferenceBlockStamp
 from src.utils import validator_state
 from src.web3py.extensions.contracts import LidoContracts
 from src.web3py.extensions.lido_validators import NodeOperatorId, StakingModuleId
@@ -251,7 +251,7 @@ def test_get_total_active_validators(ejector: Ejector) -> None:
 
 @pytest.mark.unit
 @pytest.mark.usefixtures("consensus_client", "lido_validators")
-def test_get_withdrawable_lido_validators(
+def test_get_withdrawable_lido_validators_balance(
     ejector: Ejector,
     ref_blockstamp: ReferenceBlockStamp,
     monkeypatch: pytest.MonkeyPatch,
@@ -344,8 +344,118 @@ def test_get_sweep_delay_in_epochs_pre_electra(
 def test_get_sweep_delay_in_epochs_post_electra(
     ejector: Ejector,
     chain_config: ChainConfig,
+    monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    pass
+    ejector.get_chain_config = Mock(return_value=chain_config)
+    ejector.consensus_version = Mock(return_value=3)
+    ejector.w3.cc = Mock()
+
+    ejector.w3.cc.get_validators = Mock(return_value=[])
+    delay = ejector._get_sweep_delay_in_epochs(Mock(ref_epoch=0))
+    assert delay == 0, "Unexpected sweep delay in epochs"
+
+    ejector.w3.cc.get_validators = Mock(return_value=[LidoValidatorFactory.build_with_balance(Gwei(31 * 10**9))] * 3)
+    with monkeypatch.context() as m:
+        m.setattr(
+            ejector_module,
+            "is_fully_withdrawable_validator",
+            Mock(return_value=False),
+        )
+        delay = ejector._get_sweep_delay_in_epochs(Mock(ref_epoch=0))
+        assert delay == 0, "Unexpected sweep delay in epochs"
+
+    ejector.w3.cc.get_validators = Mock(
+        return_value=[
+            LidoValidatorFactory.build_with_balance(Gwei(32 * 10**9)),
+            LidoValidatorFactory.build_with_balance(Gwei(33 * 10**9)),
+            LidoValidatorFactory.build_with_balance(Gwei(31 * 10**9)),
+        ],
+    )
+    with monkeypatch.context() as m:
+        m.setattr(
+            ejector_module,
+            "is_fully_withdrawable_validator",
+            Mock(return_value=False),
+        )
+        delay = ejector._get_sweep_delay_in_epochs(Mock(ref_epoch=0))
+        assert delay == 1, "Unexpected sweep delay in epochs"
+
+    ejector.w3.cc.get_validators = Mock(
+        return_value=[
+            LidoValidatorFactory.build_with_balance(Gwei(31 * 10**9)),
+            LidoValidatorFactory.build_with_balance(Gwei(31 * 10**9)),
+            LidoValidatorFactory.build_with_balance(Gwei(31 * 10**9)),
+        ],
+    )
+    with monkeypatch.context() as m:
+        m.setattr(
+            ejector_module,
+            "is_fully_withdrawable_validator",
+            Mock(return_value=True),
+        )
+        delay = ejector._get_sweep_delay_in_epochs(Mock(ref_epoch=0))
+        assert delay == 1, "Unexpected sweep delay in epochs"
+
+    ejector.w3.cc.get_validators = Mock(
+        return_value=[
+            LidoValidatorFactory.build_with_balance(Gwei(32 * 10**9)),
+            LidoValidatorFactory.build_with_balance(Gwei(33 * 10**9)),
+        ]
+        * 513,
+    )
+    with monkeypatch.context() as m:
+        m.setattr(
+            ejector_module,
+            "is_fully_withdrawable_validator",
+            Mock(return_value=False),
+        )
+        delay = ejector._get_sweep_delay_in_epochs(Mock(ref_epoch=0))
+        assert delay == 2, "Unexpected sweep delay in epochs"
+
+
+@pytest.mark.unit
+def test_get_withdrawable_validators(ejector: Ejector, monkeypatch) -> None:
+    ejector.w3.cc = Mock()
+    ejector.w3.cc.get_validators = Mock(
+        return_value=[
+            LidoValidatorFactory.build_with_balance(Gwei(32 * 10**9), index=1),
+            LidoValidatorFactory.build_with_balance(Gwei(33 * 10**9), index=2),
+            LidoValidatorFactory.build_with_balance(Gwei(31 * 10**9), index=3),
+        ],
+    )
+
+    with monkeypatch.context() as m:
+        m.setattr(
+            ejector_module,
+            "is_fully_withdrawable_validator",
+            Mock(return_value=False),
+        )
+        withdrawable = ejector._get_withdrawable_validators(Mock())
+
+    assert [v.index for v in withdrawable] == [2]
+
+
+@pytest.mark.unit
+def test_get_expected_withdrawable_validators(ejector: Ejector, monkeypatch) -> None:
+    ejector.w3.cc = Mock()
+    ejector.w3.cc.get_validators = Mock(
+        return_value=[
+            LidoValidatorFactory.build_with_balance(Gwei(32 * 10**9), index=1),
+            LidoValidatorFactory.build_with_balance(Gwei(33 * 10**9), index=2),
+            LidoValidatorFactory.build_with_balance(Gwei(31 * 10**9), index=3),
+        ],
+    )
+
+    with monkeypatch.context() as m:
+        m.setattr(
+            ejector_module,
+            "is_fully_withdrawable_validator",
+            Mock(return_value=False),
+        )
+        withdrawable = ejector._get_expected_withdrawable_validators(Mock())
+
+    assert [v.index for v in withdrawable] == [1, 2]
+
 
 @pytest.mark.usefixtures("contracts")
 def test_get_total_balance(ejector: Ejector, blockstamp: BlockStamp) -> None: