Add Paginated Attestation Pool to Prysm (#4827)

* added pagination to atts

* tests pass for atts

* add mock

* fix

* add to val

* fix

* add in proper mock

Co-authored-by: prylabs-bulldozer[bot] <58059840+prylabs-bulldozer[bot]@users.noreply.github.com>

WORKSPACE

@@ -1272,7 +1272,7 @@ go_repository(
 
 go_repository(
     name = "com_github_prysmaticlabs_ethereumapis",
-    commit = "a90fbe4a333e538a897fb6008b35271ef6b7b124",
+    commit = "6720aaf759152b73b1162d7c2156ecc4e788a8d3",
     importpath = "github.com/prysmaticlabs/ethereumapis",
     patch_args = ["-p1"],
     patches = [

beacon-chain/rpc/beacon/attestations.go

@@ -143,10 +143,32 @@ func (bs *Server) StreamAttestations(
 // attestations are processed and when they are no longer valid.
 // https://github.com/ethereum/eth2.0-specs/blob/dev/specs/core/0_beacon-chain.md#attestations
 func (bs *Server) AttestationPool(
-	ctx context.Context, _ *ptypes.Empty,
+	ctx context.Context, req *ethpb.AttestationPoolRequest,
 ) (*ethpb.AttestationPoolResponse, error) {
+	if int(req.PageSize) > flags.Get().MaxPageSize {
+		return nil, status.Errorf(
+			codes.InvalidArgument,
+			"Requested page size %d can not be greater than max size %d",
+			req.PageSize,
+			flags.Get().MaxPageSize,
+		)
+	}
 	atts := bs.Pool.AggregatedAttestations()
+	numAtts := len(atts)
+	if numAtts == 0 {
+		return &ethpb.AttestationPoolResponse{
+			Attestations:  make([]*ethpb.Attestation, 0),
+			TotalSize:     int32(0),
+			NextPageToken: strconv.Itoa(0),
+		}, nil
+	}
+	start, end, nextPageToken, err := pagination.StartAndEndPage(req.PageToken, int(req.PageSize), numAtts)
+	if err != nil {
+		return nil, status.Errorf(codes.Internal, "Could not paginate attestations: %v", err)
+	}
 	return &ethpb.AttestationPoolResponse{
-		Attestations: atts,
+		Attestations:  atts[start:end],
+		TotalSize:     int32(numAtts),
+		NextPageToken: nextPageToken,
 	}, nil
 }

beacon-chain/rpc/beacon/attestations_test.go

@@ -536,6 +536,143 @@ func TestServer_ListAttestations_Pagination_DefaultPageSize(t *testing.T) {
 	}
 }
 
+func TestServer_AttestationPool_Pagination_ExceedsMaxPageSize(t *testing.T) {
+	ctx := context.Background()
+	bs := &Server{}
+	exceedsMax := int32(flags.Get().MaxPageSize + 1)
+
+	wanted := fmt.Sprintf("Requested page size %d can not be greater than max size %d", exceedsMax, flags.Get().MaxPageSize)
+	req := &ethpb.AttestationPoolRequest{PageToken: strconv.Itoa(0), PageSize: exceedsMax}
+	if _, err := bs.AttestationPool(ctx, req); err != nil && !strings.Contains(err.Error(), wanted) {
+		t.Errorf("Expected error %v, received %v", wanted, err)
+	}
+}
+
+func TestServer_AttestationPool_Pagination_OutOfRange(t *testing.T) {
+	ctx := context.Background()
+	bs := &Server{
+		Pool: attestations.NewPool(),
+	}
+
+	atts := []*ethpb.Attestation{
+		{Data: &ethpb.AttestationData{Slot: 1}, AggregationBits: bitfield.Bitlist{0b1101}},
+		{Data: &ethpb.AttestationData{Slot: 2}, AggregationBits: bitfield.Bitlist{0b1101}},
+		{Data: &ethpb.AttestationData{Slot: 3}, AggregationBits: bitfield.Bitlist{0b1101}},
+	}
+	if err := bs.Pool.SaveAggregatedAttestations(atts); err != nil {
+		t.Fatal(err)
+	}
+
+	req := &ethpb.AttestationPoolRequest{
+		PageToken: strconv.Itoa(1),
+		PageSize:  100,
+	}
+	wanted := fmt.Sprintf("page start %d >= list %d", req.PageSize, len(atts))
+	if _, err := bs.AttestationPool(ctx, req); err != nil && !strings.Contains(err.Error(), wanted) {
+		t.Errorf("Expected error %v, received %v", wanted, err)
+	}
+}
+
+func TestServer_AttestationPool_Pagination_DefaultPageSize(t *testing.T) {
+	ctx := context.Background()
+	bs := &Server{
+		Pool: attestations.NewPool(),
+	}
+
+	atts := make([]*ethpb.Attestation, params.BeaconConfig().DefaultPageSize+1)
+	for i := 0; i < len(atts); i++ {
+		atts[i] = &ethpb.Attestation{
+			Data:            &ethpb.AttestationData{Slot: uint64(i)},
+			AggregationBits: bitfield.Bitlist{0b1101},
+		}
+	}
+	if err := bs.Pool.SaveAggregatedAttestations(atts); err != nil {
+		t.Fatal(err)
+	}
+
+	req := &ethpb.AttestationPoolRequest{}
+	res, err := bs.AttestationPool(ctx, req)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(res.Attestations) != params.BeaconConfig().DefaultPageSize {
+		t.Errorf(
+			"Wanted %d attestations in response, received %d",
+			params.BeaconConfig().DefaultPageSize,
+			len(res.Attestations),
+		)
+	}
+	if int(res.TotalSize) != params.BeaconConfig().DefaultPageSize+1 {
+		t.Errorf("Wanted total size %d, received %d", params.BeaconConfig().DefaultPageSize+1, res.TotalSize)
+	}
+}
+
+func TestServer_AttestationPool_Pagination_CustomPageSize(t *testing.T) {
+	ctx := context.Background()
+	bs := &Server{
+		Pool: attestations.NewPool(),
+	}
+
+	numAtts := 100
+	atts := make([]*ethpb.Attestation, numAtts)
+	for i := 0; i < len(atts); i++ {
+		atts[i] = &ethpb.Attestation{
+			Data:            &ethpb.AttestationData{Slot: uint64(i)},
+			AggregationBits: bitfield.Bitlist{0b1101},
+		}
+	}
+	if err := bs.Pool.SaveAggregatedAttestations(atts); err != nil {
+		t.Fatal(err)
+	}
+	tests := []struct {
+		req *ethpb.AttestationPoolRequest
+		res *ethpb.AttestationPoolResponse
+	}{
+		{
+			req: &ethpb.AttestationPoolRequest{
+				PageToken: strconv.Itoa(1),
+				PageSize:  3,
+			},
+			res: &ethpb.AttestationPoolResponse{
+				NextPageToken: "2",
+				TotalSize:     int32(numAtts),
+			},
+		},
+		{
+			req: &ethpb.AttestationPoolRequest{
+				PageToken: strconv.Itoa(3),
+				PageSize:  30,
+			},
+			res: &ethpb.AttestationPoolResponse{
+				NextPageToken: "",
+				TotalSize:     int32(numAtts),
+			},
+		},
+		{
+			req: &ethpb.AttestationPoolRequest{
+				PageToken: strconv.Itoa(0),
+				PageSize:  int32(numAtts),
+			},
+			res: &ethpb.AttestationPoolResponse{
+				NextPageToken: "1",
+				TotalSize:     int32(numAtts),
+			},
+		},
+	}
+	for _, tt := range tests {
+		res, err := bs.AttestationPool(ctx, tt.req)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if res.TotalSize != tt.res.TotalSize {
+			t.Errorf("Wanted total size %d, received %d", tt.res.TotalSize, res.TotalSize)
+		}
+		if res.NextPageToken != tt.res.NextPageToken {
+			t.Errorf("Wanted next page token %s, received %s", tt.res.NextPageToken, res.NextPageToken)
+		}
+	}
+}
+
 func TestServer_StreamAttestations_ContextCanceled(t *testing.T) {
 	db := dbTest.SetupDB(t)
 	defer dbTest.TeardownDB(t, db)

shared/mock/beacon_chain_service_mock.go

@@ -6,13 +6,12 @@ package mock
 
 import (
 	context "context"
-	reflect "reflect"
-
 	empty "github.com/gogo/protobuf/types"
 	gomock "github.com/golang/mock/gomock"
 	v1alpha1 "github.com/prysmaticlabs/ethereumapis/eth/v1alpha1"
 	grpc "google.golang.org/grpc"
 	metadata "google.golang.org/grpc/metadata"
+	reflect "reflect"
 )
 
 // MockBeaconChainClient is a mock of BeaconChainClient interface
@@ -39,7 +38,7 @@ func (m *MockBeaconChainClient) EXPECT() *MockBeaconChainClientMockRecorder {
 }
 
 // AttestationPool mocks base method
-func (m *MockBeaconChainClient) AttestationPool(arg0 context.Context, arg1 *empty.Empty, arg2 ...grpc.CallOption) (*v1alpha1.AttestationPoolResponse, error) {
+func (m *MockBeaconChainClient) AttestationPool(arg0 context.Context, arg1 *v1alpha1.AttestationPoolRequest, arg2 ...grpc.CallOption) (*v1alpha1.AttestationPoolResponse, error) {
 	m.ctrl.T.Helper()
 	varargs := []interface{}{arg0, arg1}
 	for _, a := range arg2 {

third_party/com_github_prysmaticlabs_ethereumapis-tags.patch

@@ -262,7 +262,7 @@ index 2ce5c34..4cbb276 100644
 +    bytes signature = 3 [(gogoproto.moretags) = "ssz-size:\"96\""];
  }
 diff --git a/eth/v1alpha1/beacon_chain.proto b/eth/v1alpha1/beacon_chain.proto
-index 1841b7d..c0d7506 100644
+index 8de1adb..ffcc8f4 100644
 --- a/eth/v1alpha1/beacon_chain.proto
 +++ b/eth/v1alpha1/beacon_chain.proto
 @@ -15,6 +15,7 @@ syntax = "proto3";