Add back archival endpoint GetValidatorBalances with fallback (#5620)

beacon-chain/rpc/beacon/attestations_test.go

@@ -612,6 +612,10 @@ func TestServer_ListIndexedAttestations_NewStateManagnmentDisabled(t *testing.T)
 func TestServer_ListIndexedAttestations_GenesisEpoch(t *testing.T) {
 	params.OverrideBeaconConfig(params.MainnetConfig())
 	defer params.OverrideBeaconConfig(params.MinimalSpecConfig())
+
+	featureconfig.Init(&featureconfig.Flags{NewStateMgmt: true})
+	defer featureconfig.Init(&featureconfig.Flags{NewStateMgmt: false})
+
 	cfg := assertNewStateMgmtIsEnabled()
 	defer featureconfig.Init(cfg)
 	db := dbTest.SetupDB(t)

beacon-chain/rpc/beacon/validators.go

@@ -31,6 +31,135 @@ func (bs *Server) ListValidatorBalances(
 			req.PageSize, flags.Get().MaxPageSize)
 	}
 
+	if !featureconfig.Get().NewStateMgmt {
+		return bs.listValidatorsBalancesUsingOldArchival(ctx, req)
+	}
+
+	currentEpoch := helpers.SlotToEpoch(bs.GenesisTimeFetcher.CurrentSlot())
+	requestedEpoch := currentEpoch
+	switch q := req.QueryFilter.(type) {
+	case *ethpb.ListValidatorBalancesRequest_Epoch:
+		requestedEpoch = q.Epoch
+	case *ethpb.ListValidatorBalancesRequest_Genesis:
+		requestedEpoch = 0
+	default:
+		requestedEpoch = currentEpoch
+	}
+
+	if requestedEpoch > currentEpoch {
+		return nil, status.Errorf(
+			codes.InvalidArgument,
+			"Cannot retrieve information about an epoch in the future, current epoch %d, requesting %d",
+			currentEpoch,
+			requestedEpoch,
+		)
+	}
+	res := make([]*ethpb.ValidatorBalances_Balance, 0)
+	filtered := map[uint64]bool{} // Track filtered validators to prevent duplication in the response.
+
+	requestedState, err := bs.StateGen.StateBySlot(ctx, helpers.StartSlot(requestedEpoch))
+	if err != nil {
+		return nil, status.Errorf(codes.Internal, "Could not get state")
+	}
+
+	validators := requestedState.Validators()
+	balances := requestedState.Balances()
+	balancesCount := len(balances)
+	for _, pubKey := range req.PublicKeys {
+		// Skip empty public key.
+		if len(pubKey) == 0 {
+			continue
+		}
+		pubkeyBytes := bytesutil.ToBytes48(pubKey)
+		index, ok := requestedState.ValidatorIndexByPubkey(pubkeyBytes)
+		if !ok {
+			return nil, status.Errorf(codes.NotFound, "Could not find validator index for public key %#x", pubkeyBytes)
+		}
+
+		filtered[index] = true
+
+		if int(index) >= len(balances) {
+			return nil, status.Errorf(codes.OutOfRange, "Validator index %d >= balance list %d",
+				index, len(balances))
+		}
+
+		res = append(res, &ethpb.ValidatorBalances_Balance{
+			PublicKey: pubKey,
+			Index:     index,
+			Balance:   balances[index],
+		})
+		balancesCount = len(res)
+	}
+
+	for _, index := range req.Indices {
+		if int(index) >= len(balances) {
+			return nil, status.Errorf(codes.OutOfRange, "Validator index %d >= balance list %d",
+				index, len(balances))
+		}
+
+		if !filtered[index] {
+			res = append(res, &ethpb.ValidatorBalances_Balance{
+				PublicKey: validators[index].PublicKey,
+				Index:     index,
+				Balance:   balances[index],
+			})
+		}
+		balancesCount = len(res)
+	}
+	// Depending on the indices and public keys given, results might not be sorted.
+	sort.Slice(res, func(i, j int) bool {
+		return res[i].Index < res[j].Index
+	})
+
+	// If there are no balances, we simply return a response specifying this.
+	// Otherwise, attempting to paginate 0 balances below would result in an error.
+	if balancesCount == 0 {
+		return &ethpb.ValidatorBalances{
+			Epoch:         requestedEpoch,
+			Balances:      make([]*ethpb.ValidatorBalances_Balance, 0),
+			TotalSize:     int32(0),
+			NextPageToken: strconv.Itoa(0),
+		}, nil
+	}
+
+	start, end, nextPageToken, err := pagination.StartAndEndPage(req.PageToken, int(req.PageSize), balancesCount)
+	if err != nil {
+		return nil, status.Errorf(
+			codes.Internal,
+			"Could not paginate results: %v",
+			err,
+		)
+	}
+
+	if len(req.Indices) == 0 && len(req.PublicKeys) == 0 {
+		// Return everything.
+		for i := start; i < end; i++ {
+			pubkey := requestedState.PubkeyAtIndex(uint64(i))
+			res = append(res, &ethpb.ValidatorBalances_Balance{
+				PublicKey: pubkey[:],
+				Index:     uint64(i),
+				Balance:   balances[i],
+			})
+		}
+		return &ethpb.ValidatorBalances{
+			Epoch:         requestedEpoch,
+			Balances:      res,
+			TotalSize:     int32(balancesCount),
+			NextPageToken: nextPageToken,
+		}, nil
+	}
+
+	return &ethpb.ValidatorBalances{
+		Epoch:         requestedEpoch,
+		Balances:      res[start:end],
+		TotalSize:     int32(balancesCount),
+		NextPageToken: nextPageToken,
+	}, nil
+}
+
+func (bs *Server) listValidatorsBalancesUsingOldArchival(
+	ctx context.Context,
+	req *ethpb.ListValidatorBalancesRequest) (*ethpb.ValidatorBalances, error) {
 	res := make([]*ethpb.ValidatorBalances_Balance, 0)
 	filtered := map[uint64]bool{} // Track filtered validators to prevent duplication in the response.