From 47b25d4b974c2618e2a75b3e7502a345868a092f Mon Sep 17 00:00:00 2001
From: Alexander Polcyn <apolcyn@google.com>
Date: Wed, 21 Dec 2016 10:29:00 -0800
Subject: [PATCH] remove proto object allocs on streaming benchmark server

---
 benchmark/benchmark.go | 28 +++++++++++++++++++---------
 1 file changed, 19 insertions(+), 9 deletions(-)

diff --git a/benchmark/benchmark.go b/benchmark/benchmark.go
index d1143270..929c1445 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,12 @@ func (s *testServer) UnaryCall(ctx context.Context, in *testpb.SimpleRequest) (*
 }
 
 func (s *testServer) StreamingCall(stream testpb.BenchmarkService_StreamingCallServer) error {
+	response := new(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 +94,9 @@ 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(payload, in.ResponseType, int(in.ResponseSize))
+		response.Payload = payload
+		if err := stream.Send(response); err != nil {
 			return err
 		}
 	}