diff --git a/beacon-chain/db/kv/blocks.go b/beacon-chain/db/kv/blocks.go index 1f59b58a5e2..d5a76a50857 100644 --- a/beacon-chain/db/kv/blocks.go +++ b/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 diff --git a/beacon-chain/db/kv/state.go b/beacon-chain/db/kv/state.go index 2d458020ef6..2692962dbb0 100644 --- a/beacon-chain/db/kv/state.go +++ b/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 diff --git a/beacon-chain/rpc/beacon/validators.go b/beacon-chain/rpc/beacon/validators.go index 5641e972c09..86381a6b1d0 100644 --- a/beacon-chain/rpc/beacon/validators.go +++ b/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) diff --git a/beacon-chain/state/stategen/replay.go b/beacon-chain/state/stategen/replay.go index 3b8179cf9a3..faf7d912366 100644 --- a/beacon-chain/state/stategen/replay.go +++ b/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