From 8e4fcf035c97407373c209494dea132ad53bcc2e Mon Sep 17 00:00:00 2001 From: Fornax <23104993+0xfornax@users.noreply.github.com> Date: Wed, 4 Mar 2026 18:57:34 -0300 Subject: [PATCH 1/2] Optimize state loading --- rocketpool/node/notify-final-balance.go | 79 ++++++++++----------- rocketpool/node/notify-validator-exit.go | 67 +++++++++-------- rocketpool/node/stake-megapool-validator.go | 49 ++++++------- 3 files changed, 100 insertions(+), 95 deletions(-) diff --git a/rocketpool/node/notify-final-balance.go b/rocketpool/node/notify-final-balance.go index 3fab917ce..8aa9105b3 100644 --- a/rocketpool/node/notify-final-balance.go +++ b/rocketpool/node/notify-final-balance.go @@ -9,7 +9,6 @@ import ( "github.com/ethereum/go-ethereum/accounts/abi/bind" "github.com/rocket-pool/smartnode/bindings/megapool" "github.com/rocket-pool/smartnode/bindings/rocketpool" - "github.com/rocket-pool/smartnode/bindings/types" "github.com/rocket-pool/smartnode/bindings/utils/eth" "github.com/urfave/cli" @@ -19,8 +18,6 @@ import ( rpgas "github.com/rocket-pool/smartnode/shared/services/gas" "github.com/rocket-pool/smartnode/shared/services/state" "github.com/rocket-pool/smartnode/shared/services/wallet" - apitypes "github.com/rocket-pool/smartnode/shared/types/api" - "github.com/rocket-pool/smartnode/shared/types/eth2" "github.com/rocket-pool/smartnode/shared/utils/api" "github.com/rocket-pool/smartnode/shared/utils/log" ) @@ -119,50 +116,48 @@ func (t *notifyFinalBalance) run(state *state.NetworkState) error { return err } - // Check if the megapool is deployed - deployed, err := megapool.GetMegapoolDeployed(t.rp, nodeAccount.Address, opts) - if err != nil { - return err + nodeDetails, exists := state.NodeDetailsByAddress[nodeAccount.Address] + if !exists { + return fmt.Errorf("node account %s not found in state", nodeAccount.Address.Hex()) } - if !deployed { + + if !nodeDetails.MegapoolDeployed { return nil } - // Get the megapool address - megapoolAddress, err := megapool.GetMegapoolExpectedAddress(t.rp, nodeAccount.Address, opts) - if err != nil { - return err - } + megapoolAddress := nodeDetails.MegapoolAddress - // Load the megapool mp, err := megapool.NewMegaPoolV1(t.rp, megapoolAddress, nil) if err != nil { return err } - // Iterate over megapool validators checking whether they're ready to submit a final balance proof - validatorCount, err := mp.GetValidatorCount(nil) - if err != nil { - return err - } - validatorInfo, err := services.GetMegapoolValidatorDetails(t.rp, t.bc, mp, megapoolAddress, uint32(validatorCount), opts) - if err != nil { - return err + validatorDetailsToProve := make(map[uint32]beacon.ValidatorStatus) + pubkeys := state.MegapoolToPubkeysMap[megapoolAddress] + for _, pubkey := range pubkeys { + validatorDetails, exists := state.MegapoolValidatorDetails[pubkey] + if !exists { + return fmt.Errorf("validator %s not found in state", pubkey.Hex()) + } + + validatorInfo := state.MegapoolValidatorInfo[pubkey] + + if validatorDetails.Status == beacon.ValidatorState_WithdrawalDone && validatorInfo.ValidatorInfo.Exiting && !validatorInfo.ValidatorInfo.Exited && validatorDetails.EffectiveBalance == 0 { + validatorDetailsToProve[validatorInfo.ValidatorId] = validatorDetails + } } - // Get the beacon state - beaconState, err := services.GetBeaconState(t.bc) - if err != nil { - return err + // Check if there are any validators to notify + if len(validatorDetailsToProve) == 0 { + return nil } - for i := uint32(0); i < uint32(validatorCount); i++ { - if validatorInfo[i].BeaconStatus.Status == "withdrawal_done" && validatorInfo[i].Exiting && !validatorInfo[i].Exited && validatorInfo[i].BeaconStatus.EffectiveBalance == 0 { - // Log - t.log.Printlnf("The validator ID %d needs a final balance proof", validatorInfo[i].ValidatorId) + // Notify the validators + for validatorId, validatorDetails := range validatorDetailsToProve { + // Log + t.log.Printlnf("The validator id %d needs a final balance proof", validatorId) - t.createFinalBalanceProof(t.rp, mp, validatorInfo[i], state, types.ValidatorPubkey(validatorInfo[i].PubKey), beaconState, opts) - } + t.createFinalBalanceProof(t.rp, mp, state, validatorId, validatorDetails, opts) } // Return @@ -170,7 +165,7 @@ func (t *notifyFinalBalance) run(state *state.NetworkState) error { } -func (t *notifyFinalBalance) createFinalBalanceProof(rp *rocketpool.RocketPool, mp megapool.Megapool, validatorInfo apitypes.MegapoolValidatorDetails, state *state.NetworkState, validatorPubkey types.ValidatorPubkey, beaconState eth2.BeaconState, callopts *bind.CallOpts) error { +func (t *notifyFinalBalance) createFinalBalanceProof(rp *rocketpool.RocketPool, mp megapool.Megapool, state *state.NetworkState, validatorId uint32, validatorDetails beacon.ValidatorStatus, callopts *bind.CallOpts) error { // Get transactor opts, err := t.w.GetNodeAccountTransactor() @@ -180,7 +175,7 @@ func (t *notifyFinalBalance) createFinalBalanceProof(rp *rocketpool.RocketPool, t.log.Printlnf("Crafting a final balance proof. This process can take several seconds and is CPU and memory intensive. If you don't see a [FINISHED] log entry your system may not have enough resources to perform this operation.") - validatorIndexStr, err := t.bc.GetValidatorIndex(validatorPubkey) + validatorIndexStr, err := t.bc.GetValidatorIndex(validatorDetails.Pubkey) if err != nil { return err } @@ -190,15 +185,15 @@ func (t *notifyFinalBalance) createFinalBalanceProof(rp *rocketpool.RocketPool, return err } - slot := validatorInfo.WithdrawableEpoch * 32 + slot := validatorDetails.WithdrawableEpoch * 32 withdrawalProof, proofSlot, stateUsed, err := services.GetWithdrawalProofForSlot(t.c, slot, validatorIndex) if err != nil { - return fmt.Errorf("error getting withdrawal proof for validator 0x%s (index: %d): %w", validatorPubkey.String(), validatorIndex, err) + return fmt.Errorf("error getting withdrawal proof for validator 0x%s (index: %d): %w", validatorDetails.Pubkey.String(), validatorIndex, err) } - t.log.Printlnf("The Beacon WithdrawalSlot for validator ID %d is: %d", validatorInfo.ValidatorId, withdrawalProof.WithdrawalSlot) + t.log.Printlnf("The Beacon WithdrawalSlot for validator index %d is: %d", validatorDetails.Index, withdrawalProof.WithdrawalSlot) - validatorProof, slotTimestamp, slotProof, err := services.GetValidatorProof(t.c, proofSlot, t.w, state.BeaconConfig, mp.GetAddress(), validatorPubkey, stateUsed) + validatorProof, slotTimestamp, slotProof, err := services.GetValidatorProof(t.c, proofSlot, t.w, state.BeaconConfig, mp.GetAddress(), validatorDetails.Pubkey, stateUsed) if err != nil { t.log.Printlnf("There was an error during the proof creation process: %w", err) return err @@ -221,9 +216,9 @@ func (t *notifyFinalBalance) createFinalBalanceProof(rp *rocketpool.RocketPool, t.log.Printlnf("The validator final balance proof has been successfully created.") // Get the gas limit - gasInfo, err := megapool.EstimateNotifyFinalBalance(rp, mp.GetAddress(), validatorInfo.ValidatorId, slotTimestamp, finalBalanceProof, validatorProof, slotProof, opts) + gasInfo, err := megapool.EstimateNotifyFinalBalance(rp, mp.GetAddress(), validatorId, slotTimestamp, finalBalanceProof, validatorProof, slotProof, opts) if err != nil { - t.log.Printlnf("Could not estimate the gas required to notify final balance on megapool validator %d: %w", validatorInfo.ValidatorId, err) + t.log.Printlnf("Could not estimate the gas required to notify final balance on megapool validator %d: %w", validatorId, err) return err } gas := big.NewInt(int64(gasInfo.SafeGasLimit)) @@ -246,7 +241,7 @@ func (t *notifyFinalBalance) createFinalBalanceProof(rp *rocketpool.RocketPool, opts.GasLimit = gas.Uint64() // Call Notify Final Balance - tx, err := megapool.NotifyFinalBalance(rp, mp.GetAddress(), validatorInfo.ValidatorId, slotTimestamp, finalBalanceProof, validatorProof, slotProof, opts) + tx, err := megapool.NotifyFinalBalance(rp, mp.GetAddress(), validatorId, slotTimestamp, finalBalanceProof, validatorProof, slotProof, opts) if err != nil { return err } @@ -258,7 +253,7 @@ func (t *notifyFinalBalance) createFinalBalanceProof(rp *rocketpool.RocketPool, } // Log - t.log.Printlnf("Successfully notified validator %d final balance.", validatorInfo.ValidatorId) + t.log.Printlnf("Successfully notified validator %d final balance.", validatorId) // Return return nil diff --git a/rocketpool/node/notify-validator-exit.go b/rocketpool/node/notify-validator-exit.go index e4570d92e..2121fd138 100644 --- a/rocketpool/node/notify-validator-exit.go +++ b/rocketpool/node/notify-validator-exit.go @@ -1,6 +1,7 @@ package node import ( + "fmt" "math/big" "github.com/docker/docker/client" @@ -118,55 +119,63 @@ func (t *notifyValidatorExit) run(state *state.NetworkState) error { return err } - // Check if the megapool is deployed - deployed, err := megapool.GetMegapoolDeployed(t.rp, nodeAccount.Address, opts) - if err != nil { - return err + nodeDetails, exists := state.NodeDetailsByAddress[nodeAccount.Address] + if !exists { + return fmt.Errorf("node account %s not found in state", nodeAccount.Address.Hex()) } - if !deployed { + + if !nodeDetails.MegapoolDeployed { return nil } - // Get the megapool address - megapoolAddress, err := megapool.GetMegapoolExpectedAddress(t.rp, nodeAccount.Address, opts) - if err != nil { - return err - } + megapoolAddress := nodeDetails.MegapoolAddress - // Load the megapool mp, err := megapool.NewMegaPoolV1(t.rp, megapoolAddress, nil) if err != nil { return err } - // Iterate over megapool validators checking whether they're ready to notify exit - validatorCount, err := mp.GetValidatorCount(opts) + var currentEpoch uint64 + + head, err := t.bc.GetBeaconHead() if err != nil { return err } - validatorInfo, err := services.GetMegapoolValidatorDetails(t.rp, t.bc, mp, megapoolAddress, uint32(validatorCount), opts) + currentEpoch = head.Epoch + + validatorDetailsToProve := make(map[uint32]beacon.ValidatorStatus) + pubkeys := state.MegapoolToPubkeysMap[megapoolAddress] + for _, pubkey := range pubkeys { + validatorDetails, exists := state.MegapoolValidatorDetails[pubkey] + if !exists { + return fmt.Errorf("validator %s not found in state", pubkey.Hex()) + } + + validatorInfo := state.MegapoolValidatorInfo[pubkey] + + if currentEpoch > validatorDetails.ActivationEpoch && validatorDetails.WithdrawableEpoch < FarFutureEpoch && validatorInfo.ValidatorInfo.Staked && !validatorInfo.ValidatorInfo.Exited && !validatorInfo.ValidatorInfo.Exiting { + validatorDetailsToProve[validatorInfo.ValidatorId] = validatorDetails + } + } + + // Check if there are any validators to notify + if len(validatorDetailsToProve) == 0 { + return nil + } + + beaconState, err := services.GetBeaconState(t.bc) if err != nil { return err } - for i := uint32(0); i < uint32(validatorCount); i++ { - if validatorInfo[i].Activated && validatorInfo[i].WithdrawableEpoch < FarFutureEpoch && validatorInfo[i].Staked && !validatorInfo[i].Exited && !validatorInfo[i].Exiting { - beaconState, err := services.GetBeaconState(t.bc) - if err != nil { - return err - } - - if beaconState.GetValidators()[validatorInfo[i].ValidatorIndex].WithdrawableEpoch < FarFutureEpoch { + for validatorId, validatorDetails := range validatorDetailsToProve { - // Log - t.log.Printlnf("The validator ID %d needs an exit proof", validatorInfo[i].ValidatorId) + // Log + t.log.Printlnf("The validator id %d needs an exit proof", validatorId) - // Call Notify Exit - t.createExitProof(t.rp, beaconState, mp, validatorInfo[i].ValidatorId, state, types.ValidatorPubkey(validatorInfo[i].PubKey), opts) - } - } + // Call Notify Exit + t.createExitProof(t.rp, beaconState, mp, validatorId, state, types.ValidatorPubkey(validatorDetails.Pubkey), opts) } - // Return return nil diff --git a/rocketpool/node/stake-megapool-validator.go b/rocketpool/node/stake-megapool-validator.go index ddbdb6f88..66985f815 100644 --- a/rocketpool/node/stake-megapool-validator.go +++ b/rocketpool/node/stake-megapool-validator.go @@ -2,7 +2,6 @@ package node import ( "math/big" - "strconv" "github.com/docker/docker/client" "github.com/ethereum/go-ethereum/accounts/abi/bind" @@ -149,38 +148,40 @@ func (t *stakeMegapoolValidator) run(state *state.NetworkState) error { return err } + // store validators that need to be staked + validatorsToStake := make(map[uint32]bool) + + for i := uint32(0); i < uint32(validatorCount); i++ { + if validatorInfo[i].InPrestake && validatorInfo[i].BeaconStatus.Index != "" { + validatorsToStake[validatorInfo[i].ValidatorId] = true + } + } + + // Check if we have any validators to stake + if len(validatorsToStake) == 0 { + return nil + } + + // Load the beacon state beaconState, err := services.GetBeaconState(t.bc) if err != nil { return err } - stakedCount := 0 - for i := uint32(0); i < uint32(validatorCount); i++ { - if validatorInfo[i].InPrestake && validatorInfo[i].BeaconStatus.Index != "" { - // Convert str to int - indexInt, err := strconv.Atoi(validatorInfo[i].BeaconStatus.Index) - if err != nil { - return err - } - if indexInt < len(beaconState.GetValidators()) { - // Log - t.log.Printlnf("The validator %d needs to be staked", validatorInfo[i].ValidatorId) - - // Call Stake - t.stakeValidator(t.rp, beaconState, mp, validatorInfo[i].ValidatorId, state, types.ValidatorPubkey(validatorInfo[i].PubKey), opts) - stakedCount++ - } - } + + // Iterate over validators to stake + for validatorId, _ := range validatorsToStake { + // Log + t.log.Printlnf("The validator id %d needs to be staked", validatorId) + + // Call Stake + t.stakeValidator(t.rp, beaconState, mp, validatorId, state, types.ValidatorPubkey(validatorInfo[validatorId].PubKey), opts) } - if stakedCount > 0 { - if err := validator.RestartValidator(t.cfg, t.bc, &t.log, t.d); err != nil { - return err - } + if err := validator.RestartValidator(t.cfg, t.bc, &t.log, t.d); err != nil { + return err } - // Return return nil - } func (t *stakeMegapoolValidator) stakeValidator(rp *rocketpool.RocketPool, beaconState eth2.BeaconState, mp megapool.Megapool, validatorId uint32, state *state.NetworkState, validatorPubkey types.ValidatorPubkey, callopts *bind.CallOpts) error { From 038dcf172abdd19848e4433925e1068587a8e073 Mon Sep 17 00:00:00 2001 From: Fornax <23104993+fornax2@users.noreply.github.com> Date: Thu, 5 Mar 2026 01:02:40 -0300 Subject: [PATCH 2/2] Skip validators in prestake --- rocketpool/node/notify-final-balance.go | 3 ++- rocketpool/node/notify-validator-exit.go | 3 ++- rocketpool/node/stake-megapool-validator.go | 8 ++++---- 3 files changed, 8 insertions(+), 6 deletions(-) diff --git a/rocketpool/node/notify-final-balance.go b/rocketpool/node/notify-final-balance.go index 8aa9105b3..0cb0c75b5 100644 --- a/rocketpool/node/notify-final-balance.go +++ b/rocketpool/node/notify-final-balance.go @@ -137,7 +137,8 @@ func (t *notifyFinalBalance) run(state *state.NetworkState) error { for _, pubkey := range pubkeys { validatorDetails, exists := state.MegapoolValidatorDetails[pubkey] if !exists { - return fmt.Errorf("validator %s not found in state", pubkey.Hex()) + // validator not in the state yet, skipping + continue } validatorInfo := state.MegapoolValidatorInfo[pubkey] diff --git a/rocketpool/node/notify-validator-exit.go b/rocketpool/node/notify-validator-exit.go index 2121fd138..4cb99abc8 100644 --- a/rocketpool/node/notify-validator-exit.go +++ b/rocketpool/node/notify-validator-exit.go @@ -148,7 +148,8 @@ func (t *notifyValidatorExit) run(state *state.NetworkState) error { for _, pubkey := range pubkeys { validatorDetails, exists := state.MegapoolValidatorDetails[pubkey] if !exists { - return fmt.Errorf("validator %s not found in state", pubkey.Hex()) + // validator not in the state yet, skipping + continue } validatorInfo := state.MegapoolValidatorInfo[pubkey] diff --git a/rocketpool/node/stake-megapool-validator.go b/rocketpool/node/stake-megapool-validator.go index 66985f815..5a445395f 100644 --- a/rocketpool/node/stake-megapool-validator.go +++ b/rocketpool/node/stake-megapool-validator.go @@ -149,11 +149,11 @@ func (t *stakeMegapoolValidator) run(state *state.NetworkState) error { } // store validators that need to be staked - validatorsToStake := make(map[uint32]bool) + validatorsToStake := make(map[uint32]types.ValidatorPubkey) for i := uint32(0); i < uint32(validatorCount); i++ { if validatorInfo[i].InPrestake && validatorInfo[i].BeaconStatus.Index != "" { - validatorsToStake[validatorInfo[i].ValidatorId] = true + validatorsToStake[validatorInfo[i].ValidatorId] = types.ValidatorPubkey(validatorInfo[i].PubKey) } } @@ -169,12 +169,12 @@ func (t *stakeMegapoolValidator) run(state *state.NetworkState) error { } // Iterate over validators to stake - for validatorId, _ := range validatorsToStake { + for validatorId, validatorPubkey := range validatorsToStake { // Log t.log.Printlnf("The validator id %d needs to be staked", validatorId) // Call Stake - t.stakeValidator(t.rp, beaconState, mp, validatorId, state, types.ValidatorPubkey(validatorInfo[validatorId].PubKey), opts) + t.stakeValidator(t.rp, beaconState, mp, validatorId, state, validatorPubkey, opts) } if err := validator.RestartValidator(t.cfg, t.bc, &t.log, t.d); err != nil {