Simplify NetworkState and StateManager

rocketpool/node/node.go

@@ -117,10 +117,7 @@ func run(c *cli.Context) error {
 	updateLog := log.NewColorLogger(UpdateColor)
 
 	// Create the state manager
-	m, err := state.NewNetworkStateManager(rp, cfg, rp.Client, bc, &updateLog)
-	if err != nil {
-		return err
-	}
+	m := state.NewNetworkStateManager(rp, cfg.Smartnode.GetStateManagerContracts(), bc, &updateLog)
 	stateLocker := collectors.NewStateLocker()
 
 	// Initialize tasks

rocketpool/watchtower/generate-rewards-tree.go

@@ -169,11 +169,7 @@ func (t *generateRewardsTree) generateRewardsTree(index uint64) {
 	address, err := client.RocketStorage.GetAddress(opts, crypto.Keccak256Hash([]byte("contract.addressrocketTokenRETH")))
 	if err == nil {
 		// Create the state manager with using the primary or fallback (not necessarily archive) EC
-		stateManager, err = state.NewNetworkStateManager(client, t.cfg, t.rp.Client, t.bc, &t.log)
-		if err != nil {
-			t.handleError(fmt.Errorf("error creating new NetworkStateManager with Archive EC: %w", err))
-			return
-		}
+		stateManager = state.NewNetworkStateManager(client, t.cfg.Smartnode.GetStateManagerContracts(), t.bc, &t.log)
 	} else {
 		// Check if an Archive EC is provided, and if using it would potentially resolve the error
 		errMessage := err.Error()
@@ -204,12 +200,16 @@ func (t *generateRewardsTree) generateRewardsTree(index uint64) {
 					t.handleError(fmt.Errorf("Error verifying rETH address with Archive EC: %w", err))
 					return
 				}
-				// Create the state manager with the archive EC
-				stateManager, err = state.NewNetworkStateManager(client, t.cfg, ec, t.bc, &t.log)
+
+				// Create a new rocketpool-go instance
+				archiveRP, err := rocketpool.NewRocketPool(ec, *t.rp.RocketStorageContract.Address)
 				if err != nil {
-					t.handleError(fmt.Errorf("Error creating new NetworkStateManager with ARchive EC: %w", err))
+					t.handleError(fmt.Errorf("Error instantiating client with Archive EC: %w", err))
 					return
 				}
+
+				// Create the state manager with the archive EC
+				stateManager = state.NewNetworkStateManager(archiveRP, t.cfg.Smartnode.GetStateManagerContracts(), t.bc, &t.log)
 			} else {
 				// No archive node specified
 				t.handleError(fmt.Errorf("***ERROR*** Primary EC cannot retrieve state for historical block %d and the Archive EC is not specified.", elBlockHeader.Number.Uint64()))

rocketpool/watchtower/submit-network-balances.go

@@ -297,10 +297,7 @@ func (t *submitNetworkBalances) getNetworkBalances(elBlockHeader *types.Header,
 	}
 
 	// Create a new state gen manager
-	mgr, err := state.NewNetworkStateManager(client, t.cfg, client.Client, t.bc, t.log)
-	if err != nil {
-		return networkBalances{}, fmt.Errorf("error creating network state manager for EL block %s, Beacon slot %d: %w", elBlock, beaconBlock, err)
-	}
+	mgr := state.NewNetworkStateManager(client, t.cfg.Smartnode.GetStateManagerContracts(), t.bc, t.log)
 
 	// Create a new state for the target block
 	state, err := mgr.GetStateForSlot(beaconBlock)

rocketpool/watchtower/submit-rewards-tree-rolling.go

@@ -282,11 +282,7 @@ func (t *submitRewardsTree_Rolling) run(headState *state.NetworkState) error {
 			}
 
 			// Generate the rewards state
-			stateMgr, err := state.NewNetworkStateManager(client, t.cfg, client.Client, t.bc, &t.log)
-			if err != nil {
-				t.handleError(fmt.Errorf("error creating state manager for rewards slot: %w", err))
-				return
-			}
+			stateMgr := state.NewNetworkStateManager(client, t.cfg.Smartnode.GetStateManagerContracts(), t.bc, &t.log)
 			state, err := stateMgr.GetStateForSlot(snapshotEnd.ConsensusBlock)
 			if err != nil {
 				t.handleError(fmt.Errorf("error getting state for rewards slot: %w", err))

rocketpool/watchtower/submit-rewards-tree-stateless.go

@@ -321,10 +321,7 @@ func (t *submitRewardsTree_Stateless) generateTreeImpl(rp *rocketpool.RocketPool
 	t.log.Printlnf("Rewards checkpoint has passed, starting Merkle tree generation for interval %d in the background.\n%s Snapshot Beacon block = %d, EL block = %d, running from %s to %s", currentIndex, t.generationPrefix, snapshotBeaconBlock, elBlockIndex, startTime, endTime)
 
 	// Create a new state gen manager
-	mgr, err := state.NewNetworkStateManager(rp, t.cfg, rp.Client, t.bc, t.log)
-	if err != nil {
-		return fmt.Errorf("error creating network state manager for EL block %d, Beacon slot %d: %w", elBlockIndex, snapshotBeaconBlock, err)
-	}
+	mgr := state.NewNetworkStateManager(rp, t.cfg.Smartnode.GetStateManagerContracts(), t.bc, t.log)
 
 	// Create a new state for the target block
 	state, err := mgr.GetStateForSlot(snapshotBeaconBlock)

rocketpool/watchtower/watchtower.go

@@ -111,10 +111,7 @@ func run(c *cli.Context) error {
 	updateLog := log.NewColorLogger(UpdateColor)
 
 	// Create the state manager
-	m, err := state.NewNetworkStateManager(rp, cfg, rp.Client, bc, &updateLog)
-	if err != nil {
-		return err
-	}
+	m := state.NewNetworkStateManager(rp, cfg.Smartnode.GetStateManagerContracts(), bc, &updateLog)
 
 	// Get the node address
 	nodeAccount, err := w.GetNodeAccount()

shared/services/config/smartnode-config.go

@@ -49,6 +49,12 @@ const (
 	RewardsExtensionSSZ  RewardsExtension = ".ssz"
 )
 
+// Contract addresses for multicall / network state manager
+type StateManagerContracts struct {
+	Multicaller    common.Address
+	BalanceBatcher common.Address
+}
+
 // Configuration for the Smartnode
 type SmartnodeConfig struct {
 	Title string `yaml:"-"`
@@ -851,6 +857,10 @@ func (cfg *SmartnodeConfig) GetRethAddress() common.Address {
 }
 
 func getDefaultDataDir(config *RocketPoolConfig) string {
+	if config == nil {
+		// Handle tests. Eventually we'll refactor so this isn't necessary.
+		return ""
+	}
 	return filepath.Join(config.RocketPoolDirectory, "data")
 }
 
@@ -1012,6 +1022,14 @@ func (cfg *SmartnodeConfig) GetBalanceBatcherAddress() string {
 	return cfg.balancebatcherAddress[cfg.Network.Value.(config.Network)]
 }
 
+// Utility function to get the state manager contracts
+func (cfg *SmartnodeConfig) GetStateManagerContracts() StateManagerContracts {
+	return StateManagerContracts{
+		Multicaller:    common.HexToAddress(cfg.GetMulticallAddress()),
+		BalanceBatcher: common.HexToAddress(cfg.GetBalanceBatcherAddress()),
+	}
+}
+
 func (cfg *SmartnodeConfig) GetFlashbotsProtectUrl() string {
 	return cfg.flashbotsProtectUrl[cfg.Network.Value.(config.Network)]
 }

shared/services/proposals/proposal-manager.go

@@ -43,10 +43,7 @@ func NewProposalManager(log *log.ColorLogger, cfg *config.RocketPoolConfig, rp *
 		return nil, fmt.Errorf("error creating node tree manager: %w", err)
 	}
 
-	stateMgr, err := state.NewNetworkStateManager(rp, cfg, rp.Client, bc, log)
-	if err != nil {
-		return nil, fmt.Errorf("error creating network state manager: %w", err)
-	}
+	stateMgr := state.NewNetworkStateManager(rp, cfg.Smartnode.GetStateManagerContracts(), bc, log)
 
 	logPrefix := "[PDAO Proposals]"
 	return &ProposalManager{

shared/services/state/manager.go

@@ -10,46 +10,51 @@ import (
 	"github.com/rocket-pool/rocketpool-go/rocketpool"
 	"github.com/rocket-pool/smartnode/shared/services/beacon"
 	"github.com/rocket-pool/smartnode/shared/services/config"
-	cfgtypes "github.com/rocket-pool/smartnode/shared/types/config"
 	"github.com/rocket-pool/smartnode/shared/utils/log"
 )
 
 type NetworkStateManager struct {
-	cfg          *config.RocketPoolConfig
-	rp           *rocketpool.RocketPool
-	ec           rocketpool.ExecutionClient
-	bc           beacon.Client
-	log          *log.ColorLogger
-	Config       *config.RocketPoolConfig
-	Network      cfgtypes.Network
-	ChainID      uint
-	BeaconConfig beacon.Eth2Config
+	rp  *rocketpool.RocketPool
+	bc  beacon.Client
+	log *log.ColorLogger
+
+	// Memoized Beacon config
+	beaconConfig *beacon.Eth2Config
+
+	// Multicaller and batch balance contract addresses
+	contracts config.StateManagerContracts
 }
 
 // Create a new manager for the network state
-func NewNetworkStateManager(rp *rocketpool.RocketPool, cfg *config.RocketPoolConfig, ec rocketpool.ExecutionClient, bc beacon.Client, log *log.ColorLogger) (*NetworkStateManager, error) {
+func NewNetworkStateManager(
+	rp *rocketpool.RocketPool,
+	contracts config.StateManagerContracts,
+	bc beacon.Client,
+	log *log.ColorLogger,
+) *NetworkStateManager {
 
 	// Create the manager
-	m := &NetworkStateManager{
-		cfg:     cfg,
-		rp:      rp,
-		ec:      ec,
-		bc:      bc,
-		log:     log,
-		Config:  cfg,
-		Network: cfg.Smartnode.Network.Value.(cfgtypes.Network),
-		ChainID: cfg.Smartnode.GetChainID(),
+	return &NetworkStateManager{
+		rp:        rp,
+		bc:        bc,
+		log:       log,
+		contracts: contracts,
+	}
+}
+
+func (m *NetworkStateManager) getBeaconConfig() (*beacon.Eth2Config, error) {
+	if m.beaconConfig != nil {
+		return m.beaconConfig, nil
 	}
 
 	// Get the Beacon config info
-	var err error
-	m.BeaconConfig, err = m.bc.GetEth2Config()
+	beaconConfig, err := m.bc.GetEth2Config()
 	if err != nil {
 		return nil, err
 	}
+	m.beaconConfig = &beaconConfig
 
-	return m, nil
-
+	return m.beaconConfig, nil
 }
 
 // Get the state of the network using the latest Execution layer block
@@ -86,27 +91,35 @@ func (m *NetworkStateManager) GetLatestBeaconBlock() (beacon.BeaconBlock, error)
 
 // Gets the latest valid finalized block
 func (m *NetworkStateManager) GetLatestFinalizedBeaconBlock() (beacon.BeaconBlock, error) {
+	beaconConfig, err := m.getBeaconConfig()
+	if err != nil {
+		return beacon.BeaconBlock{}, fmt.Errorf("error getting Beacon config: %w", err)
+	}
 	head, err := m.bc.GetBeaconHead()
 	if err != nil {
 		return beacon.BeaconBlock{}, fmt.Errorf("error getting Beacon chain head: %w", err)
 	}
-	targetSlot := head.FinalizedEpoch*m.BeaconConfig.SlotsPerEpoch + (m.BeaconConfig.SlotsPerEpoch - 1)
+	targetSlot := head.FinalizedEpoch*beaconConfig.SlotsPerEpoch + (beaconConfig.SlotsPerEpoch - 1)
 	return m.GetLatestProposedBeaconBlock(targetSlot)
 }
 
 // Gets the Beacon slot for the latest execution layer block
 func (m *NetworkStateManager) GetHeadSlot() (uint64, error) {
+	beaconConfig, err := m.getBeaconConfig()
+	if err != nil {
+		return 0, fmt.Errorf("error getting Beacon config: %w", err)
+	}
 	// Get the latest EL block
-	latestBlockHeader, err := m.ec.HeaderByNumber(context.Background(), nil)
+	latestBlockHeader, err := m.rp.Client.HeaderByNumber(context.Background(), nil)
 	if err != nil {
 		return 0, fmt.Errorf("error getting latest EL block: %w", err)
 	}
 
 	// Get the corresponding Beacon slot based on the timestamp
 	latestBlockTime := time.Unix(int64(latestBlockHeader.Time), 0)
-	genesisTime := time.Unix(int64(m.BeaconConfig.GenesisTime), 0)
+	genesisTime := time.Unix(int64(beaconConfig.GenesisTime), 0)
 	secondsSinceGenesis := uint64(latestBlockTime.Sub(genesisTime).Seconds())
-	targetSlot := secondsSinceGenesis / m.BeaconConfig.SecondsPerSlot
+	targetSlot := secondsSinceGenesis / beaconConfig.SecondsPerSlot
 	return targetSlot, nil
 }
 
@@ -131,7 +144,11 @@ func (m *NetworkStateManager) GetLatestProposedBeaconBlock(targetSlot uint64) (b
 
 // Get the state of the network at the provided Beacon slot
 func (m *NetworkStateManager) getState(slotNumber uint64) (*NetworkState, error) {
-	state, err := CreateNetworkState(m.cfg, m.rp, m.ec, m.bc, m.log, slotNumber, m.BeaconConfig)
+	beaconConfig, err := m.getBeaconConfig()
+	if err != nil {
+		return nil, fmt.Errorf("error getting Beacon config: %w", err)
+	}
+	state, err := createNetworkState(m.contracts, m.rp, m.bc, m.log, slotNumber, beaconConfig)
 	if err != nil {
 		return nil, err
 	}
@@ -140,7 +157,11 @@ func (m *NetworkStateManager) getState(slotNumber uint64) (*NetworkState, error)
 
 // Get the state of the network for a specific node only at the provided Beacon slot
 func (m *NetworkStateManager) getStateForNode(nodeAddress common.Address, slotNumber uint64, calculateTotalEffectiveStake bool) (*NetworkState, *big.Int, error) {
-	state, totalEffectiveStake, err := CreateNetworkStateForNode(m.cfg, m.rp, m.ec, m.bc, m.log, slotNumber, m.BeaconConfig, nodeAddress, calculateTotalEffectiveStake)
+	beaconConfig, err := m.getBeaconConfig()
+	if err != nil {
+		return nil, nil, fmt.Errorf("error getting Beacon config: %w", err)
+	}
+	state, totalEffectiveStake, err := createNetworkStateForNode(m.contracts, m.rp, m.bc, m.log, slotNumber, beaconConfig, nodeAddress, calculateTotalEffectiveStake)
 	if err != nil {
 		return nil, nil, err
 	}

shared/services/state/network-state.go

@@ -65,10 +65,7 @@ type NetworkState struct {
 }
 
 // Creates a snapshot of the entire Rocket Pool network state, on both the Execution and Consensus layers
-func CreateNetworkState(cfg *config.RocketPoolConfig, rp *rocketpool.RocketPool, ec rocketpool.ExecutionClient, bc beacon.Client, log *log.ColorLogger, slotNumber uint64, beaconConfig beacon.Eth2Config) (*NetworkState, error) {
-	// Get the relevant network contracts
-	multicallerAddress := common.HexToAddress(cfg.Smartnode.GetMulticallAddress())
-	balanceBatcherAddress := common.HexToAddress(cfg.Smartnode.GetBalanceBatcherAddress())
+func createNetworkState(batchContracts config.StateManagerContracts, rp *rocketpool.RocketPool, bc beacon.Client, log *log.ColorLogger, slotNumber uint64, beaconConfig *beacon.Eth2Config) (*NetworkState, error) {
 
 	// Get the execution block for the given slot
 	beaconBlock, exists, err := bc.GetBeaconBlock(fmt.Sprintf("%d", slotNumber))
@@ -92,15 +89,15 @@ func CreateNetworkState(cfg *config.RocketPoolConfig, rp *rocketpool.RocketPool,
 		MinipoolDetailsByNode:    map[common.Address][]*rpstate.NativeMinipoolDetails{},
 		BeaconSlotNumber:         slotNumber,
 		ElBlockNumber:            elBlockNumber,
-		BeaconConfig:             beaconConfig,
+		BeaconConfig:             *beaconConfig,
 		log:                      log,
 	}
 
 	state.logLine("Getting network state for EL block %d, Beacon slot %d", elBlockNumber, slotNumber)
 	start := time.Now()
 
 	// Network contracts and details
-	contracts, err := rpstate.NewNetworkContracts(rp, multicallerAddress, balanceBatcherAddress, opts)
+	contracts, err := rpstate.NewNetworkContracts(rp, batchContracts.Multicaller, batchContracts.BalanceBatcher, opts)
 	if err != nil {
 		return nil, fmt.Errorf("error getting network contracts: %w", err)
 	}
@@ -193,16 +190,12 @@ func CreateNetworkState(cfg *config.RocketPoolConfig, rp *rocketpool.RocketPool,
 
 // Creates a snapshot of the Rocket Pool network, but only for a single node
 // Also gets the total effective RPL stake of the network for convenience since this is required by several node routines
-func CreateNetworkStateForNode(cfg *config.RocketPoolConfig, rp *rocketpool.RocketPool, ec rocketpool.ExecutionClient, bc beacon.Client, log *log.ColorLogger, slotNumber uint64, beaconConfig beacon.Eth2Config, nodeAddress common.Address, calculateTotalEffectiveStake bool) (*NetworkState, *big.Int, error) {
+func createNetworkStateForNode(batchContracts config.StateManagerContracts, rp *rocketpool.RocketPool, bc beacon.Client, log *log.ColorLogger, slotNumber uint64, beaconConfig *beacon.Eth2Config, nodeAddress common.Address, calculateTotalEffectiveStake bool) (*NetworkState, *big.Int, error) {
 	steps := 5
 	if calculateTotalEffectiveStake {
 		steps++
 	}
 
-	// Get the relevant network contracts
-	multicallerAddress := common.HexToAddress(cfg.Smartnode.GetMulticallAddress())
-	balanceBatcherAddress := common.HexToAddress(cfg.Smartnode.GetBalanceBatcherAddress())
-
 	// Get the execution block for the given slot
 	beaconBlock, exists, err := bc.GetBeaconBlock(fmt.Sprintf("%d", slotNumber))
 	if err != nil {
@@ -225,15 +218,15 @@ func CreateNetworkStateForNode(cfg *config.RocketPoolConfig, rp *rocketpool.Rock
 		MinipoolDetailsByNode:    map[common.Address][]*rpstate.NativeMinipoolDetails{},
 		BeaconSlotNumber:         slotNumber,
 		ElBlockNumber:            elBlockNumber,
-		BeaconConfig:             beaconConfig,
+		BeaconConfig:             *beaconConfig,
 		log:                      log,
 	}
 
 	state.logLine("Getting network state for EL block %d, Beacon slot %d", elBlockNumber, slotNumber)
 	start := time.Now()
 
 	// Network contracts and details
-	contracts, err := rpstate.NewNetworkContracts(rp, multicallerAddress, balanceBatcherAddress, opts)
+	contracts, err := rpstate.NewNetworkContracts(rp, batchContracts.Multicaller, batchContracts.BalanceBatcher, opts)
 	if err != nil {
 		return nil, nil, fmt.Errorf("error getting network contracts: %w", err)
 	}