From 297247d9155b598a11e6f3c04c3a0b19e538523d Mon Sep 17 00:00:00 2001 From: Raul Jordan Date: Tue, 11 Feb 2020 13:18:30 -0600 Subject: [PATCH] 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 | 2 +- beacon-chain/rpc/beacon/attestations.go | 26 +++- beacon-chain/rpc/beacon/attestations_test.go | 137 ++++++++++++++++++ shared/mock/beacon_chain_service_mock.go | 5 +- ...thub_prysmaticlabs_ethereumapis-tags.patch | 2 +- 5 files changed, 165 insertions(+), 7 deletions(-) diff --git a/WORKSPACE b/WORKSPACE index 03fcb9d3b..1f8d78fb5 100644 --- a/WORKSPACE +++ b/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 = [ diff --git a/beacon-chain/rpc/beacon/attestations.go b/beacon-chain/rpc/beacon/attestations.go index 3f1874327..5f22a4a9b 100644 --- a/beacon-chain/rpc/beacon/attestations.go +++ b/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 ðpb.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 ðpb.AttestationPoolResponse{ - Attestations: atts, + Attestations: atts[start:end], + TotalSize: int32(numAtts), + NextPageToken: nextPageToken, }, nil } diff --git a/beacon-chain/rpc/beacon/attestations_test.go b/beacon-chain/rpc/beacon/attestations_test.go index f4d2150b8..be3185325 100644 --- a/beacon-chain/rpc/beacon/attestations_test.go +++ b/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 := ðpb.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: ðpb.AttestationData{Slot: 1}, AggregationBits: bitfield.Bitlist{0b1101}}, + {Data: ðpb.AttestationData{Slot: 2}, AggregationBits: bitfield.Bitlist{0b1101}}, + {Data: ðpb.AttestationData{Slot: 3}, AggregationBits: bitfield.Bitlist{0b1101}}, + } + if err := bs.Pool.SaveAggregatedAttestations(atts); err != nil { + t.Fatal(err) + } + + req := ðpb.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] = ðpb.Attestation{ + Data: ðpb.AttestationData{Slot: uint64(i)}, + AggregationBits: bitfield.Bitlist{0b1101}, + } + } + if err := bs.Pool.SaveAggregatedAttestations(atts); err != nil { + t.Fatal(err) + } + + req := ðpb.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] = ðpb.Attestation{ + Data: ðpb.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: ðpb.AttestationPoolRequest{ + PageToken: strconv.Itoa(1), + PageSize: 3, + }, + res: ðpb.AttestationPoolResponse{ + NextPageToken: "2", + TotalSize: int32(numAtts), + }, + }, + { + req: ðpb.AttestationPoolRequest{ + PageToken: strconv.Itoa(3), + PageSize: 30, + }, + res: ðpb.AttestationPoolResponse{ + NextPageToken: "", + TotalSize: int32(numAtts), + }, + }, + { + req: ðpb.AttestationPoolRequest{ + PageToken: strconv.Itoa(0), + PageSize: int32(numAtts), + }, + res: ðpb.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) diff --git a/shared/mock/beacon_chain_service_mock.go b/shared/mock/beacon_chain_service_mock.go index 4bd5914c5..cc829f439 100644 --- a/shared/mock/beacon_chain_service_mock.go +++ b/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 { diff --git a/third_party/com_github_prysmaticlabs_ethereumapis-tags.patch b/third_party/com_github_prysmaticlabs_ethereumapis-tags.patch index 09bbea3b2..7d52016a6 100644 --- a/third_party/com_github_prysmaticlabs_ethereumapis-tags.patch +++ b/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";