Context Cancelation Handling for Hot States (#7596)

* attempt deadline for state info

* cancel check

beacon-chain/db/kv/blocks.go

@@ -288,6 +288,9 @@ func (s *Store) HighestSlotBlocksBelow(ctx context.Context, slot uint64) ([]*eth
 		// Iterate through the index, which is in byte sorted order.
 		c := bkt.Cursor()
 		for s, root := c.First(); s != nil; s, root = c.Next() {
+			if ctx.Err() != nil {
+				return ctx.Err()
+			}
 			key := bytesutil.BytesToUint64BigEndian(s)
 			if root == nil {
 				continue

beacon-chain/db/kv/state.go

@@ -296,6 +296,9 @@ func (s *Store) HighestSlotStatesBelow(ctx context.Context, slot uint64) ([]*sta
 		bkt := tx.Bucket(stateSlotIndicesBucket)
 		c := bkt.Cursor()
 		for s, root := c.First(); s != nil; s, root = c.Next() {
+			if ctx.Err() != nil {
+				return ctx.Err()
+			}
 			key := bytesutil.BytesToUint64BigEndian(s)
 			if root == nil {
 				continue

beacon-chain/rpc/beacon/validators.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"sort"
 	"strconv"
+	"time"
 
 	ptypes "github.com/gogo/protobuf/types"
 	ethpb "github.com/prysmaticlabs/ethereumapis/eth/v1alpha1"
@@ -20,12 +21,18 @@ import (
 	"google.golang.org/grpc/status"
 )
 
+// BalancesTimeout for gRPC requests to ListValidatorBalances.
+const BalancesTimeout = time.Second * 30
+
 // ListValidatorBalances retrieves the validator balances for a given set of public keys.
 // An optional Epoch parameter is provided to request historical validator balances from
 // archived, persistent data.
 func (bs *Server) ListValidatorBalances(
 	ctx context.Context,
-	req *ethpb.ListValidatorBalancesRequest) (*ethpb.ValidatorBalances, error) {
+	req *ethpb.ListValidatorBalancesRequest,
+) (*ethpb.ValidatorBalances, error) {
+	ctx, cancel := context.WithTimeout(ctx, BalancesTimeout)
+	defer cancel()
 	if int(req.PageSize) > cmd.Get().MaxRPCPageSize {
 		return nil, status.Errorf(codes.InvalidArgument, "Requested page size %d can not be greater than max size %d",
 			req.PageSize, cmd.Get().MaxRPCPageSize)

beacon-chain/state/stategen/replay.go

@@ -22,6 +22,9 @@ func (s *State) ReplayBlocks(ctx context.Context, state *stateTrie.BeaconState,
 	// The input block list is sorted in decreasing slots order.
 	if len(signed) > 0 {
 		for i := len(signed) - 1; i >= 0; i-- {
+			if ctx.Err() != nil {
+				return nil, ctx.Err()
+			}
 			if state.Slot() >= targetSlot {
 				break
 			}
@@ -69,6 +72,9 @@ func (s *State) LoadBlocks(ctx context.Context, startSlot, endSlot uint64, endBl
 	// Covers the edge case if there's multiple blocks on the same end slot,
 	// the end root may not be the last index in `blockRoots`.
 	for length >= 3 && blocks[length-1].Block.Slot == blocks[length-2].Block.Slot && blockRoots[length-1] != endBlockRoot {
+		if ctx.Err() != nil {
+			return nil, ctx.Err()
+		}
 		length--
 		if blockRoots[length-2] == endBlockRoot {
 			length--
@@ -83,6 +89,9 @@ func (s *State) LoadBlocks(ctx context.Context, startSlot, endSlot uint64, endBl
 	filteredBlocks := []*ethpb.SignedBeaconBlock{blocks[length-1]}
 	// Starting from second to last index because the last block is already in the filtered block list.
 	for i := length - 2; i >= 0; i-- {
+		if ctx.Err() != nil {
+			return nil, ctx.Err()
+		}
 		b := filteredBlocks[len(filteredBlocks)-1]
 		if bytesutil.ToBytes32(b.Block.ParentRoot) != blockRoots[i] {
 			continue