From 4258b32de7285f6bb85cc7e63eed08af4ffa68bb Mon Sep 17 00:00:00 2001
From: iamqizhao <toqizhao@gmail.com>
Date: Mon, 25 Jan 2016 11:47:04 -0800
Subject: [PATCH] Don't create buffer if compressor is nil

---
 call.go   |  6 +++++-
 server.go |  9 ++++++++-
 stream.go | 24 ++++++++++++++++++------
 3 files changed, 31 insertions(+), 8 deletions(-)

diff --git a/call.go b/call.go
index 3cc8f1fd0..f29396aea 100644
--- a/call.go
+++ b/call.go
@@ -81,7 +81,11 @@ func sendRequest(ctx context.Context, codec Codec, compressor Compressor, callHd
 			}
 		}
 	}()
-	outBuf, err := encode(codec, args, compressor, new(bytes.Buffer))
+	var cbuf *bytes.Buffer
+	if compressor != nil {
+		cbuf = new(bytes.Buffer)
+	}
+	outBuf, err := encode(codec, args, compressor, cbuf)
 	if err != nil {
 		return nil, transport.StreamErrorf(codes.Internal, "grpc: %v", err)
 	}
diff --git a/server.go b/server.go
index e846f3ba9..a7ff16c51 100644
--- a/server.go
+++ b/server.go
@@ -303,7 +303,11 @@ func (s *Server) Serve(lis net.Listener) error {
 }
 
 func (s *Server) sendResponse(t transport.ServerTransport, stream *transport.Stream, msg interface{}, cp Compressor, opts *transport.Options) error {
-	p, err := encode(s.opts.codec, msg, cp, new(bytes.Buffer))
+	var cbuf *bytes.Buffer
+	if cp != nil {
+		cbuf = new(bytes.Buffer)
+	}
+	p, err := encode(s.opts.codec, msg, cp, cbuf)
 	if err != nil {
 		// This typically indicates a fatal issue (e.g., memory
 		// corruption or hardware faults) the application program
@@ -457,6 +461,9 @@ func (s *Server) processStreamingRPC(t transport.ServerTransport, stream *transp
 		dg:     s.opts.dg,
 		trInfo: trInfo,
 	}
+	if cp != nil {
+		ss.cbuf = new(bytes.Buffer)
+	}
 	if trInfo != nil {
 		trInfo.tr.LazyLog(&trInfo.firstLine, false)
 		defer func() {
diff --git a/stream.go b/stream.go
index 1a8d0e459..63f934dda 100644
--- a/stream.go
+++ b/stream.go
@@ -124,6 +124,10 @@ func NewClientStream(ctx context.Context, desc *StreamDesc, cc *ClientConn, meth
 		dg:      cc.dopts.dg,
 		tracing: EnableTracing,
 	}
+	if cp != nil {
+		callHdr.SendCompress = cp.Type()
+		cs.cbuf = new(bytes.Buffer)
+	}
 	if cs.tracing {
 		cs.trInfo.tr = trace.New("grpc.Sent."+methodFamily(method), method)
 		cs.trInfo.firstLine.client = true
@@ -164,7 +168,7 @@ type clientStream struct {
 	desc  *StreamDesc
 	codec Codec
 	cp    Compressor
-	cbuf  bytes.Buffer
+	cbuf  *bytes.Buffer
 	dg    DecompressorGenerator
 
 	tracing bool // set to EnableTracing when the clientStream is created.
@@ -211,8 +215,12 @@ func (cs *clientStream) SendMsg(m interface{}) (err error) {
 		}
 		err = toRPCErr(err)
 	}()
-	out, err := encode(cs.codec, m, cs.cp, &cs.cbuf)
-	defer cs.cbuf.Reset()
+	out, err := encode(cs.codec, m, cs.cp, cs.cbuf)
+	defer func() {
+		if cs.cbuf != nil {
+			cs.cbuf.Reset()
+		}
+	}()
 	if err != nil {
 		return transport.StreamErrorf(codes.Internal, "grpc: %v", err)
 	}
@@ -326,7 +334,7 @@ type serverStream struct {
 	codec      Codec
 	cp         Compressor
 	dg         DecompressorGenerator
-	cbuf       bytes.Buffer
+	cbuf       *bytes.Buffer
 	statusCode codes.Code
 	statusDesc string
 	trInfo     *traceInfo
@@ -365,8 +373,12 @@ func (ss *serverStream) SendMsg(m interface{}) (err error) {
 			ss.mu.Unlock()
 		}
 	}()
-	out, err := encode(ss.codec, m, ss.cp, &ss.cbuf)
-	defer ss.cbuf.Reset()
+	out, err := encode(ss.codec, m, ss.cp, ss.cbuf)
+	defer func() {
+		if ss.cbuf != nil {
+			ss.cbuf.Reset()
+		}
+	}()
 	if err != nil {
 		err = transport.StreamErrorf(codes.Internal, "grpc: %v", err)
 		return err