Optimize proof memory

rocketpool/node/stake-megapool-validator.go

@@ -181,8 +181,22 @@ func (t *stakeMegapoolValidator) run(state *state.NetworkState) error {
 		// Log
 		t.log.Printlnf("The validator id %d needs to be staked", validatorId)
 
+		// Check if the validator is included in the finalized beacon state before attempting proof generation
+		validatorIndexStr, err := t.bc.GetValidatorIndex(validatorPubkey)
+		if err != nil {
+			return err
+		}
+		validatorIndex, err := strconv.ParseUint(validatorIndexStr, 10, 64)
+		if err != nil {
+			return err
+		}
+		if validatorIndex >= uint64(len(beaconState.GetValidators())) {
+			t.log.Printlnf("Validator id %d (beacon index %d) is not yet included in the finalized beacon state. Will retry on next cycle.", validatorId, validatorIndex)
+			continue
+		}
+
 		// Call Stake
-		err := t.stakeValidator(t.rp, beaconState, mp, validatorId, state, validatorPubkey, opts)
+		err = t.stakeValidator(t.rp, beaconState, mp, validatorId, state, validatorPubkey, opts)
 		if err != nil {
 			t.log.Printlnf("Error staking validator %d: %w", validatorId, err)
 			break
@@ -208,20 +222,6 @@ func (t *stakeMegapoolValidator) stakeValidator(rp *rocketpool.RocketPool, beaco
 		return err
 	}
 
-	// Check if the validator is included in the finalized beacon state before attempting proof generation
-	validatorIndexStr, err := t.bc.GetValidatorIndex(validatorPubkey)
-	if err != nil {
-		return err
-	}
-	validatorIndex, err := strconv.ParseUint(validatorIndexStr, 10, 64)
-	if err != nil {
-		return err
-	}
-	if validatorIndex >= uint64(len(beaconState.GetValidators())) {
-		t.log.Printlnf("Validator id %d (beacon index %d) is not yet included in the finalized beacon state. Will retry on next cycle.", validatorId, validatorIndex)
-		return nil
-	}
-
 	t.log.Printlnf("Crafting a proof that the correct credentials were used on the first beacon chain deposit. This process can take several seconds and is CPU and memory intensive.")
 
 	validatorProof, slotTimestamp, slotProof, err := services.GetValidatorProof(t.c, 0, t.w, state.BeaconConfig, mp.GetAddress(), validatorPubkey, beaconState)

shared/services/beacon/client.go

@@ -1,6 +1,7 @@
 package beacon
 
 import (
+	"io"
 	"math/big"
 	"sort"
 
@@ -158,12 +159,12 @@ const (
 // Fork is the consensus version, e.g, "deneb" or "electra"
 
 type BeaconStateSSZ struct {
-	Data []byte
+	Data io.ReadCloser
 	Fork string
 }
 
 type BeaconBlockSSZ struct {
-	Data []byte
+	Data io.ReadCloser
 	Fork string
 }
 

shared/services/beacon/client/std-http-client.go

@@ -1026,14 +1026,8 @@ func (c *StandardHttpClient) GetBeaconStateSSZ(slot uint64) (*beacon.BeaconState
 		return nil, fmt.Errorf("Could not get beacon state data: HTTP status %d", response.StatusCode)
 	}
 
-	// Slurp the body
-	body, err := io.ReadAll(response.Body)
-	if err != nil {
-		return nil, fmt.Errorf("Could not get beacon state data: %w", err)
-	}
-
 	return &beacon.BeaconStateSSZ{
-		Data: body,
+		Data: response.Body,
 		Fork: response.Header.Get(ResponseConsensusVersionHeader),
 	}, nil
 }
@@ -1054,13 +1048,8 @@ func (c *StandardHttpClient) GetBeaconBlockSSZ(slot uint64) (*beacon.BeaconBlock
 		return nil, false, fmt.Errorf("Could not get beacon block data: HTTP status %d; response body: '%s'", response.StatusCode, string(responseBody))
 	}
 
-	// Slurp the body
-	body, err := io.ReadAll(response.Body)
-	if err != nil {
-		return nil, false, fmt.Errorf("Could not get beacon block data: %w", err)
-	}
 	return &beacon.BeaconBlockSSZ{
-		Data: body,
+		Data: response.Body,
 		Fork: response.Header.Get(ResponseConsensusVersionHeader),
 	}, true, nil
 }

shared/services/megapools.go

@@ -78,12 +78,8 @@ func GetValidatorProof(c *cli.Command, slot uint64, wallet wallet.Wallet, eth2Co
 		}
 	}
 
-	slotProofBytes, err := beaconState.SlotProof(beaconState.GetSlot())
-	if err != nil {
-		return megapool.ValidatorProof{}, 0, megapool.SlotProof{}, err
-	}
-
-	proofBytes, err := beaconState.ValidatorProof(validatorIndex64)
+	// Build the validator and slot proofs from a single state proof tree
+	proofBytes, slotProofBytes, err := beaconState.ValidatorAndSlotProof(validatorIndex64)
 	if err != nil {
 		return megapool.ValidatorProof{}, 0, megapool.SlotProof{}, err
 	}
@@ -179,7 +175,7 @@ func GetWithdrawableEpochProof(c *cli.Command, wallet *wallet.Wallet, eth2Config
 		return api.ValidatorWithdrawableEpochProof{}, fmt.Errorf("validator %d is not withdrawable", validatorIndex64)
 	}
 
-	proofBytes, err := beaconState.ValidatorProof(validatorIndex64)
+	proofBytes, _, err := beaconState.ValidatorAndSlotProof(validatorIndex64)
 	if err != nil {
 		return api.ValidatorWithdrawableEpochProof{}, err
 	}

shared/types/eth2/fork/electra/state_electra.go

@@ -82,105 +82,60 @@ func GetGeneralizedIndexForSlot() uint64 {
 	return math.GetPowerOfTwoCeil(getStateChunkSize()) + generic.BeaconStateSlotIndex
 }
 
-func (state *BeaconState) validatorStateProof(index uint64) ([][]byte, error) {
+// ValidatorAndSlotProof produces both the validator proof and the slot proof
+// for the state's current slot lowering the memory cost for building the proofs.
+func (state *BeaconState) ValidatorAndSlotProof(validatorIndex uint64) ([][]byte, [][]byte, error) {
 
-	// Convert the state to a proof tree
-	root, err := state.GetTree()
-	if err != nil {
-		return nil, fmt.Errorf("could not get state tree: %w", err)
+	if validatorIndex >= uint64(len(state.Validators)) {
+		return nil, nil, errors.New("validator index out of bounds")
 	}
 
-	// Find the validator's generalized index
-	generalizedIndex := generic.GetGeneralizedIndexForValidator(index, GetGeneralizedIndexForValidators())
-
-	// Grab the proof for that index
-	proof, err := root.Prove(int(generalizedIndex))
+	stateTree, err := state.GetTree()
 	if err != nil {
-		return nil, fmt.Errorf("could not get proof for validator: %w", err)
+		return nil, nil, fmt.Errorf("could not get state tree: %w", err)
 	}
 
-	// Sanity check that the proof leaf matches the expected validator
-	validatorHashTreeRoot, err := state.Validators[index].HashTreeRoot()
-	if err != nil {
-		return nil, fmt.Errorf("could not get hash tree root for validator: %w", err)
-	}
-	if !bytes.Equal(proof.Leaf, validatorHashTreeRoot[:]) {
-		return nil, fmt.Errorf("proof leaf does not match expected validator")
-	}
-
-	return proof.Hashes, nil
-
-}
-
-func (state *BeaconState) slotStateProof() ([][]byte, error) {
-
-	// Convert the state to a proof tree
-	root, err := state.GetTree()
+	validatorGid := generic.GetGeneralizedIndexForValidator(validatorIndex, GetGeneralizedIndexForValidators())
+	validatorStateProof, err := stateTree.Prove(int(validatorGid))
 	if err != nil {
-		return nil, fmt.Errorf("could not get state tree: %w", err)
+		return nil, nil, fmt.Errorf("could not get proof for validator: %w", err)
 	}
 
-	// Find the slot field generalized index
-	generalizedIndex := GetGeneralizedIndexForSlot()
-
-	// Grab the proof for that index
-	proof, err := root.Prove(int(generalizedIndex))
+	// Sanity check that the proof leaf matches the expected validator
+	validatorHashTreeRoot, err := state.Validators[validatorIndex].HashTreeRoot()
 	if err != nil {
-		return nil, fmt.Errorf("could not get proof for slot: %w", err)
+		return nil, nil, fmt.Errorf("could not get hash tree root for validator: %w", err)
 	}
-
-	return proof.Hashes, nil
-
-}
-
-func (state *BeaconState) SlotProof(slot uint64) ([][]byte, error) {
-
-	if slot != state.Slot {
-		return nil, errors.New("slot requested does not match state slot")
+	if !bytes.Equal(validatorStateProof.Leaf, validatorHashTreeRoot[:]) {
+		return nil, nil, fmt.Errorf("proof leaf does not match expected validator")
 	}
 
-	proof, err := state.slotStateProof()
+	slotStateProof, err := stateTree.Prove(int(GetGeneralizedIndexForSlot()))
 	if err != nil {
-		return nil, fmt.Errorf("could not get slot state proof: %w", err)
+		return nil, nil, fmt.Errorf("could not get proof for slot: %w", err)
 	}
 
-	// The EL proves against BeaconBlockHeader root, so we need to merge the state proof with that.
-	generalizedIndex := generic.BeaconBlockHeaderStateRootGeneralizedIndex
-	root, err := state.LatestBlockHeader.GetTree()
+	// Drop the state tree before doing more work so the GC can reclaim it.
+	stateTree = nil
+
+	bhTree, err := state.LatestBlockHeader.GetTree()
 	if err != nil {
-		return nil, fmt.Errorf("could not get block header tree: %w", err)
+		return nil, nil, fmt.Errorf("could not get block header tree: %w", err)
 	}
-	blockHeaderProof, err := root.Prove(int(generalizedIndex))
+	blockHeaderProof, err := bhTree.Prove(int(generic.BeaconBlockHeaderStateRootGeneralizedIndex))
 	if err != nil {
-		return nil, fmt.Errorf("could not get proof for block header: %w", err)
-	}
-
-	return append(proof, blockHeaderProof.Hashes...), nil
-}
-
-func (state *BeaconState) ValidatorProof(index uint64) ([][]byte, error) {
-
-	if index >= uint64(len(state.Validators)) {
-		return nil, errors.New("validator index out of bounds")
+		return nil, nil, fmt.Errorf("could not get proof for block header: %w", err)
 	}
 
-	proof, err := state.validatorStateProof(index)
-	if err != nil {
-		return nil, fmt.Errorf("could not get validator state proof: %w", err)
-	}
+	validatorProof := make([][]byte, 0, len(validatorStateProof.Hashes)+len(blockHeaderProof.Hashes))
+	validatorProof = append(validatorProof, validatorStateProof.Hashes...)
+	validatorProof = append(validatorProof, blockHeaderProof.Hashes...)
 
-	// The EL proves against BeaconBlockHeader root, so we need to merge the state proof with that.
-	generalizedIndex := generic.BeaconBlockHeaderStateRootGeneralizedIndex
-	root, err := state.LatestBlockHeader.GetTree()
-	if err != nil {
-		return nil, fmt.Errorf("could not get block header tree: %w", err)
-	}
-	blockHeaderProof, err := root.Prove(int(generalizedIndex))
-	if err != nil {
-		return nil, fmt.Errorf("could not get proof for block header: %w", err)
-	}
+	slotProof := make([][]byte, 0, len(slotStateProof.Hashes)+len(blockHeaderProof.Hashes))
+	slotProof = append(slotProof, slotStateProof.Hashes...)
+	slotProof = append(slotProof, blockHeaderProof.Hashes...)
 
-	return append(proof, blockHeaderProof.Hashes...), nil
+	return validatorProof, slotProof, nil
 }
 
 func (state *BeaconState) blockHeaderToStateProof(blockHeader *generic.BeaconBlockHeader) ([][]byte, error) {

shared/types/eth2/fork/fulu/state_fulu.go

@@ -85,105 +85,60 @@ func GetGeneralizedIndexForSlot() uint64 {
 	return math.GetPowerOfTwoCeil(getStateChunkSize()) + generic.BeaconStateSlotIndex
 }
 
-func (state *BeaconState) validatorStateProof(index uint64) ([][]byte, error) {
+// ValidatorAndSlotProof produces both the validator proof and the slot proof
+// for the state's current slot
+func (state *BeaconState) ValidatorAndSlotProof(validatorIndex uint64) ([][]byte, [][]byte, error) {
 
-	// Convert the state to a proof tree
-	root, err := state.GetTree()
-	if err != nil {
-		return nil, fmt.Errorf("could not get state tree: %w", err)
+	if validatorIndex >= uint64(len(state.Validators)) {
+		return nil, nil, errors.New("validator index out of bounds")
 	}
 
-	// Find the validator's generalized index
-	generalizedIndex := generic.GetGeneralizedIndexForValidator(index, GetGeneralizedIndexForValidators())
-
-	// Grab the proof for that index
-	proof, err := root.Prove(int(generalizedIndex))
+	stateTree, err := state.GetTree()
 	if err != nil {
-		return nil, fmt.Errorf("could not get proof for validator: %w", err)
+		return nil, nil, fmt.Errorf("could not get state tree: %w", err)
 	}
 
-	// Sanity check that the proof leaf matches the expected validator
-	validatorHashTreeRoot, err := state.Validators[index].HashTreeRoot()
-	if err != nil {
-		return nil, fmt.Errorf("could not get hash tree root for validator: %w", err)
-	}
-	if !bytes.Equal(proof.Leaf, validatorHashTreeRoot[:]) {
-		return nil, fmt.Errorf("proof leaf does not match expected validator")
-	}
-
-	return proof.Hashes, nil
-
-}
-
-func (state *BeaconState) slotStateProof() ([][]byte, error) {
-
-	// Convert the state to a proof tree
-	root, err := state.GetTree()
+	validatorGid := generic.GetGeneralizedIndexForValidator(validatorIndex, GetGeneralizedIndexForValidators())
+	validatorStateProof, err := stateTree.Prove(int(validatorGid))
 	if err != nil {
-		return nil, fmt.Errorf("could not get state tree: %w", err)
+		return nil, nil, fmt.Errorf("could not get proof for validator: %w", err)
 	}
 
-	// Find the slot field generalized index
-	generalizedIndex := GetGeneralizedIndexForSlot()
-
-	// Grab the proof for that index
-	proof, err := root.Prove(int(generalizedIndex))
+	// Sanity check that the proof leaf matches the expected validator
+	validatorHashTreeRoot, err := state.Validators[validatorIndex].HashTreeRoot()
 	if err != nil {
-		return nil, fmt.Errorf("could not get proof for slot: %w", err)
+		return nil, nil, fmt.Errorf("could not get hash tree root for validator: %w", err)
 	}
-
-	return proof.Hashes, nil
-
-}
-
-func (state *BeaconState) SlotProof(slot uint64) ([][]byte, error) {
-
-	if slot != state.Slot {
-		return nil, errors.New("slot requested does not match state slot")
+	if !bytes.Equal(validatorStateProof.Leaf, validatorHashTreeRoot[:]) {
+		return nil, nil, fmt.Errorf("proof leaf does not match expected validator")
 	}
 
-	proof, err := state.slotStateProof()
+	slotStateProof, err := stateTree.Prove(int(GetGeneralizedIndexForSlot()))
 	if err != nil {
-		return nil, fmt.Errorf("could not get slot state proof: %w", err)
+		return nil, nil, fmt.Errorf("could not get proof for slot: %w", err)
 	}
 
-	// The EL proves against BeaconBlockHeader root, so we need to merge the state proof with that.
-	generalizedIndex := generic.BeaconBlockHeaderStateRootGeneralizedIndex
-	root, err := state.LatestBlockHeader.GetTree()
+	// Drop the state tree before doing more work so the GC can reclaim it.
+	stateTree = nil
+
+	bhTree, err := state.LatestBlockHeader.GetTree()
 	if err != nil {
-		return nil, fmt.Errorf("could not get block header tree: %w", err)
+		return nil, nil, fmt.Errorf("could not get block header tree: %w", err)
 	}
-	blockHeaderProof, err := root.Prove(int(generalizedIndex))
+	blockHeaderProof, err := bhTree.Prove(int(generic.BeaconBlockHeaderStateRootGeneralizedIndex))
 	if err != nil {
-		return nil, fmt.Errorf("could not get proof for block header: %w", err)
-	}
-
-	return append(proof, blockHeaderProof.Hashes...), nil
-}
-
-func (state *BeaconState) ValidatorProof(index uint64) ([][]byte, error) {
-
-	if index >= uint64(len(state.Validators)) {
-		return nil, errors.New("validator index out of bounds")
+		return nil, nil, fmt.Errorf("could not get proof for block header: %w", err)
 	}
 
-	proof, err := state.validatorStateProof(index)
-	if err != nil {
-		return nil, fmt.Errorf("could not get validator state proof: %w", err)
-	}
+	validatorProof := make([][]byte, 0, len(validatorStateProof.Hashes)+len(blockHeaderProof.Hashes))
+	validatorProof = append(validatorProof, validatorStateProof.Hashes...)
+	validatorProof = append(validatorProof, blockHeaderProof.Hashes...)
 
-	// The EL proves against BeaconBlockHeader root, so we need to merge the state proof with that.
-	generalizedIndex := generic.BeaconBlockHeaderStateRootGeneralizedIndex
-	root, err := state.LatestBlockHeader.GetTree()
-	if err != nil {
-		return nil, fmt.Errorf("could not get block header tree: %w", err)
-	}
-	blockHeaderProof, err := root.Prove(int(generalizedIndex))
-	if err != nil {
-		return nil, fmt.Errorf("could not get proof for block header: %w", err)
-	}
+	slotProof := make([][]byte, 0, len(slotStateProof.Hashes)+len(blockHeaderProof.Hashes))
+	slotProof = append(slotProof, slotStateProof.Hashes...)
+	slotProof = append(slotProof, blockHeaderProof.Hashes...)
 
-	return append(proof, blockHeaderProof.Hashes...), nil
+	return validatorProof, slotProof, nil
 }
 
 func (state *BeaconState) blockHeaderToStateProof(blockHeader *generic.BeaconBlockHeader) ([][]byte, error) {