From c3d0395871c992edef1f57ee57840ced992a8ee8 Mon Sep 17 00:00:00 2001 From: vgorkavenko Date: Wed, 22 Jan 2025 16:51:18 +0100 Subject: [PATCH] refactor: hide consensus and fork checks in functions --- src/modules/accounting/accounting.py | 54 ++++++---------- .../bunker_cases/abnormal_cl_rebase.py | 64 ++++++------------- 2 files changed, 39 insertions(+), 79 deletions(-) diff --git a/src/modules/accounting/accounting.py b/src/modules/accounting/accounting.py index 703d64b58..3efacead5 100644 --- a/src/modules/accounting/accounting.py +++ b/src/modules/accounting/accounting.py @@ -210,28 +210,26 @@ def get_updated_modules_stats( return list(module_stats.keys()), list(module_stats.values()) @lru_cache(maxsize=1) - def _get_consensus_lido_state_pre_electra(self, blockstamp: ReferenceBlockStamp) -> tuple[ValidatorsCount, ValidatorsBalance]: - lido_validators = self.w3.lido_validators.get_lido_validators(blockstamp) - - count = len(lido_validators) - total_balance = Gwei(sum(int(validator.balance) for validator in lido_validators)) + def _get_consensus_lido_state(self, blockstamp: ReferenceBlockStamp) -> tuple[ValidatorsCount, ValidatorsBalance]: - logger.info({'msg': 'Calculate lido state on CL. (Validators count, Total balance in gwei)', 'value': (count, total_balance)}) - return ValidatorsCount(count), ValidatorsBalance(total_balance) - - @lru_cache(maxsize=1) - def _get_consensus_lido_state_post_electra(self, blockstamp: ReferenceBlockStamp) -> tuple[ValidatorsCount, ValidatorsBalance]: lido_validators = self.w3.lido_validators.get_lido_validators(blockstamp) + logger.info({'msg': 'Calculate Lido validators count', 'value': len(lido_validators)}) - validators_count = len(lido_validators) - active_balance = sum(int(validator.balance) for validator in lido_validators) - state = self.w3.cc.get_state_view(blockstamp) - pending_deposits = self.w3.lido_validators.calculate_pending_deposits_sum(lido_validators, state.pending_deposits) - total_balance = Gwei(active_balance + pending_deposits) + total_balance = sum(int(validator.balance) for validator in lido_validators) + logger.info({'msg': 'Calculate Lido active balance (in Gwei)', 'value': total_balance}) + + consensus_version = self.w3.lido_contracts.accounting_oracle.get_consensus_version(blockstamp.block_hash) + if consensus_version > 2: + spec = self.w3.cc.get_config_spec() + if blockstamp.ref_epoch >= int(spec.ELECTRA_FORK_EPOCH): + state = self.w3.cc.get_state_view(blockstamp) + pending_deposits = self.w3.lido_validators.calculate_pending_deposits_sum( + lido_validators, state.pending_deposits + ) + logger.info({'msg': 'Calculate Lido pending deposits (in Gwei)', 'value': pending_deposits}) + total_balance += pending_deposits - logger.info( - {'msg': f'Calculate Lido state on CL. {validators_count=}, {active_balance=}, {pending_deposits=}, {total_balance=} (Gwei)'}) - return ValidatorsCount(validators_count), ValidatorsBalance(total_balance) + return ValidatorsCount(len(lido_validators)), ValidatorsBalance(Gwei(total_balance)) def _get_finalization_data(self, blockstamp: ReferenceBlockStamp) -> tuple[FinalizationShareRate, FinalizationBatches]: simulation = self.simulate_full_rebase(blockstamp) @@ -278,15 +276,7 @@ def simulate_rebase_after_report( """ To calculate how much withdrawal request protocol can finalize - needs finalization share rate after this report """ - consensus_version = self.w3.lido_contracts.accounting_oracle.get_consensus_version(blockstamp.block_hash) - if consensus_version in (1, 2): - validators_count, cl_balance = self._get_consensus_lido_state_pre_electra(blockstamp) - else: - spec = self.w3.cc.get_config_spec() - if blockstamp.ref_epoch < int(spec.ELECTRA_FORK_EPOCH): - validators_count, cl_balance = self._get_consensus_lido_state_pre_electra(blockstamp) - else: - validators_count, cl_balance = self._get_consensus_lido_state_post_electra(blockstamp) + validators_count, cl_balance = self._get_consensus_lido_state(blockstamp) chain_conf = self.get_chain_config(blockstamp) @@ -399,15 +389,7 @@ def _calculate_report_latest_version(self, consensus_version: int, blockstamp: R # fetches validators_count, cl_balance, withdrawal_balance, el_vault_balance, shares_to_burn def _calculate_rebase_report(self, blockstamp: ReferenceBlockStamp) -> RebaseReport: - consensus_version = self.w3.lido_contracts.accounting_oracle.get_consensus_version(blockstamp.block_hash) - if consensus_version in (1, 2): - validators_count, cl_balance = self._get_consensus_lido_state_pre_electra(blockstamp) - else: - spec = self.w3.cc.get_config_spec() - if blockstamp.ref_epoch < int(spec.ELECTRA_FORK_EPOCH): - validators_count, cl_balance = self._get_consensus_lido_state_pre_electra(blockstamp) - else: - validators_count, cl_balance = self._get_consensus_lido_state_post_electra(blockstamp) + validators_count, cl_balance = self._get_consensus_lido_state(blockstamp) withdrawal_vault_balance = self.w3.lido_contracts.get_withdrawal_balance(blockstamp) el_rewards_vault_balance = self.w3.lido_contracts.get_el_vault_balance(blockstamp) shares_requested_to_burn = self.get_shares_to_burn(blockstamp) diff --git a/src/services/bunker_cases/abnormal_cl_rebase.py b/src/services/bunker_cases/abnormal_cl_rebase.py index 3a0dc9068..8e434579f 100644 --- a/src/services/bunker_cases/abnormal_cl_rebase.py +++ b/src/services/bunker_cases/abnormal_cl_rebase.py @@ -182,35 +182,13 @@ def _calculate_cl_rebase_between_blocks( self.w3.cc.get_validators_no_cache(prev_blockstamp), ) - # Get Lido validators' balances with WithdrawalVault balance - consensus_version = self.w3.lido_contracts.accounting_oracle.get_consensus_version(ref_blockstamp.block_hash) - - def _pre_electra_fork(): - ref = self._get_lido_validators_balance_with_vault_pre_electra( - ref_blockstamp, self.lido_validators - ) - prev = self._get_lido_validators_balance_with_vault_pre_electra( - prev_blockstamp, prev_lido_validators - ) - return ref, prev - - def _post_electra_fork(): - ref = self._get_lido_validators_balance_with_vault_post_electra( - ref_blockstamp, self.lido_validators - ) - prev = self._get_lido_validators_balance_with_vault_post_electra( - prev_blockstamp, prev_lido_validators - ) - return ref, prev + ref_lido_balance_with_vault = self._get_lido_validators_balance_with_vault( + ref_blockstamp, self.lido_validators + ) - if consensus_version in (1, 2): - ref_lido_balance_with_vault, prev_lido_balance_with_vault = _pre_electra_fork() - else: - spec = self.w3.cc.get_config_spec() - if ref_blockstamp.ref_epoch < int(spec.ELECTRA_FORK_EPOCH): - ref_lido_balance_with_vault, prev_lido_balance_with_vault = _pre_electra_fork() - else: - ref_lido_balance_with_vault, prev_lido_balance_with_vault = _post_electra_fork() + prev_lido_balance_with_vault = self._get_lido_validators_balance_with_vault( + prev_blockstamp, prev_lido_validators + ) # Raw CL rebase is calculated as difference between reference and previous Lido validators' balances # Without accounting withdrawals from WithdrawalVault @@ -234,31 +212,31 @@ def _post_electra_fork(): return cl_rebase - def _get_lido_validators_balance_with_vault_pre_electra( + def _get_lido_validators_balance_with_vault( self, blockstamp: BlockStamp, lido_validators: list[LidoValidator] ) -> Gwei: """ Get Lido validator balance with withdrawals vault balance """ - real_cl_balance = AbnormalClRebase.calculate_validators_balance_sum(lido_validators) - withdrawals_vault_balance = int( - self.w3.from_wei(self.w3.lido_contracts.get_withdrawal_balance_no_cache(blockstamp), "gwei") - ) - return Gwei(real_cl_balance + withdrawals_vault_balance) + consensus_version = self.w3.lido_contracts.accounting_oracle.get_consensus_version(blockstamp.block_hash) - def _get_lido_validators_balance_with_vault_post_electra( - self, blockstamp: BlockStamp, lido_validators: list[LidoValidator] - ) -> Gwei: - """ - Get Lido validator balance with withdrawals vault balance - """ real_cl_balance = AbnormalClRebase.calculate_validators_balance_sum(lido_validators) - state = self.w3.cc.get_state_view(blockstamp) - pending_deposits_sum = LidoValidatorsProvider.calculate_pending_deposits_sum(lido_validators, state.pending_deposits) withdrawals_vault_balance = int( self.w3.from_wei(self.w3.lido_contracts.get_withdrawal_balance_no_cache(blockstamp), "gwei") ) - return Gwei(real_cl_balance + pending_deposits_sum + withdrawals_vault_balance) + total_balance = real_cl_balance + withdrawals_vault_balance + + if consensus_version > 2: + epoch = EpochNumber(blockstamp.slot_number // self.c_conf.slots_per_epoch) + spec = self.w3.cc.get_config_spec() + if epoch >= int(spec.ELECTRA_FORK_EPOCH): + state = self.w3.cc.get_state_view(blockstamp) + pending_deposits_sum = LidoValidatorsProvider.calculate_pending_deposits_sum( + lido_validators, state.pending_deposits + ) + total_balance += pending_deposits_sum + + return Gwei(total_balance) def _get_withdrawn_from_vault_between_blocks( self, prev_blockstamp: BlockStamp, ref_blockstamp: ReferenceBlockStamp