Source file src/google.golang.org/grpc/stream.go

Documentation: google.golang.org/grpc

     1  /*
     2   *
     3   * Copyright 2014 gRPC authors.
     4   *
     5   * Licensed under the Apache License, Version 2.0 (the "License");
     6   * you may not use this file except in compliance with the License.
     7   * You may obtain a copy of the License at
     8   *
     9   *     http://www.apache.org/licenses/LICENSE-2.0
    10   *
    11   * Unless required by applicable law or agreed to in writing, software
    12   * distributed under the License is distributed on an "AS IS" BASIS,
    13   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14   * See the License for the specific language governing permissions and
    15   * limitations under the License.
    16   *
    17   */
    18  
    19  package grpc
    20  
    21  import (
    22  	"context"
    23  	"errors"
    24  	"io"
    25  	"math"
    26  	"strconv"
    27  	"sync"
    28  	"time"
    29  
    30  	"google.golang.org/grpc/balancer"
    31  	"google.golang.org/grpc/codes"
    32  	"google.golang.org/grpc/encoding"
    33  	"google.golang.org/grpc/internal"
    34  	"google.golang.org/grpc/internal/balancerload"
    35  	"google.golang.org/grpc/internal/binarylog"
    36  	"google.golang.org/grpc/internal/channelz"
    37  	"google.golang.org/grpc/internal/grpcrand"
    38  	"google.golang.org/grpc/internal/grpcutil"
    39  	imetadata "google.golang.org/grpc/internal/metadata"
    40  	iresolver "google.golang.org/grpc/internal/resolver"
    41  	"google.golang.org/grpc/internal/serviceconfig"
    42  	istatus "google.golang.org/grpc/internal/status"
    43  	"google.golang.org/grpc/internal/transport"
    44  	"google.golang.org/grpc/metadata"
    45  	"google.golang.org/grpc/peer"
    46  	"google.golang.org/grpc/stats"
    47  	"google.golang.org/grpc/status"
    48  )
    49  
    50  var metadataFromOutgoingContextRaw = internal.FromOutgoingContextRaw.(func(context.Context) (metadata.MD, [][]string, bool))
    51  
    52  // StreamHandler defines the handler called by gRPC server to complete the
    53  // execution of a streaming RPC.
    54  //
    55  // If a StreamHandler returns an error, it should either be produced by the
    56  // status package, or be one of the context errors. Otherwise, gRPC will use
    57  // codes.Unknown as the status code and err.Error() as the status message of the
    58  // RPC.
    59  type StreamHandler func(srv any, stream ServerStream) error
    60  
    61  // StreamDesc represents a streaming RPC service's method specification.  Used
    62  // on the server when registering services and on the client when initiating
    63  // new streams.
    64  type StreamDesc struct {
    65  	// StreamName and Handler are only used when registering handlers on a
    66  	// server.
    67  	StreamName string        // the name of the method excluding the service
    68  	Handler    StreamHandler // the handler called for the method
    69  
    70  	// ServerStreams and ClientStreams are used for registering handlers on a
    71  	// server as well as defining RPC behavior when passed to NewClientStream
    72  	// and ClientConn.NewStream.  At least one must be true.
    73  	ServerStreams bool // indicates the server can perform streaming sends
    74  	ClientStreams bool // indicates the client can perform streaming sends
    75  }
    76  
    77  // Stream defines the common interface a client or server stream has to satisfy.
    78  //
    79  // Deprecated: See ClientStream and ServerStream documentation instead.
    80  type Stream interface {
    81  	// Deprecated: See ClientStream and ServerStream documentation instead.
    82  	Context() context.Context
    83  	// Deprecated: See ClientStream and ServerStream documentation instead.
    84  	SendMsg(m any) error
    85  	// Deprecated: See ClientStream and ServerStream documentation instead.
    86  	RecvMsg(m any) error
    87  }
    88  
    89  // ClientStream defines the client-side behavior of a streaming RPC.
    90  //
    91  // All errors returned from ClientStream methods are compatible with the
    92  // status package.
    93  type ClientStream interface {
    94  	// Header returns the header metadata received from the server if there
    95  	// is any. It blocks if the metadata is not ready to read.  If the metadata
    96  	// is nil and the error is also nil, then the stream was terminated without
    97  	// headers, and the status can be discovered by calling RecvMsg.
    98  	Header() (metadata.MD, error)
    99  	// Trailer returns the trailer metadata from the server, if there is any.
   100  	// It must only be called after stream.CloseAndRecv has returned, or
   101  	// stream.Recv has returned a non-nil error (including io.EOF).
   102  	Trailer() metadata.MD
   103  	// CloseSend closes the send direction of the stream. It closes the stream
   104  	// when non-nil error is met. It is also not safe to call CloseSend
   105  	// concurrently with SendMsg.
   106  	CloseSend() error
   107  	// Context returns the context for this stream.
   108  	//
   109  	// It should not be called until after Header or RecvMsg has returned. Once
   110  	// called, subsequent client-side retries are disabled.
   111  	Context() context.Context
   112  	// SendMsg is generally called by generated code. On error, SendMsg aborts
   113  	// the stream. If the error was generated by the client, the status is
   114  	// returned directly; otherwise, io.EOF is returned and the status of
   115  	// the stream may be discovered using RecvMsg.
   116  	//
   117  	// SendMsg blocks until:
   118  	//   - There is sufficient flow control to schedule m with the transport, or
   119  	//   - The stream is done, or
   120  	//   - The stream breaks.
   121  	//
   122  	// SendMsg does not wait until the message is received by the server. An
   123  	// untimely stream closure may result in lost messages. To ensure delivery,
   124  	// users should ensure the RPC completed successfully using RecvMsg.
   125  	//
   126  	// It is safe to have a goroutine calling SendMsg and another goroutine
   127  	// calling RecvMsg on the same stream at the same time, but it is not safe
   128  	// to call SendMsg on the same stream in different goroutines. It is also
   129  	// not safe to call CloseSend concurrently with SendMsg.
   130  	//
   131  	// It is not safe to modify the message after calling SendMsg. Tracing
   132  	// libraries and stats handlers may use the message lazily.
   133  	SendMsg(m any) error
   134  	// RecvMsg blocks until it receives a message into m or the stream is
   135  	// done. It returns io.EOF when the stream completes successfully. On
   136  	// any other error, the stream is aborted and the error contains the RPC
   137  	// status.
   138  	//
   139  	// It is safe to have a goroutine calling SendMsg and another goroutine
   140  	// calling RecvMsg on the same stream at the same time, but it is not
   141  	// safe to call RecvMsg on the same stream in different goroutines.
   142  	RecvMsg(m any) error
   143  }
   144  
   145  // NewStream creates a new Stream for the client side. This is typically
   146  // called by generated code. ctx is used for the lifetime of the stream.
   147  //
   148  // To ensure resources are not leaked due to the stream returned, one of the following
   149  // actions must be performed:
   150  //
   151  //  1. Call Close on the ClientConn.
   152  //  2. Cancel the context provided.
   153  //  3. Call RecvMsg until a non-nil error is returned. A protobuf-generated
   154  //     client-streaming RPC, for instance, might use the helper function
   155  //     CloseAndRecv (note that CloseSend does not Recv, therefore is not
   156  //     guaranteed to release all resources).
   157  //  4. Receive a non-nil, non-io.EOF error from Header or SendMsg.
   158  //
   159  // If none of the above happen, a goroutine and a context will be leaked, and grpc
   160  // will not call the optionally-configured stats handler with a stats.End message.
   161  func (cc *ClientConn) NewStream(ctx context.Context, desc *StreamDesc, method string, opts ...CallOption) (ClientStream, error) {
   162  	// allow interceptor to see all applicable call options, which means those
   163  	// configured as defaults from dial option as well as per-call options
   164  	opts = combine(cc.dopts.callOptions, opts)
   165  
   166  	if cc.dopts.streamInt != nil {
   167  		return cc.dopts.streamInt(ctx, desc, cc, method, newClientStream, opts...)
   168  	}
   169  	return newClientStream(ctx, desc, cc, method, opts...)
   170  }
   171  
   172  // NewClientStream is a wrapper for ClientConn.NewStream.
   173  func NewClientStream(ctx context.Context, desc *StreamDesc, cc *ClientConn, method string, opts ...CallOption) (ClientStream, error) {
   174  	return cc.NewStream(ctx, desc, method, opts...)
   175  }
   176  
   177  func newClientStream(ctx context.Context, desc *StreamDesc, cc *ClientConn, method string, opts ...CallOption) (_ ClientStream, err error) {
   178  	// Start tracking the RPC for idleness purposes. This is where a stream is
   179  	// created for both streaming and unary RPCs, and hence is a good place to
   180  	// track active RPC count.
   181  	if err := cc.idlenessMgr.OnCallBegin(); err != nil {
   182  		return nil, err
   183  	}
   184  	// Add a calloption, to decrement the active call count, that gets executed
   185  	// when the RPC completes.
   186  	opts = append([]CallOption{OnFinish(func(error) { cc.idlenessMgr.OnCallEnd() })}, opts...)
   187  
   188  	if md, added, ok := metadataFromOutgoingContextRaw(ctx); ok {
   189  		// validate md
   190  		if err := imetadata.Validate(md); err != nil {
   191  			return nil, status.Error(codes.Internal, err.Error())
   192  		}
   193  		// validate added
   194  		for _, kvs := range added {
   195  			for i := 0; i < len(kvs); i += 2 {
   196  				if err := imetadata.ValidatePair(kvs[i], kvs[i+1]); err != nil {
   197  					return nil, status.Error(codes.Internal, err.Error())
   198  				}
   199  			}
   200  		}
   201  	}
   202  	if channelz.IsOn() {
   203  		cc.incrCallsStarted()
   204  		defer func() {
   205  			if err != nil {
   206  				cc.incrCallsFailed()
   207  			}
   208  		}()
   209  	}
   210  	// Provide an opportunity for the first RPC to see the first service config
   211  	// provided by the resolver.
   212  	if err := cc.waitForResolvedAddrs(ctx); err != nil {
   213  		return nil, err
   214  	}
   215  
   216  	var mc serviceconfig.MethodConfig
   217  	var onCommit func()
   218  	var newStream = func(ctx context.Context, done func()) (iresolver.ClientStream, error) {
   219  		return newClientStreamWithParams(ctx, desc, cc, method, mc, onCommit, done, opts...)
   220  	}
   221  
   222  	rpcInfo := iresolver.RPCInfo{Context: ctx, Method: method}
   223  	rpcConfig, err := cc.safeConfigSelector.SelectConfig(rpcInfo)
   224  	if err != nil {
   225  		if st, ok := status.FromError(err); ok {
   226  			// Restrict the code to the list allowed by gRFC A54.
   227  			if istatus.IsRestrictedControlPlaneCode(st) {
   228  				err = status.Errorf(codes.Internal, "config selector returned illegal status: %v", err)
   229  			}
   230  			return nil, err
   231  		}
   232  		return nil, toRPCErr(err)
   233  	}
   234  
   235  	if rpcConfig != nil {
   236  		if rpcConfig.Context != nil {
   237  			ctx = rpcConfig.Context
   238  		}
   239  		mc = rpcConfig.MethodConfig
   240  		onCommit = rpcConfig.OnCommitted
   241  		if rpcConfig.Interceptor != nil {
   242  			rpcInfo.Context = nil
   243  			ns := newStream
   244  			newStream = func(ctx context.Context, done func()) (iresolver.ClientStream, error) {
   245  				cs, err := rpcConfig.Interceptor.NewStream(ctx, rpcInfo, done, ns)
   246  				if err != nil {
   247  					return nil, toRPCErr(err)
   248  				}
   249  				return cs, nil
   250  			}
   251  		}
   252  	}
   253  
   254  	return newStream(ctx, func() {})
   255  }
   256  
   257  func newClientStreamWithParams(ctx context.Context, desc *StreamDesc, cc *ClientConn, method string, mc serviceconfig.MethodConfig, onCommit, doneFunc func(), opts ...CallOption) (_ iresolver.ClientStream, err error) {
   258  	c := defaultCallInfo()
   259  	if mc.WaitForReady != nil {
   260  		c.failFast = !*mc.WaitForReady
   261  	}
   262  
   263  	// Possible context leak:
   264  	// The cancel function for the child context we create will only be called
   265  	// when RecvMsg returns a non-nil error, if the ClientConn is closed, or if
   266  	// an error is generated by SendMsg.
   267  	// https://github.com/grpc/grpc-go/issues/1818.
   268  	var cancel context.CancelFunc
   269  	if mc.Timeout != nil && *mc.Timeout >= 0 {
   270  		ctx, cancel = context.WithTimeout(ctx, *mc.Timeout)
   271  	} else {
   272  		ctx, cancel = context.WithCancel(ctx)
   273  	}
   274  	defer func() {
   275  		if err != nil {
   276  			cancel()
   277  		}
   278  	}()
   279  
   280  	for _, o := range opts {
   281  		if err := o.before(c); err != nil {
   282  			return nil, toRPCErr(err)
   283  		}
   284  	}
   285  	c.maxSendMessageSize = getMaxSize(mc.MaxReqSize, c.maxSendMessageSize, defaultClientMaxSendMessageSize)
   286  	c.maxReceiveMessageSize = getMaxSize(mc.MaxRespSize, c.maxReceiveMessageSize, defaultClientMaxReceiveMessageSize)
   287  	if err := setCallInfoCodec(c); err != nil {
   288  		return nil, err
   289  	}
   290  
   291  	callHdr := &transport.CallHdr{
   292  		Host:           cc.authority,
   293  		Method:         method,
   294  		ContentSubtype: c.contentSubtype,
   295  		DoneFunc:       doneFunc,
   296  	}
   297  
   298  	// Set our outgoing compression according to the UseCompressor CallOption, if
   299  	// set.  In that case, also find the compressor from the encoding package.
   300  	// Otherwise, use the compressor configured by the WithCompressor DialOption,
   301  	// if set.
   302  	var cp Compressor
   303  	var comp encoding.Compressor
   304  	if ct := c.compressorType; ct != "" {
   305  		callHdr.SendCompress = ct
   306  		if ct != encoding.Identity {
   307  			comp = encoding.GetCompressor(ct)
   308  			if comp == nil {
   309  				return nil, status.Errorf(codes.Internal, "grpc: Compressor is not installed for requested grpc-encoding %q", ct)
   310  			}
   311  		}
   312  	} else if cc.dopts.cp != nil {
   313  		callHdr.SendCompress = cc.dopts.cp.Type()
   314  		cp = cc.dopts.cp
   315  	}
   316  	if c.creds != nil {
   317  		callHdr.Creds = c.creds
   318  	}
   319  
   320  	cs := &clientStream{
   321  		callHdr:      callHdr,
   322  		ctx:          ctx,
   323  		methodConfig: &mc,
   324  		opts:         opts,
   325  		callInfo:     c,
   326  		cc:           cc,
   327  		desc:         desc,
   328  		codec:        c.codec,
   329  		cp:           cp,
   330  		comp:         comp,
   331  		cancel:       cancel,
   332  		firstAttempt: true,
   333  		onCommit:     onCommit,
   334  	}
   335  	if !cc.dopts.disableRetry {
   336  		cs.retryThrottler = cc.retryThrottler.Load().(*retryThrottler)
   337  	}
   338  	if ml := binarylog.GetMethodLogger(method); ml != nil {
   339  		cs.binlogs = append(cs.binlogs, ml)
   340  	}
   341  	if cc.dopts.binaryLogger != nil {
   342  		if ml := cc.dopts.binaryLogger.GetMethodLogger(method); ml != nil {
   343  			cs.binlogs = append(cs.binlogs, ml)
   344  		}
   345  	}
   346  
   347  	// Pick the transport to use and create a new stream on the transport.
   348  	// Assign cs.attempt upon success.
   349  	op := func(a *csAttempt) error {
   350  		if err := a.getTransport(); err != nil {
   351  			return err
   352  		}
   353  		if err := a.newStream(); err != nil {
   354  			return err
   355  		}
   356  		// Because this operation is always called either here (while creating
   357  		// the clientStream) or by the retry code while locked when replaying
   358  		// the operation, it is safe to access cs.attempt directly.
   359  		cs.attempt = a
   360  		return nil
   361  	}
   362  	if err := cs.withRetry(op, func() { cs.bufferForRetryLocked(0, op) }); err != nil {
   363  		return nil, err
   364  	}
   365  
   366  	if len(cs.binlogs) != 0 {
   367  		md, _ := metadata.FromOutgoingContext(ctx)
   368  		logEntry := &binarylog.ClientHeader{
   369  			OnClientSide: true,
   370  			Header:       md,
   371  			MethodName:   method,
   372  			Authority:    cs.cc.authority,
   373  		}
   374  		if deadline, ok := ctx.Deadline(); ok {
   375  			logEntry.Timeout = time.Until(deadline)
   376  			if logEntry.Timeout < 0 {
   377  				logEntry.Timeout = 0
   378  			}
   379  		}
   380  		for _, binlog := range cs.binlogs {
   381  			binlog.Log(cs.ctx, logEntry)
   382  		}
   383  	}
   384  
   385  	if desc != unaryStreamDesc {
   386  		// Listen on cc and stream contexts to cleanup when the user closes the
   387  		// ClientConn or cancels the stream context.  In all other cases, an error
   388  		// should already be injected into the recv buffer by the transport, which
   389  		// the client will eventually receive, and then we will cancel the stream's
   390  		// context in clientStream.finish.
   391  		go func() {
   392  			select {
   393  			case <-cc.ctx.Done():
   394  				cs.finish(ErrClientConnClosing)
   395  			case <-ctx.Done():
   396  				cs.finish(toRPCErr(ctx.Err()))
   397  			}
   398  		}()
   399  	}
   400  	return cs, nil
   401  }
   402  
   403  // newAttemptLocked creates a new csAttempt without a transport or stream.
   404  func (cs *clientStream) newAttemptLocked(isTransparent bool) (*csAttempt, error) {
   405  	if err := cs.ctx.Err(); err != nil {
   406  		return nil, toRPCErr(err)
   407  	}
   408  	if err := cs.cc.ctx.Err(); err != nil {
   409  		return nil, ErrClientConnClosing
   410  	}
   411  
   412  	ctx := newContextWithRPCInfo(cs.ctx, cs.callInfo.failFast, cs.callInfo.codec, cs.cp, cs.comp)
   413  	method := cs.callHdr.Method
   414  	var beginTime time.Time
   415  	shs := cs.cc.dopts.copts.StatsHandlers
   416  	for _, sh := range shs {
   417  		ctx = sh.TagRPC(ctx, &stats.RPCTagInfo{FullMethodName: method, FailFast: cs.callInfo.failFast})
   418  		beginTime = time.Now()
   419  		begin := &stats.Begin{
   420  			Client:                    true,
   421  			BeginTime:                 beginTime,
   422  			FailFast:                  cs.callInfo.failFast,
   423  			IsClientStream:            cs.desc.ClientStreams,
   424  			IsServerStream:            cs.desc.ServerStreams,
   425  			IsTransparentRetryAttempt: isTransparent,
   426  		}
   427  		sh.HandleRPC(ctx, begin)
   428  	}
   429  
   430  	var trInfo *traceInfo
   431  	if EnableTracing {
   432  		trInfo = &traceInfo{
   433  			tr: newTrace("grpc.Sent."+methodFamily(method), method),
   434  			firstLine: firstLine{
   435  				client: true,
   436  			},
   437  		}
   438  		if deadline, ok := ctx.Deadline(); ok {
   439  			trInfo.firstLine.deadline = time.Until(deadline)
   440  		}
   441  		trInfo.tr.LazyLog(&trInfo.firstLine, false)
   442  		ctx = newTraceContext(ctx, trInfo.tr)
   443  	}
   444  
   445  	if cs.cc.parsedTarget.URL.Scheme == internal.GRPCResolverSchemeExtraMetadata {
   446  		// Add extra metadata (metadata that will be added by transport) to context
   447  		// so the balancer can see them.
   448  		ctx = grpcutil.WithExtraMetadata(ctx, metadata.Pairs(
   449  			"content-type", grpcutil.ContentType(cs.callHdr.ContentSubtype),
   450  		))
   451  	}
   452  
   453  	return &csAttempt{
   454  		ctx:           ctx,
   455  		beginTime:     beginTime,
   456  		cs:            cs,
   457  		dc:            cs.cc.dopts.dc,
   458  		statsHandlers: shs,
   459  		trInfo:        trInfo,
   460  	}, nil
   461  }
   462  
   463  func (a *csAttempt) getTransport() error {
   464  	cs := a.cs
   465  
   466  	var err error
   467  	a.t, a.pickResult, err = cs.cc.getTransport(a.ctx, cs.callInfo.failFast, cs.callHdr.Method)
   468  	if err != nil {
   469  		if de, ok := err.(dropError); ok {
   470  			err = de.error
   471  			a.drop = true
   472  		}
   473  		return err
   474  	}
   475  	if a.trInfo != nil {
   476  		a.trInfo.firstLine.SetRemoteAddr(a.t.RemoteAddr())
   477  	}
   478  	return nil
   479  }
   480  
   481  func (a *csAttempt) newStream() error {
   482  	cs := a.cs
   483  	cs.callHdr.PreviousAttempts = cs.numRetries
   484  
   485  	// Merge metadata stored in PickResult, if any, with existing call metadata.
   486  	// It is safe to overwrite the csAttempt's context here, since all state
   487  	// maintained in it are local to the attempt. When the attempt has to be
   488  	// retried, a new instance of csAttempt will be created.
   489  	if a.pickResult.Metadata != nil {
   490  		// We currently do not have a function it the metadata package which
   491  		// merges given metadata with existing metadata in a context. Existing
   492  		// function `AppendToOutgoingContext()` takes a variadic argument of key
   493  		// value pairs.
   494  		//
   495  		// TODO: Make it possible to retrieve key value pairs from metadata.MD
   496  		// in a form passable to AppendToOutgoingContext(), or create a version
   497  		// of AppendToOutgoingContext() that accepts a metadata.MD.
   498  		md, _ := metadata.FromOutgoingContext(a.ctx)
   499  		md = metadata.Join(md, a.pickResult.Metadata)
   500  		a.ctx = metadata.NewOutgoingContext(a.ctx, md)
   501  	}
   502  
   503  	s, err := a.t.NewStream(a.ctx, cs.callHdr)
   504  	if err != nil {
   505  		nse, ok := err.(*transport.NewStreamError)
   506  		if !ok {
   507  			// Unexpected.
   508  			return err
   509  		}
   510  
   511  		if nse.AllowTransparentRetry {
   512  			a.allowTransparentRetry = true
   513  		}
   514  
   515  		// Unwrap and convert error.
   516  		return toRPCErr(nse.Err)
   517  	}
   518  	a.s = s
   519  	a.ctx = s.Context()
   520  	a.p = &parser{r: s, recvBufferPool: a.cs.cc.dopts.recvBufferPool}
   521  	return nil
   522  }
   523  
   524  // clientStream implements a client side Stream.
   525  type clientStream struct {
   526  	callHdr  *transport.CallHdr
   527  	opts     []CallOption
   528  	callInfo *callInfo
   529  	cc       *ClientConn
   530  	desc     *StreamDesc
   531  
   532  	codec baseCodec
   533  	cp    Compressor
   534  	comp  encoding.Compressor
   535  
   536  	cancel context.CancelFunc // cancels all attempts
   537  
   538  	sentLast bool // sent an end stream
   539  
   540  	methodConfig *MethodConfig
   541  
   542  	ctx context.Context // the application's context, wrapped by stats/tracing
   543  
   544  	retryThrottler *retryThrottler // The throttler active when the RPC began.
   545  
   546  	binlogs []binarylog.MethodLogger
   547  	// serverHeaderBinlogged is a boolean for whether server header has been
   548  	// logged. Server header will be logged when the first time one of those
   549  	// happens: stream.Header(), stream.Recv().
   550  	//
   551  	// It's only read and used by Recv() and Header(), so it doesn't need to be
   552  	// synchronized.
   553  	serverHeaderBinlogged bool
   554  
   555  	mu                      sync.Mutex
   556  	firstAttempt            bool // if true, transparent retry is valid
   557  	numRetries              int  // exclusive of transparent retry attempt(s)
   558  	numRetriesSincePushback int  // retries since pushback; to reset backoff
   559  	finished                bool // TODO: replace with atomic cmpxchg or sync.Once?
   560  	// attempt is the active client stream attempt.
   561  	// The only place where it is written is the newAttemptLocked method and this method never writes nil.
   562  	// So, attempt can be nil only inside newClientStream function when clientStream is first created.
   563  	// One of the first things done after clientStream's creation, is to call newAttemptLocked which either
   564  	// assigns a non nil value to the attempt or returns an error. If an error is returned from newAttemptLocked,
   565  	// then newClientStream calls finish on the clientStream and returns. So, finish method is the only
   566  	// place where we need to check if the attempt is nil.
   567  	attempt *csAttempt
   568  	// TODO(hedging): hedging will have multiple attempts simultaneously.
   569  	committed  bool // active attempt committed for retry?
   570  	onCommit   func()
   571  	buffer     []func(a *csAttempt) error // operations to replay on retry
   572  	bufferSize int                        // current size of buffer
   573  }
   574  
   575  // csAttempt implements a single transport stream attempt within a
   576  // clientStream.
   577  type csAttempt struct {
   578  	ctx        context.Context
   579  	cs         *clientStream
   580  	t          transport.ClientTransport
   581  	s          *transport.Stream
   582  	p          *parser
   583  	pickResult balancer.PickResult
   584  
   585  	finished  bool
   586  	dc        Decompressor
   587  	decomp    encoding.Compressor
   588  	decompSet bool
   589  
   590  	mu sync.Mutex // guards trInfo.tr
   591  	// trInfo may be nil (if EnableTracing is false).
   592  	// trInfo.tr is set when created (if EnableTracing is true),
   593  	// and cleared when the finish method is called.
   594  	trInfo *traceInfo
   595  
   596  	statsHandlers []stats.Handler
   597  	beginTime     time.Time
   598  
   599  	// set for newStream errors that may be transparently retried
   600  	allowTransparentRetry bool
   601  	// set for pick errors that are returned as a status
   602  	drop bool
   603  }
   604  
   605  func (cs *clientStream) commitAttemptLocked() {
   606  	if !cs.committed && cs.onCommit != nil {
   607  		cs.onCommit()
   608  	}
   609  	cs.committed = true
   610  	cs.buffer = nil
   611  }
   612  
   613  func (cs *clientStream) commitAttempt() {
   614  	cs.mu.Lock()
   615  	cs.commitAttemptLocked()
   616  	cs.mu.Unlock()
   617  }
   618  
   619  // shouldRetry returns nil if the RPC should be retried; otherwise it returns
   620  // the error that should be returned by the operation.  If the RPC should be
   621  // retried, the bool indicates whether it is being retried transparently.
   622  func (a *csAttempt) shouldRetry(err error) (bool, error) {
   623  	cs := a.cs
   624  
   625  	if cs.finished || cs.committed || a.drop {
   626  		// RPC is finished or committed or was dropped by the picker; cannot retry.
   627  		return false, err
   628  	}
   629  	if a.s == nil && a.allowTransparentRetry {
   630  		return true, nil
   631  	}
   632  	// Wait for the trailers.
   633  	unprocessed := false
   634  	if a.s != nil {
   635  		<-a.s.Done()
   636  		unprocessed = a.s.Unprocessed()
   637  	}
   638  	if cs.firstAttempt && unprocessed {
   639  		// First attempt, stream unprocessed: transparently retry.
   640  		return true, nil
   641  	}
   642  	if cs.cc.dopts.disableRetry {
   643  		return false, err
   644  	}
   645  
   646  	pushback := 0
   647  	hasPushback := false
   648  	if a.s != nil {
   649  		if !a.s.TrailersOnly() {
   650  			return false, err
   651  		}
   652  
   653  		// TODO(retry): Move down if the spec changes to not check server pushback
   654  		// before considering this a failure for throttling.
   655  		sps := a.s.Trailer()["grpc-retry-pushback-ms"]
   656  		if len(sps) == 1 {
   657  			var e error
   658  			if pushback, e = strconv.Atoi(sps[0]); e != nil || pushback < 0 {
   659  				channelz.Infof(logger, cs.cc.channelz, "Server retry pushback specified to abort (%q).", sps[0])
   660  				cs.retryThrottler.throttle() // This counts as a failure for throttling.
   661  				return false, err
   662  			}
   663  			hasPushback = true
   664  		} else if len(sps) > 1 {
   665  			channelz.Warningf(logger, cs.cc.channelz, "Server retry pushback specified multiple values (%q); not retrying.", sps)
   666  			cs.retryThrottler.throttle() // This counts as a failure for throttling.
   667  			return false, err
   668  		}
   669  	}
   670  
   671  	var code codes.Code
   672  	if a.s != nil {
   673  		code = a.s.Status().Code()
   674  	} else {
   675  		code = status.Code(err)
   676  	}
   677  
   678  	rp := cs.methodConfig.RetryPolicy
   679  	if rp == nil || !rp.RetryableStatusCodes[code] {
   680  		return false, err
   681  	}
   682  
   683  	// Note: the ordering here is important; we count this as a failure
   684  	// only if the code matched a retryable code.
   685  	if cs.retryThrottler.throttle() {
   686  		return false, err
   687  	}
   688  	if cs.numRetries+1 >= rp.MaxAttempts {
   689  		return false, err
   690  	}
   691  
   692  	var dur time.Duration
   693  	if hasPushback {
   694  		dur = time.Millisecond * time.Duration(pushback)
   695  		cs.numRetriesSincePushback = 0
   696  	} else {
   697  		fact := math.Pow(rp.BackoffMultiplier, float64(cs.numRetriesSincePushback))
   698  		cur := float64(rp.InitialBackoff) * fact
   699  		if max := float64(rp.MaxBackoff); cur > max {
   700  			cur = max
   701  		}
   702  		dur = time.Duration(grpcrand.Int63n(int64(cur)))
   703  		cs.numRetriesSincePushback++
   704  	}
   705  
   706  	// TODO(dfawley): we could eagerly fail here if dur puts us past the
   707  	// deadline, but unsure if it is worth doing.
   708  	t := time.NewTimer(dur)
   709  	select {
   710  	case <-t.C:
   711  		cs.numRetries++
   712  		return false, nil
   713  	case <-cs.ctx.Done():
   714  		t.Stop()
   715  		return false, status.FromContextError(cs.ctx.Err()).Err()
   716  	}
   717  }
   718  
   719  // Returns nil if a retry was performed and succeeded; error otherwise.
   720  func (cs *clientStream) retryLocked(attempt *csAttempt, lastErr error) error {
   721  	for {
   722  		attempt.finish(toRPCErr(lastErr))
   723  		isTransparent, err := attempt.shouldRetry(lastErr)
   724  		if err != nil {
   725  			cs.commitAttemptLocked()
   726  			return err
   727  		}
   728  		cs.firstAttempt = false
   729  		attempt, err = cs.newAttemptLocked(isTransparent)
   730  		if err != nil {
   731  			// Only returns error if the clientconn is closed or the context of
   732  			// the stream is canceled.
   733  			return err
   734  		}
   735  		// Note that the first op in the replay buffer always sets cs.attempt
   736  		// if it is able to pick a transport and create a stream.
   737  		if lastErr = cs.replayBufferLocked(attempt); lastErr == nil {
   738  			return nil
   739  		}
   740  	}
   741  }
   742  
   743  func (cs *clientStream) Context() context.Context {
   744  	cs.commitAttempt()
   745  	// No need to lock before using attempt, since we know it is committed and
   746  	// cannot change.
   747  	if cs.attempt.s != nil {
   748  		return cs.attempt.s.Context()
   749  	}
   750  	return cs.ctx
   751  }
   752  
   753  func (cs *clientStream) withRetry(op func(a *csAttempt) error, onSuccess func()) error {
   754  	cs.mu.Lock()
   755  	for {
   756  		if cs.committed {
   757  			cs.mu.Unlock()
   758  			// toRPCErr is used in case the error from the attempt comes from
   759  			// NewClientStream, which intentionally doesn't return a status
   760  			// error to allow for further inspection; all other errors should
   761  			// already be status errors.
   762  			return toRPCErr(op(cs.attempt))
   763  		}
   764  		if len(cs.buffer) == 0 {
   765  			// For the first op, which controls creation of the stream and
   766  			// assigns cs.attempt, we need to create a new attempt inline
   767  			// before executing the first op.  On subsequent ops, the attempt
   768  			// is created immediately before replaying the ops.
   769  			var err error
   770  			if cs.attempt, err = cs.newAttemptLocked(false /* isTransparent */); err != nil {
   771  				cs.mu.Unlock()
   772  				cs.finish(err)
   773  				return err
   774  			}
   775  		}
   776  		a := cs.attempt
   777  		cs.mu.Unlock()
   778  		err := op(a)
   779  		cs.mu.Lock()
   780  		if a != cs.attempt {
   781  			// We started another attempt already.
   782  			continue
   783  		}
   784  		if err == io.EOF {
   785  			<-a.s.Done()
   786  		}
   787  		if err == nil || (err == io.EOF && a.s.Status().Code() == codes.OK) {
   788  			onSuccess()
   789  			cs.mu.Unlock()
   790  			return err
   791  		}
   792  		if err := cs.retryLocked(a, err); err != nil {
   793  			cs.mu.Unlock()
   794  			return err
   795  		}
   796  	}
   797  }
   798  
   799  func (cs *clientStream) Header() (metadata.MD, error) {
   800  	var m metadata.MD
   801  	err := cs.withRetry(func(a *csAttempt) error {
   802  		var err error
   803  		m, err = a.s.Header()
   804  		return toRPCErr(err)
   805  	}, cs.commitAttemptLocked)
   806  
   807  	if m == nil && err == nil {
   808  		// The stream ended with success.  Finish the clientStream.
   809  		err = io.EOF
   810  	}
   811  
   812  	if err != nil {
   813  		cs.finish(err)
   814  		// Do not return the error.  The user should get it by calling Recv().
   815  		return nil, nil
   816  	}
   817  
   818  	if len(cs.binlogs) != 0 && !cs.serverHeaderBinlogged && m != nil {
   819  		// Only log if binary log is on and header has not been logged, and
   820  		// there is actually headers to log.
   821  		logEntry := &binarylog.ServerHeader{
   822  			OnClientSide: true,
   823  			Header:       m,
   824  			PeerAddr:     nil,
   825  		}
   826  		if peer, ok := peer.FromContext(cs.Context()); ok {
   827  			logEntry.PeerAddr = peer.Addr
   828  		}
   829  		cs.serverHeaderBinlogged = true
   830  		for _, binlog := range cs.binlogs {
   831  			binlog.Log(cs.ctx, logEntry)
   832  		}
   833  	}
   834  
   835  	return m, nil
   836  }
   837  
   838  func (cs *clientStream) Trailer() metadata.MD {
   839  	// On RPC failure, we never need to retry, because usage requires that
   840  	// RecvMsg() returned a non-nil error before calling this function is valid.
   841  	// We would have retried earlier if necessary.
   842  	//
   843  	// Commit the attempt anyway, just in case users are not following those
   844  	// directions -- it will prevent races and should not meaningfully impact
   845  	// performance.
   846  	cs.commitAttempt()
   847  	if cs.attempt.s == nil {
   848  		return nil
   849  	}
   850  	return cs.attempt.s.Trailer()
   851  }
   852  
   853  func (cs *clientStream) replayBufferLocked(attempt *csAttempt) error {
   854  	for _, f := range cs.buffer {
   855  		if err := f(attempt); err != nil {
   856  			return err
   857  		}
   858  	}
   859  	return nil
   860  }
   861  
   862  func (cs *clientStream) bufferForRetryLocked(sz int, op func(a *csAttempt) error) {
   863  	// Note: we still will buffer if retry is disabled (for transparent retries).
   864  	if cs.committed {
   865  		return
   866  	}
   867  	cs.bufferSize += sz
   868  	if cs.bufferSize > cs.callInfo.maxRetryRPCBufferSize {
   869  		cs.commitAttemptLocked()
   870  		return
   871  	}
   872  	cs.buffer = append(cs.buffer, op)
   873  }
   874  
   875  func (cs *clientStream) SendMsg(m any) (err error) {
   876  	defer func() {
   877  		if err != nil && err != io.EOF {
   878  			// Call finish on the client stream for errors generated by this SendMsg
   879  			// call, as these indicate problems created by this client.  (Transport
   880  			// errors are converted to an io.EOF error in csAttempt.sendMsg; the real
   881  			// error will be returned from RecvMsg eventually in that case, or be
   882  			// retried.)
   883  			cs.finish(err)
   884  		}
   885  	}()
   886  	if cs.sentLast {
   887  		return status.Errorf(codes.Internal, "SendMsg called after CloseSend")
   888  	}
   889  	if !cs.desc.ClientStreams {
   890  		cs.sentLast = true
   891  	}
   892  
   893  	// load hdr, payload, data
   894  	hdr, payload, data, err := prepareMsg(m, cs.codec, cs.cp, cs.comp)
   895  	if err != nil {
   896  		return err
   897  	}
   898  
   899  	// TODO(dfawley): should we be checking len(data) instead?
   900  	if len(payload) > *cs.callInfo.maxSendMessageSize {
   901  		return status.Errorf(codes.ResourceExhausted, "trying to send message larger than max (%d vs. %d)", len(payload), *cs.callInfo.maxSendMessageSize)
   902  	}
   903  	op := func(a *csAttempt) error {
   904  		return a.sendMsg(m, hdr, payload, data)
   905  	}
   906  	err = cs.withRetry(op, func() { cs.bufferForRetryLocked(len(hdr)+len(payload), op) })
   907  	if len(cs.binlogs) != 0 && err == nil {
   908  		cm := &binarylog.ClientMessage{
   909  			OnClientSide: true,
   910  			Message:      data,
   911  		}
   912  		for _, binlog := range cs.binlogs {
   913  			binlog.Log(cs.ctx, cm)
   914  		}
   915  	}
   916  	return err
   917  }
   918  
   919  func (cs *clientStream) RecvMsg(m any) error {
   920  	if len(cs.binlogs) != 0 && !cs.serverHeaderBinlogged {
   921  		// Call Header() to binary log header if it's not already logged.
   922  		cs.Header()
   923  	}
   924  	var recvInfo *payloadInfo
   925  	if len(cs.binlogs) != 0 {
   926  		recvInfo = &payloadInfo{}
   927  	}
   928  	err := cs.withRetry(func(a *csAttempt) error {
   929  		return a.recvMsg(m, recvInfo)
   930  	}, cs.commitAttemptLocked)
   931  	if len(cs.binlogs) != 0 && err == nil {
   932  		sm := &binarylog.ServerMessage{
   933  			OnClientSide: true,
   934  			Message:      recvInfo.uncompressedBytes,
   935  		}
   936  		for _, binlog := range cs.binlogs {
   937  			binlog.Log(cs.ctx, sm)
   938  		}
   939  	}
   940  	if err != nil || !cs.desc.ServerStreams {
   941  		// err != nil or non-server-streaming indicates end of stream.
   942  		cs.finish(err)
   943  	}
   944  	return err
   945  }
   946  
   947  func (cs *clientStream) CloseSend() error {
   948  	if cs.sentLast {
   949  		// TODO: return an error and finish the stream instead, due to API misuse?
   950  		return nil
   951  	}
   952  	cs.sentLast = true
   953  	op := func(a *csAttempt) error {
   954  		a.t.Write(a.s, nil, nil, &transport.Options{Last: true})
   955  		// Always return nil; io.EOF is the only error that might make sense
   956  		// instead, but there is no need to signal the client to call RecvMsg
   957  		// as the only use left for the stream after CloseSend is to call
   958  		// RecvMsg.  This also matches historical behavior.
   959  		return nil
   960  	}
   961  	cs.withRetry(op, func() { cs.bufferForRetryLocked(0, op) })
   962  	if len(cs.binlogs) != 0 {
   963  		chc := &binarylog.ClientHalfClose{
   964  			OnClientSide: true,
   965  		}
   966  		for _, binlog := range cs.binlogs {
   967  			binlog.Log(cs.ctx, chc)
   968  		}
   969  	}
   970  	// We never returned an error here for reasons.
   971  	return nil
   972  }
   973  
   974  func (cs *clientStream) finish(err error) {
   975  	if err == io.EOF {
   976  		// Ending a stream with EOF indicates a success.
   977  		err = nil
   978  	}
   979  	cs.mu.Lock()
   980  	if cs.finished {
   981  		cs.mu.Unlock()
   982  		return
   983  	}
   984  	cs.finished = true
   985  	for _, onFinish := range cs.callInfo.onFinish {
   986  		onFinish(err)
   987  	}
   988  	cs.commitAttemptLocked()
   989  	if cs.attempt != nil {
   990  		cs.attempt.finish(err)
   991  		// after functions all rely upon having a stream.
   992  		if cs.attempt.s != nil {
   993  			for _, o := range cs.opts {
   994  				o.after(cs.callInfo, cs.attempt)
   995  			}
   996  		}
   997  	}
   998  
   999  	cs.mu.Unlock()
  1000  	// Only one of cancel or trailer needs to be logged.
  1001  	if len(cs.binlogs) != 0 {
  1002  		switch err {
  1003  		case errContextCanceled, errContextDeadline, ErrClientConnClosing:
  1004  			c := &binarylog.Cancel{
  1005  				OnClientSide: true,
  1006  			}
  1007  			for _, binlog := range cs.binlogs {
  1008  				binlog.Log(cs.ctx, c)
  1009  			}
  1010  		default:
  1011  			logEntry := &binarylog.ServerTrailer{
  1012  				OnClientSide: true,
  1013  				Trailer:      cs.Trailer(),
  1014  				Err:          err,
  1015  			}
  1016  			if peer, ok := peer.FromContext(cs.Context()); ok {
  1017  				logEntry.PeerAddr = peer.Addr
  1018  			}
  1019  			for _, binlog := range cs.binlogs {
  1020  				binlog.Log(cs.ctx, logEntry)
  1021  			}
  1022  		}
  1023  	}
  1024  	if err == nil {
  1025  		cs.retryThrottler.successfulRPC()
  1026  	}
  1027  	if channelz.IsOn() {
  1028  		if err != nil {
  1029  			cs.cc.incrCallsFailed()
  1030  		} else {
  1031  			cs.cc.incrCallsSucceeded()
  1032  		}
  1033  	}
  1034  	cs.cancel()
  1035  }
  1036  
  1037  func (a *csAttempt) sendMsg(m any, hdr, payld, data []byte) error {
  1038  	cs := a.cs
  1039  	if a.trInfo != nil {
  1040  		a.mu.Lock()
  1041  		if a.trInfo.tr != nil {
  1042  			a.trInfo.tr.LazyLog(&payload{sent: true, msg: m}, true)
  1043  		}
  1044  		a.mu.Unlock()
  1045  	}
  1046  	if err := a.t.Write(a.s, hdr, payld, &transport.Options{Last: !cs.desc.ClientStreams}); err != nil {
  1047  		if !cs.desc.ClientStreams {
  1048  			// For non-client-streaming RPCs, we return nil instead of EOF on error
  1049  			// because the generated code requires it.  finish is not called; RecvMsg()
  1050  			// will call it with the stream's status independently.
  1051  			return nil
  1052  		}
  1053  		return io.EOF
  1054  	}
  1055  	for _, sh := range a.statsHandlers {
  1056  		sh.HandleRPC(a.ctx, outPayload(true, m, data, payld, time.Now()))
  1057  	}
  1058  	if channelz.IsOn() {
  1059  		a.t.IncrMsgSent()
  1060  	}
  1061  	return nil
  1062  }
  1063  
  1064  func (a *csAttempt) recvMsg(m any, payInfo *payloadInfo) (err error) {
  1065  	cs := a.cs
  1066  	if len(a.statsHandlers) != 0 && payInfo == nil {
  1067  		payInfo = &payloadInfo{}
  1068  	}
  1069  
  1070  	if !a.decompSet {
  1071  		// Block until we receive headers containing received message encoding.
  1072  		if ct := a.s.RecvCompress(); ct != "" && ct != encoding.Identity {
  1073  			if a.dc == nil || a.dc.Type() != ct {
  1074  				// No configured decompressor, or it does not match the incoming
  1075  				// message encoding; attempt to find a registered compressor that does.
  1076  				a.dc = nil
  1077  				a.decomp = encoding.GetCompressor(ct)
  1078  			}
  1079  		} else {
  1080  			// No compression is used; disable our decompressor.
  1081  			a.dc = nil
  1082  		}
  1083  		// Only initialize this state once per stream.
  1084  		a.decompSet = true
  1085  	}
  1086  	err = recv(a.p, cs.codec, a.s, a.dc, m, *cs.callInfo.maxReceiveMessageSize, payInfo, a.decomp)
  1087  	if err != nil {
  1088  		if err == io.EOF {
  1089  			if statusErr := a.s.Status().Err(); statusErr != nil {
  1090  				return statusErr
  1091  			}
  1092  			return io.EOF // indicates successful end of stream.
  1093  		}
  1094  
  1095  		return toRPCErr(err)
  1096  	}
  1097  	if a.trInfo != nil {
  1098  		a.mu.Lock()
  1099  		if a.trInfo.tr != nil {
  1100  			a.trInfo.tr.LazyLog(&payload{sent: false, msg: m}, true)
  1101  		}
  1102  		a.mu.Unlock()
  1103  	}
  1104  	for _, sh := range a.statsHandlers {
  1105  		sh.HandleRPC(a.ctx, &stats.InPayload{
  1106  			Client:   true,
  1107  			RecvTime: time.Now(),
  1108  			Payload:  m,
  1109  			// TODO truncate large payload.
  1110  			Data:             payInfo.uncompressedBytes,
  1111  			WireLength:       payInfo.compressedLength + headerLen,
  1112  			CompressedLength: payInfo.compressedLength,
  1113  			Length:           len(payInfo.uncompressedBytes),
  1114  		})
  1115  	}
  1116  	if channelz.IsOn() {
  1117  		a.t.IncrMsgRecv()
  1118  	}
  1119  	if cs.desc.ServerStreams {
  1120  		// Subsequent messages should be received by subsequent RecvMsg calls.
  1121  		return nil
  1122  	}
  1123  	// Special handling for non-server-stream rpcs.
  1124  	// This recv expects EOF or errors, so we don't collect inPayload.
  1125  	err = recv(a.p, cs.codec, a.s, a.dc, m, *cs.callInfo.maxReceiveMessageSize, nil, a.decomp)
  1126  	if err == nil {
  1127  		return toRPCErr(errors.New("grpc: client streaming protocol violation: get <nil>, want <EOF>"))
  1128  	}
  1129  	if err == io.EOF {
  1130  		return a.s.Status().Err() // non-server streaming Recv returns nil on success
  1131  	}
  1132  	return toRPCErr(err)
  1133  }
  1134  
  1135  func (a *csAttempt) finish(err error) {
  1136  	a.mu.Lock()
  1137  	if a.finished {
  1138  		a.mu.Unlock()
  1139  		return
  1140  	}
  1141  	a.finished = true
  1142  	if err == io.EOF {
  1143  		// Ending a stream with EOF indicates a success.
  1144  		err = nil
  1145  	}
  1146  	var tr metadata.MD
  1147  	if a.s != nil {
  1148  		a.t.CloseStream(a.s, err)
  1149  		tr = a.s.Trailer()
  1150  	}
  1151  
  1152  	if a.pickResult.Done != nil {
  1153  		br := false
  1154  		if a.s != nil {
  1155  			br = a.s.BytesReceived()
  1156  		}
  1157  		a.pickResult.Done(balancer.DoneInfo{
  1158  			Err:           err,
  1159  			Trailer:       tr,
  1160  			BytesSent:     a.s != nil,
  1161  			BytesReceived: br,
  1162  			ServerLoad:    balancerload.Parse(tr),
  1163  		})
  1164  	}
  1165  	for _, sh := range a.statsHandlers {
  1166  		end := &stats.End{
  1167  			Client:    true,
  1168  			BeginTime: a.beginTime,
  1169  			EndTime:   time.Now(),
  1170  			Trailer:   tr,
  1171  			Error:     err,
  1172  		}
  1173  		sh.HandleRPC(a.ctx, end)
  1174  	}
  1175  	if a.trInfo != nil && a.trInfo.tr != nil {
  1176  		if err == nil {
  1177  			a.trInfo.tr.LazyPrintf("RPC: [OK]")
  1178  		} else {
  1179  			a.trInfo.tr.LazyPrintf("RPC: [%v]", err)
  1180  			a.trInfo.tr.SetError()
  1181  		}
  1182  		a.trInfo.tr.Finish()
  1183  		a.trInfo.tr = nil
  1184  	}
  1185  	a.mu.Unlock()
  1186  }
  1187  
  1188  // newClientStream creates a ClientStream with the specified transport, on the
  1189  // given addrConn.
  1190  //
  1191  // It's expected that the given transport is either the same one in addrConn, or
  1192  // is already closed. To avoid race, transport is specified separately, instead
  1193  // of using ac.transpot.
  1194  //
  1195  // Main difference between this and ClientConn.NewStream:
  1196  // - no retry
  1197  // - no service config (or wait for service config)
  1198  // - no tracing or stats
  1199  func newNonRetryClientStream(ctx context.Context, desc *StreamDesc, method string, t transport.ClientTransport, ac *addrConn, opts ...CallOption) (_ ClientStream, err error) {
  1200  	if t == nil {
  1201  		// TODO: return RPC error here?
  1202  		return nil, errors.New("transport provided is nil")
  1203  	}
  1204  	// defaultCallInfo contains unnecessary info(i.e. failfast, maxRetryRPCBufferSize), so we just initialize an empty struct.
  1205  	c := &callInfo{}
  1206  
  1207  	// Possible context leak:
  1208  	// The cancel function for the child context we create will only be called
  1209  	// when RecvMsg returns a non-nil error, if the ClientConn is closed, or if
  1210  	// an error is generated by SendMsg.
  1211  	// https://github.com/grpc/grpc-go/issues/1818.
  1212  	ctx, cancel := context.WithCancel(ctx)
  1213  	defer func() {
  1214  		if err != nil {
  1215  			cancel()
  1216  		}
  1217  	}()
  1218  
  1219  	for _, o := range opts {
  1220  		if err := o.before(c); err != nil {
  1221  			return nil, toRPCErr(err)
  1222  		}
  1223  	}
  1224  	c.maxReceiveMessageSize = getMaxSize(nil, c.maxReceiveMessageSize, defaultClientMaxReceiveMessageSize)
  1225  	c.maxSendMessageSize = getMaxSize(nil, c.maxSendMessageSize, defaultServerMaxSendMessageSize)
  1226  	if err := setCallInfoCodec(c); err != nil {
  1227  		return nil, err
  1228  	}
  1229  
  1230  	callHdr := &transport.CallHdr{
  1231  		Host:           ac.cc.authority,
  1232  		Method:         method,
  1233  		ContentSubtype: c.contentSubtype,
  1234  	}
  1235  
  1236  	// Set our outgoing compression according to the UseCompressor CallOption, if
  1237  	// set.  In that case, also find the compressor from the encoding package.
  1238  	// Otherwise, use the compressor configured by the WithCompressor DialOption,
  1239  	// if set.
  1240  	var cp Compressor
  1241  	var comp encoding.Compressor
  1242  	if ct := c.compressorType; ct != "" {
  1243  		callHdr.SendCompress = ct
  1244  		if ct != encoding.Identity {
  1245  			comp = encoding.GetCompressor(ct)
  1246  			if comp == nil {
  1247  				return nil, status.Errorf(codes.Internal, "grpc: Compressor is not installed for requested grpc-encoding %q", ct)
  1248  			}
  1249  		}
  1250  	} else if ac.cc.dopts.cp != nil {
  1251  		callHdr.SendCompress = ac.cc.dopts.cp.Type()
  1252  		cp = ac.cc.dopts.cp
  1253  	}
  1254  	if c.creds != nil {
  1255  		callHdr.Creds = c.creds
  1256  	}
  1257  
  1258  	// Use a special addrConnStream to avoid retry.
  1259  	as := &addrConnStream{
  1260  		callHdr:  callHdr,
  1261  		ac:       ac,
  1262  		ctx:      ctx,
  1263  		cancel:   cancel,
  1264  		opts:     opts,
  1265  		callInfo: c,
  1266  		desc:     desc,
  1267  		codec:    c.codec,
  1268  		cp:       cp,
  1269  		comp:     comp,
  1270  		t:        t,
  1271  	}
  1272  
  1273  	s, err := as.t.NewStream(as.ctx, as.callHdr)
  1274  	if err != nil {
  1275  		err = toRPCErr(err)
  1276  		return nil, err
  1277  	}
  1278  	as.s = s
  1279  	as.p = &parser{r: s, recvBufferPool: ac.dopts.recvBufferPool}
  1280  	ac.incrCallsStarted()
  1281  	if desc != unaryStreamDesc {
  1282  		// Listen on stream context to cleanup when the stream context is
  1283  		// canceled.  Also listen for the addrConn's context in case the
  1284  		// addrConn is closed or reconnects to a different address.  In all
  1285  		// other cases, an error should already be injected into the recv
  1286  		// buffer by the transport, which the client will eventually receive,
  1287  		// and then we will cancel the stream's context in
  1288  		// addrConnStream.finish.
  1289  		go func() {
  1290  			ac.mu.Lock()
  1291  			acCtx := ac.ctx
  1292  			ac.mu.Unlock()
  1293  			select {
  1294  			case <-acCtx.Done():
  1295  				as.finish(status.Error(codes.Canceled, "grpc: the SubConn is closing"))
  1296  			case <-ctx.Done():
  1297  				as.finish(toRPCErr(ctx.Err()))
  1298  			}
  1299  		}()
  1300  	}
  1301  	return as, nil
  1302  }
  1303  
  1304  type addrConnStream struct {
  1305  	s         *transport.Stream
  1306  	ac        *addrConn
  1307  	callHdr   *transport.CallHdr
  1308  	cancel    context.CancelFunc
  1309  	opts      []CallOption
  1310  	callInfo  *callInfo
  1311  	t         transport.ClientTransport
  1312  	ctx       context.Context
  1313  	sentLast  bool
  1314  	desc      *StreamDesc
  1315  	codec     baseCodec
  1316  	cp        Compressor
  1317  	comp      encoding.Compressor
  1318  	decompSet bool
  1319  	dc        Decompressor
  1320  	decomp    encoding.Compressor
  1321  	p         *parser
  1322  	mu        sync.Mutex
  1323  	finished  bool
  1324  }
  1325  
  1326  func (as *addrConnStream) Header() (metadata.MD, error) {
  1327  	m, err := as.s.Header()
  1328  	if err != nil {
  1329  		as.finish(toRPCErr(err))
  1330  	}
  1331  	return m, err
  1332  }
  1333  
  1334  func (as *addrConnStream) Trailer() metadata.MD {
  1335  	return as.s.Trailer()
  1336  }
  1337  
  1338  func (as *addrConnStream) CloseSend() error {
  1339  	if as.sentLast {
  1340  		// TODO: return an error and finish the stream instead, due to API misuse?
  1341  		return nil
  1342  	}
  1343  	as.sentLast = true
  1344  
  1345  	as.t.Write(as.s, nil, nil, &transport.Options{Last: true})
  1346  	// Always return nil; io.EOF is the only error that might make sense
  1347  	// instead, but there is no need to signal the client to call RecvMsg
  1348  	// as the only use left for the stream after CloseSend is to call
  1349  	// RecvMsg.  This also matches historical behavior.
  1350  	return nil
  1351  }
  1352  
  1353  func (as *addrConnStream) Context() context.Context {
  1354  	return as.s.Context()
  1355  }
  1356  
  1357  func (as *addrConnStream) SendMsg(m any) (err error) {
  1358  	defer func() {
  1359  		if err != nil && err != io.EOF {
  1360  			// Call finish on the client stream for errors generated by this SendMsg
  1361  			// call, as these indicate problems created by this client.  (Transport
  1362  			// errors are converted to an io.EOF error in csAttempt.sendMsg; the real
  1363  			// error will be returned from RecvMsg eventually in that case, or be
  1364  			// retried.)
  1365  			as.finish(err)
  1366  		}
  1367  	}()
  1368  	if as.sentLast {
  1369  		return status.Errorf(codes.Internal, "SendMsg called after CloseSend")
  1370  	}
  1371  	if !as.desc.ClientStreams {
  1372  		as.sentLast = true
  1373  	}
  1374  
  1375  	// load hdr, payload, data
  1376  	hdr, payld, _, err := prepareMsg(m, as.codec, as.cp, as.comp)
  1377  	if err != nil {
  1378  		return err
  1379  	}
  1380  
  1381  	// TODO(dfawley): should we be checking len(data) instead?
  1382  	if len(payld) > *as.callInfo.maxSendMessageSize {
  1383  		return status.Errorf(codes.ResourceExhausted, "trying to send message larger than max (%d vs. %d)", len(payld), *as.callInfo.maxSendMessageSize)
  1384  	}
  1385  
  1386  	if err := as.t.Write(as.s, hdr, payld, &transport.Options{Last: !as.desc.ClientStreams}); err != nil {
  1387  		if !as.desc.ClientStreams {
  1388  			// For non-client-streaming RPCs, we return nil instead of EOF on error
  1389  			// because the generated code requires it.  finish is not called; RecvMsg()
  1390  			// will call it with the stream's status independently.
  1391  			return nil
  1392  		}
  1393  		return io.EOF
  1394  	}
  1395  
  1396  	if channelz.IsOn() {
  1397  		as.t.IncrMsgSent()
  1398  	}
  1399  	return nil
  1400  }
  1401  
  1402  func (as *addrConnStream) RecvMsg(m any) (err error) {
  1403  	defer func() {
  1404  		if err != nil || !as.desc.ServerStreams {
  1405  			// err != nil or non-server-streaming indicates end of stream.
  1406  			as.finish(err)
  1407  		}
  1408  	}()
  1409  
  1410  	if !as.decompSet {
  1411  		// Block until we receive headers containing received message encoding.
  1412  		if ct := as.s.RecvCompress(); ct != "" && ct != encoding.Identity {
  1413  			if as.dc == nil || as.dc.Type() != ct {
  1414  				// No configured decompressor, or it does not match the incoming
  1415  				// message encoding; attempt to find a registered compressor that does.
  1416  				as.dc = nil
  1417  				as.decomp = encoding.GetCompressor(ct)
  1418  			}
  1419  		} else {
  1420  			// No compression is used; disable our decompressor.
  1421  			as.dc = nil
  1422  		}
  1423  		// Only initialize this state once per stream.
  1424  		as.decompSet = true
  1425  	}
  1426  	err = recv(as.p, as.codec, as.s, as.dc, m, *as.callInfo.maxReceiveMessageSize, nil, as.decomp)
  1427  	if err != nil {
  1428  		if err == io.EOF {
  1429  			if statusErr := as.s.Status().Err(); statusErr != nil {
  1430  				return statusErr
  1431  			}
  1432  			return io.EOF // indicates successful end of stream.
  1433  		}
  1434  		return toRPCErr(err)
  1435  	}
  1436  
  1437  	if channelz.IsOn() {
  1438  		as.t.IncrMsgRecv()
  1439  	}
  1440  	if as.desc.ServerStreams {
  1441  		// Subsequent messages should be received by subsequent RecvMsg calls.
  1442  		return nil
  1443  	}
  1444  
  1445  	// Special handling for non-server-stream rpcs.
  1446  	// This recv expects EOF or errors, so we don't collect inPayload.
  1447  	err = recv(as.p, as.codec, as.s, as.dc, m, *as.callInfo.maxReceiveMessageSize, nil, as.decomp)
  1448  	if err == nil {
  1449  		return toRPCErr(errors.New("grpc: client streaming protocol violation: get <nil>, want <EOF>"))
  1450  	}
  1451  	if err == io.EOF {
  1452  		return as.s.Status().Err() // non-server streaming Recv returns nil on success
  1453  	}
  1454  	return toRPCErr(err)
  1455  }
  1456  
  1457  func (as *addrConnStream) finish(err error) {
  1458  	as.mu.Lock()
  1459  	if as.finished {
  1460  		as.mu.Unlock()
  1461  		return
  1462  	}
  1463  	as.finished = true
  1464  	if err == io.EOF {
  1465  		// Ending a stream with EOF indicates a success.
  1466  		err = nil
  1467  	}
  1468  	if as.s != nil {
  1469  		as.t.CloseStream(as.s, err)
  1470  	}
  1471  
  1472  	if err != nil {
  1473  		as.ac.incrCallsFailed()
  1474  	} else {
  1475  		as.ac.incrCallsSucceeded()
  1476  	}
  1477  	as.cancel()
  1478  	as.mu.Unlock()
  1479  }
  1480  
  1481  // ServerStream defines the server-side behavior of a streaming RPC.
  1482  //
  1483  // Errors returned from ServerStream methods are compatible with the status
  1484  // package.  However, the status code will often not match the RPC status as
  1485  // seen by the client application, and therefore, should not be relied upon for
  1486  // this purpose.
  1487  type ServerStream interface {
  1488  	// SetHeader sets the header metadata. It may be called multiple times.
  1489  	// When call multiple times, all the provided metadata will be merged.
  1490  	// All the metadata will be sent out when one of the following happens:
  1491  	//  - ServerStream.SendHeader() is called;
  1492  	//  - The first response is sent out;
  1493  	//  - An RPC status is sent out (error or success).
  1494  	SetHeader(metadata.MD) error
  1495  	// SendHeader sends the header metadata.
  1496  	// The provided md and headers set by SetHeader() will be sent.
  1497  	// It fails if called multiple times.
  1498  	SendHeader(metadata.MD) error
  1499  	// SetTrailer sets the trailer metadata which will be sent with the RPC status.
  1500  	// When called more than once, all the provided metadata will be merged.
  1501  	SetTrailer(metadata.MD)
  1502  	// Context returns the context for this stream.
  1503  	Context() context.Context
  1504  	// SendMsg sends a message. On error, SendMsg aborts the stream and the
  1505  	// error is returned directly.
  1506  	//
  1507  	// SendMsg blocks until:
  1508  	//   - There is sufficient flow control to schedule m with the transport, or
  1509  	//   - The stream is done, or
  1510  	//   - The stream breaks.
  1511  	//
  1512  	// SendMsg does not wait until the message is received by the client. An
  1513  	// untimely stream closure may result in lost messages.
  1514  	//
  1515  	// It is safe to have a goroutine calling SendMsg and another goroutine
  1516  	// calling RecvMsg on the same stream at the same time, but it is not safe
  1517  	// to call SendMsg on the same stream in different goroutines.
  1518  	//
  1519  	// It is not safe to modify the message after calling SendMsg. Tracing
  1520  	// libraries and stats handlers may use the message lazily.
  1521  	SendMsg(m any) error
  1522  	// RecvMsg blocks until it receives a message into m or the stream is
  1523  	// done. It returns io.EOF when the client has performed a CloseSend. On
  1524  	// any non-EOF error, the stream is aborted and the error contains the
  1525  	// RPC status.
  1526  	//
  1527  	// It is safe to have a goroutine calling SendMsg and another goroutine
  1528  	// calling RecvMsg on the same stream at the same time, but it is not
  1529  	// safe to call RecvMsg on the same stream in different goroutines.
  1530  	RecvMsg(m any) error
  1531  }
  1532  
  1533  // serverStream implements a server side Stream.
  1534  type serverStream struct {
  1535  	ctx   context.Context
  1536  	t     transport.ServerTransport
  1537  	s     *transport.Stream
  1538  	p     *parser
  1539  	codec baseCodec
  1540  
  1541  	cp     Compressor
  1542  	dc     Decompressor
  1543  	comp   encoding.Compressor
  1544  	decomp encoding.Compressor
  1545  
  1546  	sendCompressorName string
  1547  
  1548  	maxReceiveMessageSize int
  1549  	maxSendMessageSize    int
  1550  	trInfo                *traceInfo
  1551  
  1552  	statsHandler []stats.Handler
  1553  
  1554  	binlogs []binarylog.MethodLogger
  1555  	// serverHeaderBinlogged indicates whether server header has been logged. It
  1556  	// will happen when one of the following two happens: stream.SendHeader(),
  1557  	// stream.Send().
  1558  	//
  1559  	// It's only checked in send and sendHeader, doesn't need to be
  1560  	// synchronized.
  1561  	serverHeaderBinlogged bool
  1562  
  1563  	mu sync.Mutex // protects trInfo.tr after the service handler runs.
  1564  }
  1565  
  1566  func (ss *serverStream) Context() context.Context {
  1567  	return ss.ctx
  1568  }
  1569  
  1570  func (ss *serverStream) SetHeader(md metadata.MD) error {
  1571  	if md.Len() == 0 {
  1572  		return nil
  1573  	}
  1574  	err := imetadata.Validate(md)
  1575  	if err != nil {
  1576  		return status.Error(codes.Internal, err.Error())
  1577  	}
  1578  	return ss.s.SetHeader(md)
  1579  }
  1580  
  1581  func (ss *serverStream) SendHeader(md metadata.MD) error {
  1582  	err := imetadata.Validate(md)
  1583  	if err != nil {
  1584  		return status.Error(codes.Internal, err.Error())
  1585  	}
  1586  
  1587  	err = ss.t.WriteHeader(ss.s, md)
  1588  	if len(ss.binlogs) != 0 && !ss.serverHeaderBinlogged {
  1589  		h, _ := ss.s.Header()
  1590  		sh := &binarylog.ServerHeader{
  1591  			Header: h,
  1592  		}
  1593  		ss.serverHeaderBinlogged = true
  1594  		for _, binlog := range ss.binlogs {
  1595  			binlog.Log(ss.ctx, sh)
  1596  		}
  1597  	}
  1598  	return err
  1599  }
  1600  
  1601  func (ss *serverStream) SetTrailer(md metadata.MD) {
  1602  	if md.Len() == 0 {
  1603  		return
  1604  	}
  1605  	if err := imetadata.Validate(md); err != nil {
  1606  		logger.Errorf("stream: failed to validate md when setting trailer, err: %v", err)
  1607  	}
  1608  	ss.s.SetTrailer(md)
  1609  }
  1610  
  1611  func (ss *serverStream) SendMsg(m any) (err error) {
  1612  	defer func() {
  1613  		if ss.trInfo != nil {
  1614  			ss.mu.Lock()
  1615  			if ss.trInfo.tr != nil {
  1616  				if err == nil {
  1617  					ss.trInfo.tr.LazyLog(&payload{sent: true, msg: m}, true)
  1618  				} else {
  1619  					ss.trInfo.tr.LazyLog(&fmtStringer{"%v", []any{err}}, true)
  1620  					ss.trInfo.tr.SetError()
  1621  				}
  1622  			}
  1623  			ss.mu.Unlock()
  1624  		}
  1625  		if err != nil && err != io.EOF {
  1626  			st, _ := status.FromError(toRPCErr(err))
  1627  			ss.t.WriteStatus(ss.s, st)
  1628  			// Non-user specified status was sent out. This should be an error
  1629  			// case (as a server side Cancel maybe).
  1630  			//
  1631  			// This is not handled specifically now. User will return a final
  1632  			// status from the service handler, we will log that error instead.
  1633  			// This behavior is similar to an interceptor.
  1634  		}
  1635  		if channelz.IsOn() && err == nil {
  1636  			ss.t.IncrMsgSent()
  1637  		}
  1638  	}()
  1639  
  1640  	// Server handler could have set new compressor by calling SetSendCompressor.
  1641  	// In case it is set, we need to use it for compressing outbound message.
  1642  	if sendCompressorsName := ss.s.SendCompress(); sendCompressorsName != ss.sendCompressorName {
  1643  		ss.comp = encoding.GetCompressor(sendCompressorsName)
  1644  		ss.sendCompressorName = sendCompressorsName
  1645  	}
  1646  
  1647  	// load hdr, payload, data
  1648  	hdr, payload, data, err := prepareMsg(m, ss.codec, ss.cp, ss.comp)
  1649  	if err != nil {
  1650  		return err
  1651  	}
  1652  
  1653  	// TODO(dfawley): should we be checking len(data) instead?
  1654  	if len(payload) > ss.maxSendMessageSize {
  1655  		return status.Errorf(codes.ResourceExhausted, "trying to send message larger than max (%d vs. %d)", len(payload), ss.maxSendMessageSize)
  1656  	}
  1657  	if err := ss.t.Write(ss.s, hdr, payload, &transport.Options{Last: false}); err != nil {
  1658  		return toRPCErr(err)
  1659  	}
  1660  	if len(ss.binlogs) != 0 {
  1661  		if !ss.serverHeaderBinlogged {
  1662  			h, _ := ss.s.Header()
  1663  			sh := &binarylog.ServerHeader{
  1664  				Header: h,
  1665  			}
  1666  			ss.serverHeaderBinlogged = true
  1667  			for _, binlog := range ss.binlogs {
  1668  				binlog.Log(ss.ctx, sh)
  1669  			}
  1670  		}
  1671  		sm := &binarylog.ServerMessage{
  1672  			Message: data,
  1673  		}
  1674  		for _, binlog := range ss.binlogs {
  1675  			binlog.Log(ss.ctx, sm)
  1676  		}
  1677  	}
  1678  	if len(ss.statsHandler) != 0 {
  1679  		for _, sh := range ss.statsHandler {
  1680  			sh.HandleRPC(ss.s.Context(), outPayload(false, m, data, payload, time.Now()))
  1681  		}
  1682  	}
  1683  	return nil
  1684  }
  1685  
  1686  func (ss *serverStream) RecvMsg(m any) (err error) {
  1687  	defer func() {
  1688  		if ss.trInfo != nil {
  1689  			ss.mu.Lock()
  1690  			if ss.trInfo.tr != nil {
  1691  				if err == nil {
  1692  					ss.trInfo.tr.LazyLog(&payload{sent: false, msg: m}, true)
  1693  				} else if err != io.EOF {
  1694  					ss.trInfo.tr.LazyLog(&fmtStringer{"%v", []any{err}}, true)
  1695  					ss.trInfo.tr.SetError()
  1696  				}
  1697  			}
  1698  			ss.mu.Unlock()
  1699  		}
  1700  		if err != nil && err != io.EOF {
  1701  			st, _ := status.FromError(toRPCErr(err))
  1702  			ss.t.WriteStatus(ss.s, st)
  1703  			// Non-user specified status was sent out. This should be an error
  1704  			// case (as a server side Cancel maybe).
  1705  			//
  1706  			// This is not handled specifically now. User will return a final
  1707  			// status from the service handler, we will log that error instead.
  1708  			// This behavior is similar to an interceptor.
  1709  		}
  1710  		if channelz.IsOn() && err == nil {
  1711  			ss.t.IncrMsgRecv()
  1712  		}
  1713  	}()
  1714  	var payInfo *payloadInfo
  1715  	if len(ss.statsHandler) != 0 || len(ss.binlogs) != 0 {
  1716  		payInfo = &payloadInfo{}
  1717  	}
  1718  	if err := recv(ss.p, ss.codec, ss.s, ss.dc, m, ss.maxReceiveMessageSize, payInfo, ss.decomp); err != nil {
  1719  		if err == io.EOF {
  1720  			if len(ss.binlogs) != 0 {
  1721  				chc := &binarylog.ClientHalfClose{}
  1722  				for _, binlog := range ss.binlogs {
  1723  					binlog.Log(ss.ctx, chc)
  1724  				}
  1725  			}
  1726  			return err
  1727  		}
  1728  		if err == io.ErrUnexpectedEOF {
  1729  			err = status.Errorf(codes.Internal, io.ErrUnexpectedEOF.Error())
  1730  		}
  1731  		return toRPCErr(err)
  1732  	}
  1733  	if len(ss.statsHandler) != 0 {
  1734  		for _, sh := range ss.statsHandler {
  1735  			sh.HandleRPC(ss.s.Context(), &stats.InPayload{
  1736  				RecvTime: time.Now(),
  1737  				Payload:  m,
  1738  				// TODO truncate large payload.
  1739  				Data:             payInfo.uncompressedBytes,
  1740  				Length:           len(payInfo.uncompressedBytes),
  1741  				WireLength:       payInfo.compressedLength + headerLen,
  1742  				CompressedLength: payInfo.compressedLength,
  1743  			})
  1744  		}
  1745  	}
  1746  	if len(ss.binlogs) != 0 {
  1747  		cm := &binarylog.ClientMessage{
  1748  			Message: payInfo.uncompressedBytes,
  1749  		}
  1750  		for _, binlog := range ss.binlogs {
  1751  			binlog.Log(ss.ctx, cm)
  1752  		}
  1753  	}
  1754  	return nil
  1755  }
  1756  
  1757  // MethodFromServerStream returns the method string for the input stream.
  1758  // The returned string is in the format of "/service/method".
  1759  func MethodFromServerStream(stream ServerStream) (string, bool) {
  1760  	return Method(stream.Context())
  1761  }
  1762  
  1763  // prepareMsg returns the hdr, payload and data
  1764  // using the compressors passed or using the
  1765  // passed preparedmsg
  1766  func prepareMsg(m any, codec baseCodec, cp Compressor, comp encoding.Compressor) (hdr, payload, data []byte, err error) {
  1767  	if preparedMsg, ok := m.(*PreparedMsg); ok {
  1768  		return preparedMsg.hdr, preparedMsg.payload, preparedMsg.encodedData, nil
  1769  	}
  1770  	// The input interface is not a prepared msg.
  1771  	// Marshal and Compress the data at this point
  1772  	data, err = encode(codec, m)
  1773  	if err != nil {
  1774  		return nil, nil, nil, err
  1775  	}
  1776  	compData, err := compress(data, cp, comp)
  1777  	if err != nil {
  1778  		return nil, nil, nil, err
  1779  	}
  1780  	hdr, payload = msgHeader(data, compData)
  1781  	return hdr, payload, data, nil
  1782  }
  1783  

View as plain text