Fix GCS

Signed-off-by: Olivier Gambier <olivier@docker.com>

vendor/google.golang.org/grpc/stream.go

@@ -34,8 +34,10 @@
 package grpc
 
 import (
+	"bytes"
 	"errors"
 	"io"
+	"sync"
 	"time"
 
 	"golang.org/x/net/context"
@@ -69,7 +71,7 @@ type Stream interface {
 	SendMsg(m interface{}) error
 	// RecvMsg blocks until it receives a message or the stream is
 	// done. On client side, it returns io.EOF when the stream is done. On
-	// any other error, it aborts the streama nd returns an RPC status. On
+	// any other error, it aborts the stream and returns an RPC status. On
 	// server side, it simply returns the error to the caller.
 	RecvMsg(m interface{}) error
 }
@@ -95,45 +97,84 @@ type ClientStream interface {
 // NewClientStream creates a new Stream for the client side. This is called
 // by generated code.
 func NewClientStream(ctx context.Context, desc *StreamDesc, cc *ClientConn, method string, opts ...CallOption) (ClientStream, error) {
+	var (
+		t   transport.ClientTransport
+		err error
+	)
+	t, err = cc.dopts.picker.Pick(ctx)
+	if err != nil {
+		return nil, toRPCErr(err)
+	}
 	// TODO(zhaoq): CallOption is omitted. Add support when it is needed.
 	callHdr := &transport.CallHdr{
 		Host:   cc.authority,
 		Method: method,
+		Flush:  desc.ServerStreams && desc.ClientStreams,
+	}
+	if cc.dopts.cp != nil {
+		callHdr.SendCompress = cc.dopts.cp.Type()
 	}
 	cs := &clientStream{
-		desc:  desc,
-		codec: cc.dopts.codec,
+		desc:    desc,
+		codec:   cc.dopts.codec,
+		cp:      cc.dopts.cp,
+		dc:      cc.dopts.dc,
+		tracing: EnableTracing,
 	}
-	if EnableTracing {
-		cs.traceInfo.tr = trace.New("Sent."+methodFamily(method), method)
-		cs.traceInfo.firstLine.client = true
+	if cc.dopts.cp != nil {
+		callHdr.SendCompress = cc.dopts.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
 		if deadline, ok := ctx.Deadline(); ok {
-			cs.traceInfo.firstLine.deadline = deadline.Sub(time.Now())
+			cs.trInfo.firstLine.deadline = deadline.Sub(time.Now())
 		}
-		cs.traceInfo.tr.LazyLog(&cs.traceInfo.firstLine, false)
-	}
-	t, _, err := cc.wait(ctx, 0)
-	if err != nil {
-		return nil, toRPCErr(err)
+		cs.trInfo.tr.LazyLog(&cs.trInfo.firstLine, false)
+		ctx = trace.NewContext(ctx, cs.trInfo.tr)
 	}
 	s, err := t.NewStream(ctx, callHdr)
 	if err != nil {
+		cs.finish(err)
 		return nil, toRPCErr(err)
 	}
 	cs.t = t
 	cs.s = s
-	cs.p = &parser{s: s}
+	cs.p = &parser{r: s}
+	// Listen on ctx.Done() to detect cancellation when there is no pending
+	// I/O operations on this stream.
+	go func() {
+		select {
+		case <-t.Error():
+			// Incur transport error, simply exit.
+		case <-s.Context().Done():
+			err := s.Context().Err()
+			cs.finish(err)
+			cs.closeTransportStream(transport.ContextErr(err))
+		}
+	}()
 	return cs, nil
 }
 
 // clientStream implements a client side Stream.
 type clientStream struct {
-	t         transport.ClientTransport
-	s         *transport.Stream
-	p         *parser
-	desc      *StreamDesc
-	codec     Codec
-	traceInfo traceInfo
+	t     transport.ClientTransport
+	s     *transport.Stream
+	p     *parser
+	desc  *StreamDesc
+	codec Codec
+	cp    Compressor
+	cbuf  *bytes.Buffer
+	dc    Decompressor
+
+	tracing bool // set to EnableTracing when the clientStream is created.
+
+	mu     sync.Mutex
+	closed bool
+	// trInfo.tr is set when the clientStream is created (if EnableTracing is true),
+	// and is set to nil when the clientStream's finish method is called.
+	trInfo traceInfo
 }
 
 func (cs *clientStream) Context() context.Context {
@@ -144,7 +185,7 @@ func (cs *clientStream) Header() (metadata.MD, error) {
 	m, err := cs.s.Header()
 	if err != nil {
 		if _, ok := err.(transport.ConnectionError); !ok {
-			cs.t.CloseStream(cs.s, err)
+			cs.closeTransportStream(err)
 		}
 	}
 	return m, err
@@ -155,16 +196,31 @@ func (cs *clientStream) Trailer() metadata.MD {
 }
 
 func (cs *clientStream) SendMsg(m interface{}) (err error) {
+	if cs.tracing {
+		cs.mu.Lock()
+		if cs.trInfo.tr != nil {
+			cs.trInfo.tr.LazyLog(&payload{sent: true, msg: m}, true)
+		}
+		cs.mu.Unlock()
+	}
 	defer func() {
+		if err != nil {
+			cs.finish(err)
+		}
 		if err == nil || err == io.EOF {
 			return
 		}
 		if _, ok := err.(transport.ConnectionError); !ok {
-			cs.t.CloseStream(cs.s, err)
+			cs.closeTransportStream(err)
 		}
 		err = toRPCErr(err)
 	}()
-	out, err := encode(cs.codec, m, compressionNone)
+	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)
 	}
@@ -172,29 +228,33 @@ func (cs *clientStream) SendMsg(m interface{}) (err error) {
 }
 
 func (cs *clientStream) RecvMsg(m interface{}) (err error) {
-	err = recv(cs.p, cs.codec, m)
+	err = recv(cs.p, cs.codec, cs.s, cs.dc, m)
 	defer func() {
 		// err != nil indicates the termination of the stream.
-		if EnableTracing && err != nil {
-			if err != io.EOF {
-				cs.traceInfo.tr.LazyLog(&fmtStringer{"%v", []interface{}{err}}, true)
-				cs.traceInfo.tr.SetError()
-			}
-			cs.traceInfo.tr.Finish()
+		if err != nil {
+			cs.finish(err)
 		}
 	}()
 	if err == nil {
+		if cs.tracing {
+			cs.mu.Lock()
+			if cs.trInfo.tr != nil {
+				cs.trInfo.tr.LazyLog(&payload{sent: false, msg: m}, true)
+			}
+			cs.mu.Unlock()
+		}
 		if !cs.desc.ClientStreams || cs.desc.ServerStreams {
 			return
 		}
 		// Special handling for client streaming rpc.
-		err = recv(cs.p, cs.codec, m)
-		cs.t.CloseStream(cs.s, err)
+		err = recv(cs.p, cs.codec, cs.s, cs.dc, m)
+		cs.closeTransportStream(err)
 		if err == nil {
 			return toRPCErr(errors.New("grpc: client streaming protocol violation: get <nil>, want <EOF>"))
 		}
 		if err == io.EOF {
 			if cs.s.StatusCode() == codes.OK {
+				cs.finish(err)
 				return nil
 			}
 			return Errorf(cs.s.StatusCode(), cs.s.StatusDesc())
@@ -202,7 +262,7 @@ func (cs *clientStream) RecvMsg(m interface{}) (err error) {
 		return toRPCErr(err)
 	}
 	if _, ok := err.(transport.ConnectionError); !ok {
-		cs.t.CloseStream(cs.s, err)
+		cs.closeTransportStream(err)
 	}
 	if err == io.EOF {
 		if cs.s.StatusCode() == codes.OK {
@@ -216,16 +276,50 @@ func (cs *clientStream) RecvMsg(m interface{}) (err error) {
 
 func (cs *clientStream) CloseSend() (err error) {
 	err = cs.t.Write(cs.s, nil, &transport.Options{Last: true})
+	defer func() {
+		if err != nil {
+			cs.finish(err)
+		}
+	}()
 	if err == nil || err == io.EOF {
 		return
 	}
 	if _, ok := err.(transport.ConnectionError); !ok {
-		cs.t.CloseStream(cs.s, err)
+		cs.closeTransportStream(err)
 	}
 	err = toRPCErr(err)
 	return
 }
 
+func (cs *clientStream) closeTransportStream(err error) {
+	cs.mu.Lock()
+	if cs.closed {
+		cs.mu.Unlock()
+		return
+	}
+	cs.closed = true
+	cs.mu.Unlock()
+	cs.t.CloseStream(cs.s, err)
+}
+
+func (cs *clientStream) finish(err error) {
+	if !cs.tracing {
+		return
+	}
+	cs.mu.Lock()
+	defer cs.mu.Unlock()
+	if cs.trInfo.tr != nil {
+		if err == nil || err == io.EOF {
+			cs.trInfo.tr.LazyPrintf("RPC: [OK]")
+		} else {
+			cs.trInfo.tr.LazyPrintf("RPC: [%v]", err)
+			cs.trInfo.tr.SetError()
+		}
+		cs.trInfo.tr.Finish()
+		cs.trInfo.tr = nil
+	}
+}
+
 // ServerStream defines the interface a server stream has to satisfy.
 type ServerStream interface {
 	// SendHeader sends the header metadata. It should not be called
@@ -244,8 +338,14 @@ type serverStream struct {
 	s          *transport.Stream
 	p          *parser
 	codec      Codec
+	cp         Compressor
+	dc         Decompressor
+	cbuf       *bytes.Buffer
 	statusCode codes.Code
 	statusDesc string
+	trInfo     *traceInfo
+
+	mu sync.Mutex // protects trInfo.tr after the service handler runs.
 }
 
 func (ss *serverStream) Context() context.Context {
@@ -264,8 +364,27 @@ func (ss *serverStream) SetTrailer(md metadata.MD) {
 	return
 }
 
-func (ss *serverStream) SendMsg(m interface{}) error {
-	out, err := encode(ss.codec, m, compressionNone)
+func (ss *serverStream) SendMsg(m interface{}) (err error) {
+	defer func() {
+		if ss.trInfo != nil {
+			ss.mu.Lock()
+			if ss.trInfo.tr != nil {
+				if err == nil {
+					ss.trInfo.tr.LazyLog(&payload{sent: true, msg: m}, true)
+				} else {
+					ss.trInfo.tr.LazyLog(&fmtStringer{"%v", []interface{}{err}}, true)
+					ss.trInfo.tr.SetError()
+				}
+			}
+			ss.mu.Unlock()
+		}
+	}()
+	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
@@ -273,6 +392,20 @@ func (ss *serverStream) SendMsg(m interface{}) error {
 	return ss.t.Write(ss.s, out, &transport.Options{Last: false})
 }
 
-func (ss *serverStream) RecvMsg(m interface{}) error {
-	return recv(ss.p, ss.codec, m)
+func (ss *serverStream) RecvMsg(m interface{}) (err error) {
+	defer func() {
+		if ss.trInfo != nil {
+			ss.mu.Lock()
+			if ss.trInfo.tr != nil {
+				if err == nil {
+					ss.trInfo.tr.LazyLog(&payload{sent: false, msg: m}, true)
+				} else if err != io.EOF {
+					ss.trInfo.tr.LazyLog(&fmtStringer{"%v", []interface{}{err}}, true)
+					ss.trInfo.tr.SetError()
+				}
+			}
+			ss.mu.Unlock()
+		}
+	}()
+	return recv(ss.p, ss.codec, ss.s, ss.dc, m)
 }