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

412 lines
11 KiB
Go
Raw Normal View History

/*
*
* Copyright 2014, Google Inc.
* All rights reserved.
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are
* met:
*
* * Redistributions of source code must retain the above copyright
* notice, this list of conditions and the following disclaimer.
* * Redistributions in binary form must reproduce the above
* copyright notice, this list of conditions and the following disclaimer
* in the documentation and/or other materials provided with the
* distribution.
* * Neither the name of Google Inc. nor the names of its
* contributors may be used to endorse or promote products derived from
* this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
* "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
* LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
* A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
* OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
* SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
* LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
* DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
* THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
* (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
*/
package grpc
import (
"bytes"
"errors"
"io"
"sync"
"time"
"golang.org/x/net/context"
"golang.org/x/net/trace"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/metadata"
"google.golang.org/grpc/transport"
)
type streamHandler func(srv interface{}, stream ServerStream) error
// StreamDesc represents a streaming RPC service's method specification.
type StreamDesc struct {
StreamName string
Handler streamHandler
// At least one of these is true.
ServerStreams bool
ClientStreams bool
}
// Stream defines the common interface a client or server stream has to satisfy.
type Stream interface {
// Context returns the context for this stream.
Context() context.Context
// SendMsg blocks until it sends m, the stream is done or the stream
// breaks.
// On error, it aborts the stream and returns an RPC status on client
// side. On server side, it simply returns the error to the caller.
// SendMsg is called by generated code.
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 stream and returns an RPC status. On
// server side, it simply returns the error to the caller.
RecvMsg(m interface{}) error
}
// ClientStream defines the interface a client stream has to satify.
type ClientStream interface {
// Header returns the header metedata received from the server if there
// is any. It blocks if the metadata is not ready to read.
Header() (metadata.MD, error)
// Trailer returns the trailer metadata from the server. It must be called
// after stream.Recv() returns non-nil error (including io.EOF) for
// bi-directional streaming and server streaming or stream.CloseAndRecv()
// returns for client streaming in order to receive trailer metadata if
// present. Otherwise, it could returns an empty MD even though trailer
// is present.
Trailer() metadata.MD
// CloseSend closes the send direction of the stream. It closes the stream
// when non-nil error is met.
CloseSend() error
Stream
}
// 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,
cp: cc.dopts.cp,
dc: cc.dopts.dc,
tracing: EnableTracing,
}
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.trInfo.firstLine.deadline = deadline.Sub(time.Now())
}
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{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
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 {
return cs.s.Context()
}
func (cs *clientStream) Header() (metadata.MD, error) {
m, err := cs.s.Header()
if err != nil {
if _, ok := err.(transport.ConnectionError); !ok {
cs.closeTransportStream(err)
}
}
return m, err
}
func (cs *clientStream) Trailer() metadata.MD {
return cs.s.Trailer()
}
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.closeTransportStream(err)
}
err = toRPCErr(err)
}()
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)
}
return cs.t.Write(cs.s, out, &transport.Options{Last: false})
}
func (cs *clientStream) RecvMsg(m interface{}) (err error) {
err = recv(cs.p, cs.codec, cs.s, cs.dc, m)
defer func() {
// err != nil indicates the termination of the stream.
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, 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())
}
return toRPCErr(err)
}
if _, ok := err.(transport.ConnectionError); !ok {
cs.closeTransportStream(err)
}
if err == io.EOF {
if cs.s.StatusCode() == codes.OK {
// Returns io.EOF to indicate the end of the stream.
return
}
return Errorf(cs.s.StatusCode(), cs.s.StatusDesc())
}
return toRPCErr(err)
}
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.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
// after SendProto. It fails if called multiple times or if
// called after SendProto.
SendHeader(metadata.MD) error
// SetTrailer sets the trailer metadata which will be sent with the
// RPC status.
SetTrailer(metadata.MD)
Stream
}
// serverStream implements a server side Stream.
type serverStream struct {
t transport.ServerTransport
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 {
return ss.s.Context()
}
func (ss *serverStream) SendHeader(md metadata.MD) error {
return ss.t.WriteHeader(ss.s, md)
}
func (ss *serverStream) SetTrailer(md metadata.MD) {
if md.Len() == 0 {
return
}
ss.s.SetTrailer(md)
return
}
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
}
return ss.t.Write(ss.s, out, &transport.Options{Last: false})
}
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)
}