Make memory improvements to state ssz parsing automatically

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

@@ -169,10 +169,6 @@ func GetWithdrawableEpochProof(c *cli.Command, wallet *wallet.Wallet, eth2Config
 	if err != nil {
 		return api.ValidatorWithdrawableEpochProof{}, err
 	}
-	// Drop the raw SSZ buffer (hundreds of MB on mainnet) now that the state
-	// has been unmarshalled into its own owned data, so the proof-tree allocation
-	// below can reuse the freed memory.
-	beaconStateResponse.Data = nil
 
 	withdrawableEpoch := beaconState.GetValidators()[validatorIndex64].WithdrawableEpoch
 	if withdrawableEpoch == math.MaxUint64 {
@@ -730,9 +726,6 @@ func GetWithdrawalProofForSlot(c *cli.Command, slot uint64, validatorIndex uint6
 	if err != nil {
 		return megapool.FinalBalanceProof{}, 0, nil, err
 	}
-	// Drop the raw SSZ buffer (hundreds of MB on mainnet) now that the state
-	// owns its own data; the proof-tree work below can reuse the freed memory.
-	stateResponse.Data = nil
 
 	fuluState, ok := beaconState.(*fulu.BeaconState)
 	if !ok {
@@ -779,9 +772,6 @@ func GetWithdrawalProofForSlot(c *cli.Command, slot uint64, validatorIndex uint6
 		if err != nil {
 			return megapool.FinalBalanceProof{}, 0, nil, err
 		}
-		// Drop the raw SSZ buffer for the block-roots state now that it has
-		// been unmarshalled, so it doesn't overlap with the proof tree below.
-		blockRootsStateResponse.Data = nil
 		summaryProof, err = blockRootsState.HistoricalSummaryBlockRootProof(int(response.WithdrawalSlot))
 		if err != nil {
 			return megapool.FinalBalanceProof{}, 0, nil, err

shared/services/services.go

@@ -232,10 +232,6 @@ func GetBeaconState(bc beacon.Client) (eth2.BeaconState, error) {
 	if err != nil {
 		return nil, err
 	}
-	// Drop the raw SSZ buffer (hundreds of MB on mainnet) now that the state
-	// has been unmarshalled into its own owned data; the proof-tree allocations
-	// downstream can reuse the freed memory.
-	beaconStateResponse.Data = nil
 	return beaconState, nil
 }
 

shared/types/eth2/types.go

@@ -2,6 +2,7 @@ package eth2
 
 import (
 	"fmt"
+	"io"
 	"strings"
 
 	"github.com/rocket-pool/smartnode/shared/types/eth2/fork/deneb"
@@ -37,20 +38,29 @@ type SignedBeaconBlock interface {
 	Withdrawals() []*generic.Withdrawal
 }
 
-func NewBeaconState(data []byte, fork string) (BeaconState, error) {
+func NewBeaconState(data io.ReadCloser, fork string) (BeaconState, error) {
 	fork = strings.ToLower(fork)
 
+	defer func() {
+		_ = data.Close()
+	}()
+
+	dataBytes, err := io.ReadAll(data)
+	if err != nil {
+		return nil, err
+	}
+
 	switch fork {
 	case "electra":
 		out := &electra.BeaconState{}
-		err := out.UnmarshalSSZ(data)
+		err := out.UnmarshalSSZ(dataBytes)
 		if err != nil {
 			return nil, err
 		}
 		return out, nil
 	case "fulu":
 		out := &fulu.BeaconState{}
-		err := out.UnmarshalSSZ(data)
+		err := out.UnmarshalSSZ(dataBytes)
 		if err != nil {
 			return nil, err
 		}
@@ -60,27 +70,36 @@ func NewBeaconState(data []byte, fork string) (BeaconState, error) {
 	}
 }
 
-func NewSignedBeaconBlock(data []byte, fork string) (SignedBeaconBlock, error) {
+func NewSignedBeaconBlock(data io.ReadCloser, fork string) (SignedBeaconBlock, error) {
 	fork = strings.ToLower(fork)
 
+	defer func() {
+		_ = data.Close()
+	}()
+
+	dataBytes, err := io.ReadAll(data)
+	if err != nil {
+		return nil, err
+	}
+
 	switch fork {
 	case "deneb":
 		out := &deneb.SignedBeaconBlock{}
-		err := out.UnmarshalSSZ(data)
+		err := out.UnmarshalSSZ(dataBytes)
 		if err != nil {
 			return nil, err
 		}
 		return out, nil
 	case "electra":
 		out := &electra.SignedBeaconBlock{}
-		err := out.UnmarshalSSZ(data)
+		err := out.UnmarshalSSZ(dataBytes)
 		if err != nil {
 			return nil, err
 		}
 		return out, nil
 	case "fulu":
 		out := &fulu.SignedBeaconBlock{}
-		err := out.UnmarshalSSZ(data)
+		err := out.UnmarshalSSZ(dataBytes)
 		if err != nil {
 			return nil, err
 		}