diff --git a/benchmark/benchmark.go b/benchmark/benchmark.go index d1143270..fc3304ac 100644 --- a/benchmark/benchmark.go +++ b/benchmark/benchmark.go @@ -47,7 +47,8 @@ import ( "google.golang.org/grpc/grpclog" ) -func newPayload(t testpb.PayloadType, size int) *testpb.Payload { +// Allows reuse of the same testpb.Payload object. +func setPayload(p *testpb.Payload, t testpb.PayloadType, size int) { if size < 0 { grpclog.Fatalf("Requested a response with invalid length %d", size) } @@ -59,10 +60,15 @@ func newPayload(t testpb.PayloadType, size int) *testpb.Payload { default: grpclog.Fatalf("Unsupported payload type: %d", t) } - return &testpb.Payload{ - Type: t, - Body: body, - } + p.Type = t + p.Body = body + return +} + +func newPayload(t testpb.PayloadType, size int) *testpb.Payload { + p := new(testpb.Payload) + setPayload(p, t, size) + return p } type testServer struct { @@ -75,8 +81,13 @@ func (s *testServer) UnaryCall(ctx context.Context, in *testpb.SimpleRequest) (* } func (s *testServer) StreamingCall(stream testpb.BenchmarkService_StreamingCallServer) error { + response := &testpb.SimpleResponse{ + Payload: new(testpb.Payload), + } + in := new(testpb.SimpleRequest) for { - in, err := stream.Recv() + // use ServerStream directly to reuse the same testpb.SimpleRequest object + err := stream.(grpc.ServerStream).RecvMsg(in) if err == io.EOF { // read done. return nil @@ -84,9 +95,8 @@ func (s *testServer) StreamingCall(stream testpb.BenchmarkService_StreamingCallS if err != nil { return err } - if err := stream.Send(&testpb.SimpleResponse{ - Payload: newPayload(in.ResponseType, int(in.ResponseSize)), - }); err != nil { + setPayload(response.Payload, in.ResponseType, int(in.ResponseSize)) + if err := stream.Send(response); err != nil { return err } }