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/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
 	}
@@ -173,6 +169,10 @@ 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,6 +730,9 @@ 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 {
@@ -776,6 +779,9 @@ 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,6 +232,10 @@ 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/fork/electra/state_electra.go

@@ -183,6 +183,62 @@ func (state *BeaconState) ValidatorProof(index uint64) ([][]byte, error) {
 	return append(proof, blockHeaderProof.Hashes...), nil
 }
 
+// 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) {
+
+	if validatorIndex >= uint64(len(state.Validators)) {
+		return nil, nil, errors.New("validator index out of bounds")
+	}
+
+	stateTree, err := state.GetTree()
+	if err != nil {
+		return nil, nil, fmt.Errorf("could not get state tree: %w", err)
+	}
+
+	validatorGid := generic.GetGeneralizedIndexForValidator(validatorIndex, GetGeneralizedIndexForValidators())
+	validatorStateProof, err := stateTree.Prove(int(validatorGid))
+	if err != nil {
+		return nil, nil, fmt.Errorf("could not get proof for validator: %w", err)
+	}
+
+	// Sanity check that the proof leaf matches the expected validator
+	validatorHashTreeRoot, err := state.Validators[validatorIndex].HashTreeRoot()
+	if err != nil {
+		return nil, nil, fmt.Errorf("could not get hash tree root for validator: %w", err)
+	}
+	if !bytes.Equal(validatorStateProof.Leaf, validatorHashTreeRoot[:]) {
+		return nil, nil, fmt.Errorf("proof leaf does not match expected validator")
+	}
+
+	slotStateProof, err := stateTree.Prove(int(GetGeneralizedIndexForSlot()))
+	if err != nil {
+		return nil, nil, fmt.Errorf("could not get proof for slot: %w", err)
+	}
+
+	// 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, nil, fmt.Errorf("could not get block header tree: %w", err)
+	}
+	blockHeaderProof, err := bhTree.Prove(int(generic.BeaconBlockHeaderStateRootGeneralizedIndex))
+	if err != nil {
+		return nil, nil, fmt.Errorf("could not get proof for block header: %w", err)
+	}
+
+	validatorProof := make([][]byte, 0, len(validatorStateProof.Hashes)+len(blockHeaderProof.Hashes))
+	validatorProof = append(validatorProof, validatorStateProof.Hashes...)
+	validatorProof = append(validatorProof, blockHeaderProof.Hashes...)
+
+	slotProof := make([][]byte, 0, len(slotStateProof.Hashes)+len(blockHeaderProof.Hashes))
+	slotProof = append(slotProof, slotStateProof.Hashes...)
+	slotProof = append(slotProof, blockHeaderProof.Hashes...)
+
+	return validatorProof, slotProof, nil
+}
+
 func (state *BeaconState) blockHeaderToStateProof(blockHeader *generic.BeaconBlockHeader) ([][]byte, error) {
 	generalizedIndex := generic.BeaconBlockHeaderStateRootGeneralizedIndex
 	root, err := blockHeader.GetTree()

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

@@ -186,6 +186,62 @@ func (state *BeaconState) ValidatorProof(index uint64) ([][]byte, error) {
 	return append(proof, blockHeaderProof.Hashes...), nil
 }
 
+// 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) {
+
+	if validatorIndex >= uint64(len(state.Validators)) {
+		return nil, nil, errors.New("validator index out of bounds")
+	}
+
+	stateTree, err := state.GetTree()
+	if err != nil {
+		return nil, nil, fmt.Errorf("could not get state tree: %w", err)
+	}
+
+	validatorGid := generic.GetGeneralizedIndexForValidator(validatorIndex, GetGeneralizedIndexForValidators())
+	validatorStateProof, err := stateTree.Prove(int(validatorGid))
+	if err != nil {
+		return nil, nil, fmt.Errorf("could not get proof for validator: %w", err)
+	}
+
+	// Sanity check that the proof leaf matches the expected validator
+	validatorHashTreeRoot, err := state.Validators[validatorIndex].HashTreeRoot()
+	if err != nil {
+		return nil, nil, fmt.Errorf("could not get hash tree root for validator: %w", err)
+	}
+	if !bytes.Equal(validatorStateProof.Leaf, validatorHashTreeRoot[:]) {
+		return nil, nil, fmt.Errorf("proof leaf does not match expected validator")
+	}
+
+	slotStateProof, err := stateTree.Prove(int(GetGeneralizedIndexForSlot()))
+	if err != nil {
+		return nil, nil, fmt.Errorf("could not get proof for slot: %w", err)
+	}
+
+	// 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, nil, fmt.Errorf("could not get block header tree: %w", err)
+	}
+	blockHeaderProof, err := bhTree.Prove(int(generic.BeaconBlockHeaderStateRootGeneralizedIndex))
+	if err != nil {
+		return nil, nil, fmt.Errorf("could not get proof for block header: %w", err)
+	}
+
+	validatorProof := make([][]byte, 0, len(validatorStateProof.Hashes)+len(blockHeaderProof.Hashes))
+	validatorProof = append(validatorProof, validatorStateProof.Hashes...)
+	validatorProof = append(validatorProof, blockHeaderProof.Hashes...)
+
+	slotProof := make([][]byte, 0, len(slotStateProof.Hashes)+len(blockHeaderProof.Hashes))
+	slotProof = append(slotProof, slotStateProof.Hashes...)
+	slotProof = append(slotProof, blockHeaderProof.Hashes...)
+
+	return validatorProof, slotProof, nil
+}
+
 func (state *BeaconState) blockHeaderToStateProof(blockHeader *generic.BeaconBlockHeader) ([][]byte, error) {
 	generalizedIndex := generic.BeaconBlockHeaderStateRootGeneralizedIndex
 	root, err := blockHeader.GetTree()

shared/types/eth2/types.go

@@ -23,6 +23,7 @@ type BeaconState interface {
 	GetSlot() uint64
 	ValidatorProof(index uint64) ([][]byte, error)
 	SlotProof(slot uint64) ([][]byte, error)
+	ValidatorAndSlotProof(validatorIndex uint64) (validatorProof [][]byte, slotProof [][]byte, err error)
 	HistoricalSummaryProof(slot uint64, capellaOffset uint64) ([][]byte, error)
 	HistoricalSummaryBlockRootProof(slot int) ([][]byte, error)
 	BlockRootProof(slot uint64) ([][]byte, error)