3pp-mirror/pomerium
1
0
Fork 0
mirror of https://github.com/pomerium/pomerium.git synced 2025-05-18 03:27:16 +02:00
Code Issues Projects Releases Packages Wiki Activity

replace GetAllPages with InitialSync, improve merge performance (#1624)

Browse source 
* replace GetAllPages with InitialSync, improve merge performance

* fmt proto

* add test for base64 function

* add sync test

* go mod tidy

Co-authored-by: Bobby DeSimone <bobbydesimone@gmail.com>
This commit is contained in:
Caleb Doxsey 2020-11-30 12:21:44 -07:00 • committed by GitHub
parent ba55fec67c
commit aad8ac2e61
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
10 changed files with 298 additions and 214 deletions
Download patch file Download diff file Expand all files Collapse all files

151
pkg/grpc/databroker/databroker_test.go
View file

@ -3,16 +3,11 @@ package databroker
import (
"context"
"net"
"strconv"
"testing"
"time"
"github.com/stretchr/testify/assert"
"google.golang.org/grpc"
"google.golang.org/grpc/codes"
status "google.golang.org/grpc/status"
"google.golang.org/protobuf/types/known/anypb"
"google.golang.org/protobuf/types/known/wrapperspb"
)
func TestApplyOffsetAndLimit(t *testing.T) {
@ -54,103 +49,61 @@ func TestApplyOffsetAndLimit(t *testing.T) {
}
}
func TestInitialSync(t *testing.T) {
ctx, clearTimeout := context.WithTimeout(context.Background(), time.Second*10)
defer clearTimeout()
li, err := net.Listen("tcp", "127.0.0.1:0")
if !assert.NoError(t, err) {
return
}
defer li.Close()
r1 := new(Record)
r2 := new(Record)
r3 := new(Record)
m := &mockServer{
sync: func(req *SyncRequest, stream DataBrokerService_SyncServer) error {
assert.Equal(t, true, req.GetNoWait())
stream.Send(&SyncResponse{
ServerVersion: "a",
Records: []*Record{r1, r2},
})
stream.Send(&SyncResponse{
ServerVersion: "b",
Records: []*Record{r3},
})
return nil
},
}
srv := grpc.NewServer()
RegisterDataBrokerServiceServer(srv, m)
go srv.Serve(li)
cc, err := grpc.Dial(li.Addr().String(), grpc.WithInsecure())
if !assert.NoError(t, err) {
return
}
defer cc.Close()
c := NewDataBrokerServiceClient(cc)
res, err := InitialSync(ctx, c, &SyncRequest{
Type: "TEST",
})
assert.NoError(t, err)
assert.Equal(t, "b", res.GetServerVersion())
assert.Equal(t, []*Record{r1, r2, r3}, res.GetRecords())
}
type mockServer struct {
DataBrokerServiceServer
getAll func(context.Context, *GetAllRequest) (*GetAllResponse, error)
sync func(*SyncRequest, DataBrokerService_SyncServer) error
}
func (m *mockServer) GetAll(ctx context.Context, req *GetAllRequest) (*GetAllResponse, error) {
return m.getAll(ctx, req)
}
func TestGetAllPages(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), time.Second*10)
defer cancel()
t.Run("resource exhausted", func(t *testing.T) {
li, err := net.Listen("tcp", "127.0.0.1:0")
if !assert.NoError(t, err) {
return
}
defer li.Close()
m := &mockServer{
getAll: func(ctx context.Context, req *GetAllRequest) (*GetAllResponse, error) {
any, _ := anypb.New(wrapperspb.String("TEST"))
var records []*Record
for i := 0; i < 1000000; i++ {
records = append(records, &Record{
Type: req.GetType(),
Data: any,
})
}
return &GetAllResponse{Records: records}, nil
},
}
srv := grpc.NewServer()
RegisterDataBrokerServiceServer(srv, m)
go srv.Serve(li)
cc, err := grpc.Dial(li.Addr().String(), grpc.WithInsecure())
if !assert.NoError(t, err) {
return
}
defer cc.Close()
c := NewDataBrokerServiceClient(cc)
res, err := GetAllPages(ctx, c, &GetAllRequest{
Type: "TEST",
})
assert.Error(t, err)
assert.Equal(t, codes.ResourceExhausted, status.Code(err))
assert.Nil(t, res)
})
t.Run("with paging", func(t *testing.T) {
li, err := net.Listen("tcp", "127.0.0.1:0")
if !assert.NoError(t, err) {
return
}
defer li.Close()
m := &mockServer{
getAll: func(ctx context.Context, req *GetAllRequest) (*GetAllResponse, error) {
pageToken, _ := strconv.Atoi(req.GetPageToken())
any, _ := anypb.New(wrapperspb.String("TEST"))
var records []*Record
for i := pageToken; i < pageToken+10000 && i < 1000000; i++ {
records = append(records, &Record{
Type: req.GetType(),
Data: any,
})
}
if len(records) == 0 {
return &GetAllResponse{}, nil
}
return &GetAllResponse{Records: records, NextPageToken: strconv.Itoa(pageToken + 10000)}, nil
},
}
srv := grpc.NewServer()
RegisterDataBrokerServiceServer(srv, m)
go srv.Serve(li)
cc, err := grpc.Dial(li.Addr().String(), grpc.WithInsecure())
if !assert.NoError(t, err) {
return
}
defer cc.Close()
c := NewDataBrokerServiceClient(cc)
res, err := GetAllPages(ctx, c, &GetAllRequest{
Type: "TEST",
})
assert.NoError(t, err)
assert.NotEqual(t, codes.ResourceExhausted, status.Code(err))
assert.Len(t, res.GetRecords(), 1000000)
})
func (m *mockServer) Sync(req *SyncRequest, stream DataBrokerService_SyncServer) error {
return m.sync(req, stream)
}