stream.go 43 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482148314841485
  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. package grpc
  19. import (
  20. "context"
  21. "errors"
  22. "io"
  23. "math"
  24. "strconv"
  25. "sync"
  26. "time"
  27. "golang.org/x/net/trace"
  28. "google.golang.org/grpc/balancer"
  29. "google.golang.org/grpc/codes"
  30. "google.golang.org/grpc/connectivity"
  31. "google.golang.org/grpc/encoding"
  32. "google.golang.org/grpc/grpclog"
  33. "google.golang.org/grpc/internal/binarylog"
  34. "google.golang.org/grpc/internal/channelz"
  35. "google.golang.org/grpc/internal/grpcrand"
  36. "google.golang.org/grpc/internal/transport"
  37. "google.golang.org/grpc/metadata"
  38. "google.golang.org/grpc/peer"
  39. "google.golang.org/grpc/stats"
  40. "google.golang.org/grpc/status"
  41. )
  42. // StreamHandler defines the handler called by gRPC server to complete the
  43. // execution of a streaming RPC. If a StreamHandler returns an error, it
  44. // should be produced by the status package, or else gRPC will use
  45. // codes.Unknown as the status code and err.Error() as the status message
  46. // of the RPC.
  47. type StreamHandler func(srv interface{}, stream ServerStream) error
  48. // StreamDesc represents a streaming RPC service's method specification.
  49. type StreamDesc struct {
  50. StreamName string
  51. Handler StreamHandler
  52. // At least one of these is true.
  53. ServerStreams bool
  54. ClientStreams bool
  55. }
  56. // Stream defines the common interface a client or server stream has to satisfy.
  57. //
  58. // Deprecated: See ClientStream and ServerStream documentation instead.
  59. type Stream interface {
  60. // Deprecated: See ClientStream and ServerStream documentation instead.
  61. Context() context.Context
  62. // Deprecated: See ClientStream and ServerStream documentation instead.
  63. SendMsg(m interface{}) error
  64. // Deprecated: See ClientStream and ServerStream documentation instead.
  65. RecvMsg(m interface{}) error
  66. }
  67. // ClientStream defines the client-side behavior of a streaming RPC.
  68. //
  69. // All errors returned from ClientStream methods are compatible with the
  70. // status package.
  71. type ClientStream interface {
  72. // Header returns the header metadata received from the server if there
  73. // is any. It blocks if the metadata is not ready to read.
  74. Header() (metadata.MD, error)
  75. // Trailer returns the trailer metadata from the server, if there is any.
  76. // It must only be called after stream.CloseAndRecv has returned, or
  77. // stream.Recv has returned a non-nil error (including io.EOF).
  78. Trailer() metadata.MD
  79. // CloseSend closes the send direction of the stream. It closes the stream
  80. // when non-nil error is met. It is also not safe to call CloseSend
  81. // concurrently with SendMsg.
  82. CloseSend() error
  83. // Context returns the context for this stream.
  84. //
  85. // It should not be called until after Header or RecvMsg has returned. Once
  86. // called, subsequent client-side retries are disabled.
  87. Context() context.Context
  88. // SendMsg is generally called by generated code. On error, SendMsg aborts
  89. // the stream. If the error was generated by the client, the status is
  90. // returned directly; otherwise, io.EOF is returned and the status of
  91. // the stream may be discovered using RecvMsg.
  92. //
  93. // SendMsg blocks until:
  94. // - There is sufficient flow control to schedule m with the transport, or
  95. // - The stream is done, or
  96. // - The stream breaks.
  97. //
  98. // SendMsg does not wait until the message is received by the server. An
  99. // untimely stream closure may result in lost messages. To ensure delivery,
  100. // users should ensure the RPC completed successfully using RecvMsg.
  101. //
  102. // It is safe to have a goroutine calling SendMsg and another goroutine
  103. // calling RecvMsg on the same stream at the same time, but it is not safe
  104. // to call SendMsg on the same stream in different goroutines. It is also
  105. // not safe to call CloseSend concurrently with SendMsg.
  106. SendMsg(m interface{}) error
  107. // RecvMsg blocks until it receives a message into m or the stream is
  108. // done. It returns io.EOF when the stream completes successfully. On
  109. // any other error, the stream is aborted and the error contains the RPC
  110. // status.
  111. //
  112. // It is safe to have a goroutine calling SendMsg and another goroutine
  113. // calling RecvMsg on the same stream at the same time, but it is not
  114. // safe to call RecvMsg on the same stream in different goroutines.
  115. RecvMsg(m interface{}) error
  116. }
  117. // NewStream creates a new Stream for the client side. This is typically
  118. // called by generated code. ctx is used for the lifetime of the stream.
  119. //
  120. // To ensure resources are not leaked due to the stream returned, one of the following
  121. // actions must be performed:
  122. //
  123. // 1. Call Close on the ClientConn.
  124. // 2. Cancel the context provided.
  125. // 3. Call RecvMsg until a non-nil error is returned. A protobuf-generated
  126. // client-streaming RPC, for instance, might use the helper function
  127. // CloseAndRecv (note that CloseSend does not Recv, therefore is not
  128. // guaranteed to release all resources).
  129. // 4. Receive a non-nil, non-io.EOF error from Header or SendMsg.
  130. //
  131. // If none of the above happen, a goroutine and a context will be leaked, and grpc
  132. // will not call the optionally-configured stats handler with a stats.End message.
  133. func (cc *ClientConn) NewStream(ctx context.Context, desc *StreamDesc, method string, opts ...CallOption) (ClientStream, error) {
  134. // allow interceptor to see all applicable call options, which means those
  135. // configured as defaults from dial option as well as per-call options
  136. opts = combine(cc.dopts.callOptions, opts)
  137. if cc.dopts.streamInt != nil {
  138. return cc.dopts.streamInt(ctx, desc, cc, method, newClientStream, opts...)
  139. }
  140. return newClientStream(ctx, desc, cc, method, opts...)
  141. }
  142. // NewClientStream is a wrapper for ClientConn.NewStream.
  143. func NewClientStream(ctx context.Context, desc *StreamDesc, cc *ClientConn, method string, opts ...CallOption) (ClientStream, error) {
  144. return cc.NewStream(ctx, desc, method, opts...)
  145. }
  146. func newClientStream(ctx context.Context, desc *StreamDesc, cc *ClientConn, method string, opts ...CallOption) (_ ClientStream, err error) {
  147. if channelz.IsOn() {
  148. cc.incrCallsStarted()
  149. defer func() {
  150. if err != nil {
  151. cc.incrCallsFailed()
  152. }
  153. }()
  154. }
  155. c := defaultCallInfo()
  156. // Provide an opportunity for the first RPC to see the first service config
  157. // provided by the resolver.
  158. if err := cc.waitForResolvedAddrs(ctx); err != nil {
  159. return nil, err
  160. }
  161. mc := cc.GetMethodConfig(method)
  162. if mc.WaitForReady != nil {
  163. c.failFast = !*mc.WaitForReady
  164. }
  165. // Possible context leak:
  166. // The cancel function for the child context we create will only be called
  167. // when RecvMsg returns a non-nil error, if the ClientConn is closed, or if
  168. // an error is generated by SendMsg.
  169. // https://github.com/grpc/grpc-go/issues/1818.
  170. var cancel context.CancelFunc
  171. if mc.Timeout != nil && *mc.Timeout >= 0 {
  172. ctx, cancel = context.WithTimeout(ctx, *mc.Timeout)
  173. } else {
  174. ctx, cancel = context.WithCancel(ctx)
  175. }
  176. defer func() {
  177. if err != nil {
  178. cancel()
  179. }
  180. }()
  181. for _, o := range opts {
  182. if err := o.before(c); err != nil {
  183. return nil, toRPCErr(err)
  184. }
  185. }
  186. c.maxSendMessageSize = getMaxSize(mc.MaxReqSize, c.maxSendMessageSize, defaultClientMaxSendMessageSize)
  187. c.maxReceiveMessageSize = getMaxSize(mc.MaxRespSize, c.maxReceiveMessageSize, defaultClientMaxReceiveMessageSize)
  188. if err := setCallInfoCodec(c); err != nil {
  189. return nil, err
  190. }
  191. callHdr := &transport.CallHdr{
  192. Host: cc.authority,
  193. Method: method,
  194. ContentSubtype: c.contentSubtype,
  195. }
  196. // Set our outgoing compression according to the UseCompressor CallOption, if
  197. // set. In that case, also find the compressor from the encoding package.
  198. // Otherwise, use the compressor configured by the WithCompressor DialOption,
  199. // if set.
  200. var cp Compressor
  201. var comp encoding.Compressor
  202. if ct := c.compressorType; ct != "" {
  203. callHdr.SendCompress = ct
  204. if ct != encoding.Identity {
  205. comp = encoding.GetCompressor(ct)
  206. if comp == nil {
  207. return nil, status.Errorf(codes.Internal, "grpc: Compressor is not installed for requested grpc-encoding %q", ct)
  208. }
  209. }
  210. } else if cc.dopts.cp != nil {
  211. callHdr.SendCompress = cc.dopts.cp.Type()
  212. cp = cc.dopts.cp
  213. }
  214. if c.creds != nil {
  215. callHdr.Creds = c.creds
  216. }
  217. var trInfo traceInfo
  218. if EnableTracing {
  219. trInfo.tr = trace.New("grpc.Sent."+methodFamily(method), method)
  220. trInfo.firstLine.client = true
  221. if deadline, ok := ctx.Deadline(); ok {
  222. trInfo.firstLine.deadline = time.Until(deadline)
  223. }
  224. trInfo.tr.LazyLog(&trInfo.firstLine, false)
  225. ctx = trace.NewContext(ctx, trInfo.tr)
  226. }
  227. ctx = newContextWithRPCInfo(ctx, c.failFast)
  228. sh := cc.dopts.copts.StatsHandler
  229. var beginTime time.Time
  230. if sh != nil {
  231. ctx = sh.TagRPC(ctx, &stats.RPCTagInfo{FullMethodName: method, FailFast: c.failFast})
  232. beginTime = time.Now()
  233. begin := &stats.Begin{
  234. Client: true,
  235. BeginTime: beginTime,
  236. FailFast: c.failFast,
  237. }
  238. sh.HandleRPC(ctx, begin)
  239. }
  240. cs := &clientStream{
  241. callHdr: callHdr,
  242. ctx: ctx,
  243. methodConfig: &mc,
  244. opts: opts,
  245. callInfo: c,
  246. cc: cc,
  247. desc: desc,
  248. codec: c.codec,
  249. cp: cp,
  250. comp: comp,
  251. cancel: cancel,
  252. beginTime: beginTime,
  253. firstAttempt: true,
  254. }
  255. if !cc.dopts.disableRetry {
  256. cs.retryThrottler = cc.retryThrottler.Load().(*retryThrottler)
  257. }
  258. cs.binlog = binarylog.GetMethodLogger(method)
  259. cs.callInfo.stream = cs
  260. // Only this initial attempt has stats/tracing.
  261. // TODO(dfawley): move to newAttempt when per-attempt stats are implemented.
  262. if err := cs.newAttemptLocked(sh, trInfo); err != nil {
  263. cs.finish(err)
  264. return nil, err
  265. }
  266. op := func(a *csAttempt) error { return a.newStream() }
  267. if err := cs.withRetry(op, func() { cs.bufferForRetryLocked(0, op) }); err != nil {
  268. cs.finish(err)
  269. return nil, err
  270. }
  271. if cs.binlog != nil {
  272. md, _ := metadata.FromOutgoingContext(ctx)
  273. logEntry := &binarylog.ClientHeader{
  274. OnClientSide: true,
  275. Header: md,
  276. MethodName: method,
  277. Authority: cs.cc.authority,
  278. }
  279. if deadline, ok := ctx.Deadline(); ok {
  280. logEntry.Timeout = time.Until(deadline)
  281. if logEntry.Timeout < 0 {
  282. logEntry.Timeout = 0
  283. }
  284. }
  285. cs.binlog.Log(logEntry)
  286. }
  287. if desc != unaryStreamDesc {
  288. // Listen on cc and stream contexts to cleanup when the user closes the
  289. // ClientConn or cancels the stream context. In all other cases, an error
  290. // should already be injected into the recv buffer by the transport, which
  291. // the client will eventually receive, and then we will cancel the stream's
  292. // context in clientStream.finish.
  293. go func() {
  294. select {
  295. case <-cc.ctx.Done():
  296. cs.finish(ErrClientConnClosing)
  297. case <-ctx.Done():
  298. cs.finish(toRPCErr(ctx.Err()))
  299. }
  300. }()
  301. }
  302. return cs, nil
  303. }
  304. func (cs *clientStream) newAttemptLocked(sh stats.Handler, trInfo traceInfo) error {
  305. cs.attempt = &csAttempt{
  306. cs: cs,
  307. dc: cs.cc.dopts.dc,
  308. statsHandler: sh,
  309. trInfo: trInfo,
  310. }
  311. if err := cs.ctx.Err(); err != nil {
  312. return toRPCErr(err)
  313. }
  314. t, done, err := cs.cc.getTransport(cs.ctx, cs.callInfo.failFast, cs.callHdr.Method)
  315. if err != nil {
  316. return err
  317. }
  318. cs.attempt.t = t
  319. cs.attempt.done = done
  320. return nil
  321. }
  322. func (a *csAttempt) newStream() error {
  323. cs := a.cs
  324. cs.callHdr.PreviousAttempts = cs.numRetries
  325. s, err := a.t.NewStream(cs.ctx, cs.callHdr)
  326. if err != nil {
  327. return toRPCErr(err)
  328. }
  329. cs.attempt.s = s
  330. cs.attempt.p = &parser{r: s}
  331. return nil
  332. }
  333. // clientStream implements a client side Stream.
  334. type clientStream struct {
  335. callHdr *transport.CallHdr
  336. opts []CallOption
  337. callInfo *callInfo
  338. cc *ClientConn
  339. desc *StreamDesc
  340. codec baseCodec
  341. cp Compressor
  342. comp encoding.Compressor
  343. cancel context.CancelFunc // cancels all attempts
  344. sentLast bool // sent an end stream
  345. beginTime time.Time
  346. methodConfig *MethodConfig
  347. ctx context.Context // the application's context, wrapped by stats/tracing
  348. retryThrottler *retryThrottler // The throttler active when the RPC began.
  349. binlog *binarylog.MethodLogger // Binary logger, can be nil.
  350. // serverHeaderBinlogged is a boolean for whether server header has been
  351. // logged. Server header will be logged when the first time one of those
  352. // happens: stream.Header(), stream.Recv().
  353. //
  354. // It's only read and used by Recv() and Header(), so it doesn't need to be
  355. // synchronized.
  356. serverHeaderBinlogged bool
  357. mu sync.Mutex
  358. firstAttempt bool // if true, transparent retry is valid
  359. numRetries int // exclusive of transparent retry attempt(s)
  360. numRetriesSincePushback int // retries since pushback; to reset backoff
  361. finished bool // TODO: replace with atomic cmpxchg or sync.Once?
  362. attempt *csAttempt // the active client stream attempt
  363. // TODO(hedging): hedging will have multiple attempts simultaneously.
  364. committed bool // active attempt committed for retry?
  365. buffer []func(a *csAttempt) error // operations to replay on retry
  366. bufferSize int // current size of buffer
  367. }
  368. // csAttempt implements a single transport stream attempt within a
  369. // clientStream.
  370. type csAttempt struct {
  371. cs *clientStream
  372. t transport.ClientTransport
  373. s *transport.Stream
  374. p *parser
  375. done func(balancer.DoneInfo)
  376. finished bool
  377. dc Decompressor
  378. decomp encoding.Compressor
  379. decompSet bool
  380. mu sync.Mutex // guards trInfo.tr
  381. // trInfo.tr is set when created (if EnableTracing is true),
  382. // and cleared when the finish method is called.
  383. trInfo traceInfo
  384. statsHandler stats.Handler
  385. }
  386. func (cs *clientStream) commitAttemptLocked() {
  387. cs.committed = true
  388. cs.buffer = nil
  389. }
  390. func (cs *clientStream) commitAttempt() {
  391. cs.mu.Lock()
  392. cs.commitAttemptLocked()
  393. cs.mu.Unlock()
  394. }
  395. // shouldRetry returns nil if the RPC should be retried; otherwise it returns
  396. // the error that should be returned by the operation.
  397. func (cs *clientStream) shouldRetry(err error) error {
  398. if cs.attempt.s == nil && !cs.callInfo.failFast {
  399. // In the event of any error from NewStream (attempt.s == nil), we
  400. // never attempted to write anything to the wire, so we can retry
  401. // indefinitely for non-fail-fast RPCs.
  402. return nil
  403. }
  404. if cs.finished || cs.committed {
  405. // RPC is finished or committed; cannot retry.
  406. return err
  407. }
  408. // Wait for the trailers.
  409. if cs.attempt.s != nil {
  410. <-cs.attempt.s.Done()
  411. }
  412. if cs.firstAttempt && !cs.callInfo.failFast && (cs.attempt.s == nil || cs.attempt.s.Unprocessed()) {
  413. // First attempt, wait-for-ready, stream unprocessed: transparently retry.
  414. cs.firstAttempt = false
  415. return nil
  416. }
  417. cs.firstAttempt = false
  418. if cs.cc.dopts.disableRetry {
  419. return err
  420. }
  421. pushback := 0
  422. hasPushback := false
  423. if cs.attempt.s != nil {
  424. if to, toErr := cs.attempt.s.TrailersOnly(); toErr != nil || !to {
  425. return err
  426. }
  427. // TODO(retry): Move down if the spec changes to not check server pushback
  428. // before considering this a failure for throttling.
  429. sps := cs.attempt.s.Trailer()["grpc-retry-pushback-ms"]
  430. if len(sps) == 1 {
  431. var e error
  432. if pushback, e = strconv.Atoi(sps[0]); e != nil || pushback < 0 {
  433. grpclog.Infof("Server retry pushback specified to abort (%q).", sps[0])
  434. cs.retryThrottler.throttle() // This counts as a failure for throttling.
  435. return err
  436. }
  437. hasPushback = true
  438. } else if len(sps) > 1 {
  439. grpclog.Warningf("Server retry pushback specified multiple values (%q); not retrying.", sps)
  440. cs.retryThrottler.throttle() // This counts as a failure for throttling.
  441. return err
  442. }
  443. }
  444. var code codes.Code
  445. if cs.attempt.s != nil {
  446. code = cs.attempt.s.Status().Code()
  447. } else {
  448. code = status.Convert(err).Code()
  449. }
  450. rp := cs.methodConfig.retryPolicy
  451. if rp == nil || !rp.retryableStatusCodes[code] {
  452. return err
  453. }
  454. // Note: the ordering here is important; we count this as a failure
  455. // only if the code matched a retryable code.
  456. if cs.retryThrottler.throttle() {
  457. return err
  458. }
  459. if cs.numRetries+1 >= rp.maxAttempts {
  460. return err
  461. }
  462. var dur time.Duration
  463. if hasPushback {
  464. dur = time.Millisecond * time.Duration(pushback)
  465. cs.numRetriesSincePushback = 0
  466. } else {
  467. fact := math.Pow(rp.backoffMultiplier, float64(cs.numRetriesSincePushback))
  468. cur := float64(rp.initialBackoff) * fact
  469. if max := float64(rp.maxBackoff); cur > max {
  470. cur = max
  471. }
  472. dur = time.Duration(grpcrand.Int63n(int64(cur)))
  473. cs.numRetriesSincePushback++
  474. }
  475. // TODO(dfawley): we could eagerly fail here if dur puts us past the
  476. // deadline, but unsure if it is worth doing.
  477. t := time.NewTimer(dur)
  478. select {
  479. case <-t.C:
  480. cs.numRetries++
  481. return nil
  482. case <-cs.ctx.Done():
  483. t.Stop()
  484. return status.FromContextError(cs.ctx.Err()).Err()
  485. }
  486. }
  487. // Returns nil if a retry was performed and succeeded; error otherwise.
  488. func (cs *clientStream) retryLocked(lastErr error) error {
  489. for {
  490. cs.attempt.finish(lastErr)
  491. if err := cs.shouldRetry(lastErr); err != nil {
  492. cs.commitAttemptLocked()
  493. return err
  494. }
  495. if err := cs.newAttemptLocked(nil, traceInfo{}); err != nil {
  496. return err
  497. }
  498. if lastErr = cs.replayBufferLocked(); lastErr == nil {
  499. return nil
  500. }
  501. }
  502. }
  503. func (cs *clientStream) Context() context.Context {
  504. cs.commitAttempt()
  505. // No need to lock before using attempt, since we know it is committed and
  506. // cannot change.
  507. return cs.attempt.s.Context()
  508. }
  509. func (cs *clientStream) withRetry(op func(a *csAttempt) error, onSuccess func()) error {
  510. cs.mu.Lock()
  511. for {
  512. if cs.committed {
  513. cs.mu.Unlock()
  514. return op(cs.attempt)
  515. }
  516. a := cs.attempt
  517. cs.mu.Unlock()
  518. err := op(a)
  519. cs.mu.Lock()
  520. if a != cs.attempt {
  521. // We started another attempt already.
  522. continue
  523. }
  524. if err == io.EOF {
  525. <-a.s.Done()
  526. }
  527. if err == nil || (err == io.EOF && a.s.Status().Code() == codes.OK) {
  528. onSuccess()
  529. cs.mu.Unlock()
  530. return err
  531. }
  532. if err := cs.retryLocked(err); err != nil {
  533. cs.mu.Unlock()
  534. return err
  535. }
  536. }
  537. }
  538. func (cs *clientStream) Header() (metadata.MD, error) {
  539. var m metadata.MD
  540. err := cs.withRetry(func(a *csAttempt) error {
  541. var err error
  542. m, err = a.s.Header()
  543. return toRPCErr(err)
  544. }, cs.commitAttemptLocked)
  545. if err != nil {
  546. cs.finish(err)
  547. return nil, err
  548. }
  549. if cs.binlog != nil && !cs.serverHeaderBinlogged {
  550. // Only log if binary log is on and header has not been logged.
  551. logEntry := &binarylog.ServerHeader{
  552. OnClientSide: true,
  553. Header: m,
  554. PeerAddr: nil,
  555. }
  556. if peer, ok := peer.FromContext(cs.Context()); ok {
  557. logEntry.PeerAddr = peer.Addr
  558. }
  559. cs.binlog.Log(logEntry)
  560. cs.serverHeaderBinlogged = true
  561. }
  562. return m, err
  563. }
  564. func (cs *clientStream) Trailer() metadata.MD {
  565. // On RPC failure, we never need to retry, because usage requires that
  566. // RecvMsg() returned a non-nil error before calling this function is valid.
  567. // We would have retried earlier if necessary.
  568. //
  569. // Commit the attempt anyway, just in case users are not following those
  570. // directions -- it will prevent races and should not meaningfully impact
  571. // performance.
  572. cs.commitAttempt()
  573. if cs.attempt.s == nil {
  574. return nil
  575. }
  576. return cs.attempt.s.Trailer()
  577. }
  578. func (cs *clientStream) replayBufferLocked() error {
  579. a := cs.attempt
  580. for _, f := range cs.buffer {
  581. if err := f(a); err != nil {
  582. return err
  583. }
  584. }
  585. return nil
  586. }
  587. func (cs *clientStream) bufferForRetryLocked(sz int, op func(a *csAttempt) error) {
  588. // Note: we still will buffer if retry is disabled (for transparent retries).
  589. if cs.committed {
  590. return
  591. }
  592. cs.bufferSize += sz
  593. if cs.bufferSize > cs.callInfo.maxRetryRPCBufferSize {
  594. cs.commitAttemptLocked()
  595. return
  596. }
  597. cs.buffer = append(cs.buffer, op)
  598. }
  599. func (cs *clientStream) SendMsg(m interface{}) (err error) {
  600. defer func() {
  601. if err != nil && err != io.EOF {
  602. // Call finish on the client stream for errors generated by this SendMsg
  603. // call, as these indicate problems created by this client. (Transport
  604. // errors are converted to an io.EOF error in csAttempt.sendMsg; the real
  605. // error will be returned from RecvMsg eventually in that case, or be
  606. // retried.)
  607. cs.finish(err)
  608. }
  609. }()
  610. if cs.sentLast {
  611. return status.Errorf(codes.Internal, "SendMsg called after CloseSend")
  612. }
  613. if !cs.desc.ClientStreams {
  614. cs.sentLast = true
  615. }
  616. data, err := encode(cs.codec, m)
  617. if err != nil {
  618. return err
  619. }
  620. compData, err := compress(data, cs.cp, cs.comp)
  621. if err != nil {
  622. return err
  623. }
  624. hdr, payload := msgHeader(data, compData)
  625. // TODO(dfawley): should we be checking len(data) instead?
  626. if len(payload) > *cs.callInfo.maxSendMessageSize {
  627. return status.Errorf(codes.ResourceExhausted, "trying to send message larger than max (%d vs. %d)", len(payload), *cs.callInfo.maxSendMessageSize)
  628. }
  629. msgBytes := data // Store the pointer before setting to nil. For binary logging.
  630. op := func(a *csAttempt) error {
  631. err := a.sendMsg(m, hdr, payload, data)
  632. // nil out the message and uncomp when replaying; they are only needed for
  633. // stats which is disabled for subsequent attempts.
  634. m, data = nil, nil
  635. return err
  636. }
  637. err = cs.withRetry(op, func() { cs.bufferForRetryLocked(len(hdr)+len(payload), op) })
  638. if cs.binlog != nil && err == nil {
  639. cs.binlog.Log(&binarylog.ClientMessage{
  640. OnClientSide: true,
  641. Message: msgBytes,
  642. })
  643. }
  644. return
  645. }
  646. func (cs *clientStream) RecvMsg(m interface{}) error {
  647. if cs.binlog != nil && !cs.serverHeaderBinlogged {
  648. // Call Header() to binary log header if it's not already logged.
  649. cs.Header()
  650. }
  651. var recvInfo *payloadInfo
  652. if cs.binlog != nil {
  653. recvInfo = &payloadInfo{}
  654. }
  655. err := cs.withRetry(func(a *csAttempt) error {
  656. return a.recvMsg(m, recvInfo)
  657. }, cs.commitAttemptLocked)
  658. if cs.binlog != nil && err == nil {
  659. cs.binlog.Log(&binarylog.ServerMessage{
  660. OnClientSide: true,
  661. Message: recvInfo.uncompressedBytes,
  662. })
  663. }
  664. if err != nil || !cs.desc.ServerStreams {
  665. // err != nil or non-server-streaming indicates end of stream.
  666. cs.finish(err)
  667. if cs.binlog != nil {
  668. // finish will not log Trailer. Log Trailer here.
  669. logEntry := &binarylog.ServerTrailer{
  670. OnClientSide: true,
  671. Trailer: cs.Trailer(),
  672. Err: err,
  673. }
  674. if logEntry.Err == io.EOF {
  675. logEntry.Err = nil
  676. }
  677. if peer, ok := peer.FromContext(cs.Context()); ok {
  678. logEntry.PeerAddr = peer.Addr
  679. }
  680. cs.binlog.Log(logEntry)
  681. }
  682. }
  683. return err
  684. }
  685. func (cs *clientStream) CloseSend() error {
  686. if cs.sentLast {
  687. // TODO: return an error and finish the stream instead, due to API misuse?
  688. return nil
  689. }
  690. cs.sentLast = true
  691. op := func(a *csAttempt) error {
  692. a.t.Write(a.s, nil, nil, &transport.Options{Last: true})
  693. // Always return nil; io.EOF is the only error that might make sense
  694. // instead, but there is no need to signal the client to call RecvMsg
  695. // as the only use left for the stream after CloseSend is to call
  696. // RecvMsg. This also matches historical behavior.
  697. return nil
  698. }
  699. cs.withRetry(op, func() { cs.bufferForRetryLocked(0, op) })
  700. if cs.binlog != nil {
  701. cs.binlog.Log(&binarylog.ClientHalfClose{
  702. OnClientSide: true,
  703. })
  704. }
  705. // We never returned an error here for reasons.
  706. return nil
  707. }
  708. func (cs *clientStream) finish(err error) {
  709. if err == io.EOF {
  710. // Ending a stream with EOF indicates a success.
  711. err = nil
  712. }
  713. cs.mu.Lock()
  714. if cs.finished {
  715. cs.mu.Unlock()
  716. return
  717. }
  718. cs.finished = true
  719. cs.commitAttemptLocked()
  720. cs.mu.Unlock()
  721. // For binary logging. only log cancel in finish (could be caused by RPC ctx
  722. // canceled or ClientConn closed). Trailer will be logged in RecvMsg.
  723. //
  724. // Only one of cancel or trailer needs to be logged. In the cases where
  725. // users don't call RecvMsg, users must have already canceled the RPC.
  726. if cs.binlog != nil && status.Code(err) == codes.Canceled {
  727. cs.binlog.Log(&binarylog.Cancel{
  728. OnClientSide: true,
  729. })
  730. }
  731. if err == nil {
  732. cs.retryThrottler.successfulRPC()
  733. }
  734. if channelz.IsOn() {
  735. if err != nil {
  736. cs.cc.incrCallsFailed()
  737. } else {
  738. cs.cc.incrCallsSucceeded()
  739. }
  740. }
  741. if cs.attempt != nil {
  742. cs.attempt.finish(err)
  743. }
  744. // after functions all rely upon having a stream.
  745. if cs.attempt.s != nil {
  746. for _, o := range cs.opts {
  747. o.after(cs.callInfo)
  748. }
  749. }
  750. cs.cancel()
  751. }
  752. func (a *csAttempt) sendMsg(m interface{}, hdr, payld, data []byte) error {
  753. cs := a.cs
  754. if EnableTracing {
  755. a.mu.Lock()
  756. if a.trInfo.tr != nil {
  757. a.trInfo.tr.LazyLog(&payload{sent: true, msg: m}, true)
  758. }
  759. a.mu.Unlock()
  760. }
  761. if err := a.t.Write(a.s, hdr, payld, &transport.Options{Last: !cs.desc.ClientStreams}); err != nil {
  762. if !cs.desc.ClientStreams {
  763. // For non-client-streaming RPCs, we return nil instead of EOF on error
  764. // because the generated code requires it. finish is not called; RecvMsg()
  765. // will call it with the stream's status independently.
  766. return nil
  767. }
  768. return io.EOF
  769. }
  770. if a.statsHandler != nil {
  771. a.statsHandler.HandleRPC(cs.ctx, outPayload(true, m, data, payld, time.Now()))
  772. }
  773. if channelz.IsOn() {
  774. a.t.IncrMsgSent()
  775. }
  776. return nil
  777. }
  778. func (a *csAttempt) recvMsg(m interface{}, payInfo *payloadInfo) (err error) {
  779. cs := a.cs
  780. if a.statsHandler != nil && payInfo == nil {
  781. payInfo = &payloadInfo{}
  782. }
  783. if !a.decompSet {
  784. // Block until we receive headers containing received message encoding.
  785. if ct := a.s.RecvCompress(); ct != "" && ct != encoding.Identity {
  786. if a.dc == nil || a.dc.Type() != ct {
  787. // No configured decompressor, or it does not match the incoming
  788. // message encoding; attempt to find a registered compressor that does.
  789. a.dc = nil
  790. a.decomp = encoding.GetCompressor(ct)
  791. }
  792. } else {
  793. // No compression is used; disable our decompressor.
  794. a.dc = nil
  795. }
  796. // Only initialize this state once per stream.
  797. a.decompSet = true
  798. }
  799. err = recv(a.p, cs.codec, a.s, a.dc, m, *cs.callInfo.maxReceiveMessageSize, payInfo, a.decomp)
  800. if err != nil {
  801. if err == io.EOF {
  802. if statusErr := a.s.Status().Err(); statusErr != nil {
  803. return statusErr
  804. }
  805. return io.EOF // indicates successful end of stream.
  806. }
  807. return toRPCErr(err)
  808. }
  809. if EnableTracing {
  810. a.mu.Lock()
  811. if a.trInfo.tr != nil {
  812. a.trInfo.tr.LazyLog(&payload{sent: false, msg: m}, true)
  813. }
  814. a.mu.Unlock()
  815. }
  816. if a.statsHandler != nil {
  817. a.statsHandler.HandleRPC(cs.ctx, &stats.InPayload{
  818. Client: true,
  819. RecvTime: time.Now(),
  820. Payload: m,
  821. // TODO truncate large payload.
  822. Data: payInfo.uncompressedBytes,
  823. Length: len(payInfo.uncompressedBytes),
  824. })
  825. }
  826. if channelz.IsOn() {
  827. a.t.IncrMsgRecv()
  828. }
  829. if cs.desc.ServerStreams {
  830. // Subsequent messages should be received by subsequent RecvMsg calls.
  831. return nil
  832. }
  833. // Special handling for non-server-stream rpcs.
  834. // This recv expects EOF or errors, so we don't collect inPayload.
  835. err = recv(a.p, cs.codec, a.s, a.dc, m, *cs.callInfo.maxReceiveMessageSize, nil, a.decomp)
  836. if err == nil {
  837. return toRPCErr(errors.New("grpc: client streaming protocol violation: get <nil>, want <EOF>"))
  838. }
  839. if err == io.EOF {
  840. return a.s.Status().Err() // non-server streaming Recv returns nil on success
  841. }
  842. return toRPCErr(err)
  843. }
  844. func (a *csAttempt) finish(err error) {
  845. a.mu.Lock()
  846. if a.finished {
  847. a.mu.Unlock()
  848. return
  849. }
  850. a.finished = true
  851. if err == io.EOF {
  852. // Ending a stream with EOF indicates a success.
  853. err = nil
  854. }
  855. if a.s != nil {
  856. a.t.CloseStream(a.s, err)
  857. }
  858. if a.done != nil {
  859. br := false
  860. var tr metadata.MD
  861. if a.s != nil {
  862. br = a.s.BytesReceived()
  863. tr = a.s.Trailer()
  864. }
  865. a.done(balancer.DoneInfo{
  866. Err: err,
  867. Trailer: tr,
  868. BytesSent: a.s != nil,
  869. BytesReceived: br,
  870. })
  871. }
  872. if a.statsHandler != nil {
  873. end := &stats.End{
  874. Client: true,
  875. BeginTime: a.cs.beginTime,
  876. EndTime: time.Now(),
  877. Error: err,
  878. }
  879. a.statsHandler.HandleRPC(a.cs.ctx, end)
  880. }
  881. if a.trInfo.tr != nil {
  882. if err == nil {
  883. a.trInfo.tr.LazyPrintf("RPC: [OK]")
  884. } else {
  885. a.trInfo.tr.LazyPrintf("RPC: [%v]", err)
  886. a.trInfo.tr.SetError()
  887. }
  888. a.trInfo.tr.Finish()
  889. a.trInfo.tr = nil
  890. }
  891. a.mu.Unlock()
  892. }
  893. func (ac *addrConn) newClientStream(ctx context.Context, desc *StreamDesc, method string, t transport.ClientTransport, opts ...CallOption) (_ ClientStream, err error) {
  894. ac.mu.Lock()
  895. if ac.transport != t {
  896. ac.mu.Unlock()
  897. return nil, status.Error(codes.Canceled, "the provided transport is no longer valid to use")
  898. }
  899. // transition to CONNECTING state when an attempt starts
  900. if ac.state != connectivity.Connecting {
  901. ac.updateConnectivityState(connectivity.Connecting)
  902. ac.cc.handleSubConnStateChange(ac.acbw, ac.state)
  903. }
  904. ac.mu.Unlock()
  905. if t == nil {
  906. // TODO: return RPC error here?
  907. return nil, errors.New("transport provided is nil")
  908. }
  909. // defaultCallInfo contains unnecessary info(i.e. failfast, maxRetryRPCBufferSize), so we just initialize an empty struct.
  910. c := &callInfo{}
  911. for _, o := range opts {
  912. if err := o.before(c); err != nil {
  913. return nil, toRPCErr(err)
  914. }
  915. }
  916. c.maxReceiveMessageSize = getMaxSize(nil, c.maxReceiveMessageSize, defaultClientMaxReceiveMessageSize)
  917. c.maxSendMessageSize = getMaxSize(nil, c.maxSendMessageSize, defaultServerMaxSendMessageSize)
  918. // Possible context leak:
  919. // The cancel function for the child context we create will only be called
  920. // when RecvMsg returns a non-nil error, if the ClientConn is closed, or if
  921. // an error is generated by SendMsg.
  922. // https://github.com/grpc/grpc-go/issues/1818.
  923. ctx, cancel := context.WithCancel(ctx)
  924. defer func() {
  925. if err != nil {
  926. cancel()
  927. }
  928. }()
  929. if err := setCallInfoCodec(c); err != nil {
  930. return nil, err
  931. }
  932. callHdr := &transport.CallHdr{
  933. Host: ac.cc.authority,
  934. Method: method,
  935. ContentSubtype: c.contentSubtype,
  936. }
  937. // Set our outgoing compression according to the UseCompressor CallOption, if
  938. // set. In that case, also find the compressor from the encoding package.
  939. // Otherwise, use the compressor configured by the WithCompressor DialOption,
  940. // if set.
  941. var cp Compressor
  942. var comp encoding.Compressor
  943. if ct := c.compressorType; ct != "" {
  944. callHdr.SendCompress = ct
  945. if ct != encoding.Identity {
  946. comp = encoding.GetCompressor(ct)
  947. if comp == nil {
  948. return nil, status.Errorf(codes.Internal, "grpc: Compressor is not installed for requested grpc-encoding %q", ct)
  949. }
  950. }
  951. } else if ac.cc.dopts.cp != nil {
  952. callHdr.SendCompress = ac.cc.dopts.cp.Type()
  953. cp = ac.cc.dopts.cp
  954. }
  955. if c.creds != nil {
  956. callHdr.Creds = c.creds
  957. }
  958. as := &addrConnStream{
  959. callHdr: callHdr,
  960. ac: ac,
  961. ctx: ctx,
  962. cancel: cancel,
  963. opts: opts,
  964. callInfo: c,
  965. desc: desc,
  966. codec: c.codec,
  967. cp: cp,
  968. comp: comp,
  969. t: t,
  970. }
  971. as.callInfo.stream = as
  972. s, err := as.t.NewStream(as.ctx, as.callHdr)
  973. if err != nil {
  974. err = toRPCErr(err)
  975. return nil, err
  976. }
  977. as.s = s
  978. as.p = &parser{r: s}
  979. ac.incrCallsStarted()
  980. if desc != unaryStreamDesc {
  981. // Listen on cc and stream contexts to cleanup when the user closes the
  982. // ClientConn or cancels the stream context. In all other cases, an error
  983. // should already be injected into the recv buffer by the transport, which
  984. // the client will eventually receive, and then we will cancel the stream's
  985. // context in clientStream.finish.
  986. go func() {
  987. select {
  988. case <-ac.ctx.Done():
  989. as.finish(status.Error(codes.Canceled, "grpc: the SubConn is closing"))
  990. case <-ctx.Done():
  991. as.finish(toRPCErr(ctx.Err()))
  992. }
  993. }()
  994. }
  995. return as, nil
  996. }
  997. type addrConnStream struct {
  998. s *transport.Stream
  999. ac *addrConn
  1000. callHdr *transport.CallHdr
  1001. cancel context.CancelFunc
  1002. opts []CallOption
  1003. callInfo *callInfo
  1004. t transport.ClientTransport
  1005. ctx context.Context
  1006. sentLast bool
  1007. desc *StreamDesc
  1008. codec baseCodec
  1009. cp Compressor
  1010. comp encoding.Compressor
  1011. decompSet bool
  1012. dc Decompressor
  1013. decomp encoding.Compressor
  1014. p *parser
  1015. mu sync.Mutex
  1016. finished bool
  1017. }
  1018. func (as *addrConnStream) Header() (metadata.MD, error) {
  1019. m, err := as.s.Header()
  1020. if err != nil {
  1021. as.finish(toRPCErr(err))
  1022. }
  1023. return m, err
  1024. }
  1025. func (as *addrConnStream) Trailer() metadata.MD {
  1026. return as.s.Trailer()
  1027. }
  1028. func (as *addrConnStream) CloseSend() error {
  1029. if as.sentLast {
  1030. // TODO: return an error and finish the stream instead, due to API misuse?
  1031. return nil
  1032. }
  1033. as.sentLast = true
  1034. as.t.Write(as.s, nil, nil, &transport.Options{Last: true})
  1035. // Always return nil; io.EOF is the only error that might make sense
  1036. // instead, but there is no need to signal the client to call RecvMsg
  1037. // as the only use left for the stream after CloseSend is to call
  1038. // RecvMsg. This also matches historical behavior.
  1039. return nil
  1040. }
  1041. func (as *addrConnStream) Context() context.Context {
  1042. return as.s.Context()
  1043. }
  1044. func (as *addrConnStream) SendMsg(m interface{}) (err error) {
  1045. defer func() {
  1046. if err != nil && err != io.EOF {
  1047. // Call finish on the client stream for errors generated by this SendMsg
  1048. // call, as these indicate problems created by this client. (Transport
  1049. // errors are converted to an io.EOF error in csAttempt.sendMsg; the real
  1050. // error will be returned from RecvMsg eventually in that case, or be
  1051. // retried.)
  1052. as.finish(err)
  1053. }
  1054. }()
  1055. if as.sentLast {
  1056. return status.Errorf(codes.Internal, "SendMsg called after CloseSend")
  1057. }
  1058. if !as.desc.ClientStreams {
  1059. as.sentLast = true
  1060. }
  1061. data, err := encode(as.codec, m)
  1062. if err != nil {
  1063. return err
  1064. }
  1065. compData, err := compress(data, as.cp, as.comp)
  1066. if err != nil {
  1067. return err
  1068. }
  1069. hdr, payld := msgHeader(data, compData)
  1070. // TODO(dfawley): should we be checking len(data) instead?
  1071. if len(payld) > *as.callInfo.maxSendMessageSize {
  1072. return status.Errorf(codes.ResourceExhausted, "trying to send message larger than max (%d vs. %d)", len(payld), *as.callInfo.maxSendMessageSize)
  1073. }
  1074. if err := as.t.Write(as.s, hdr, payld, &transport.Options{Last: !as.desc.ClientStreams}); err != nil {
  1075. if !as.desc.ClientStreams {
  1076. // For non-client-streaming RPCs, we return nil instead of EOF on error
  1077. // because the generated code requires it. finish is not called; RecvMsg()
  1078. // will call it with the stream's status independently.
  1079. return nil
  1080. }
  1081. return io.EOF
  1082. }
  1083. if channelz.IsOn() {
  1084. as.t.IncrMsgSent()
  1085. }
  1086. return nil
  1087. }
  1088. func (as *addrConnStream) RecvMsg(m interface{}) (err error) {
  1089. defer func() {
  1090. if err != nil || !as.desc.ServerStreams {
  1091. // err != nil or non-server-streaming indicates end of stream.
  1092. as.finish(err)
  1093. }
  1094. }()
  1095. if !as.decompSet {
  1096. // Block until we receive headers containing received message encoding.
  1097. if ct := as.s.RecvCompress(); ct != "" && ct != encoding.Identity {
  1098. if as.dc == nil || as.dc.Type() != ct {
  1099. // No configured decompressor, or it does not match the incoming
  1100. // message encoding; attempt to find a registered compressor that does.
  1101. as.dc = nil
  1102. as.decomp = encoding.GetCompressor(ct)
  1103. }
  1104. } else {
  1105. // No compression is used; disable our decompressor.
  1106. as.dc = nil
  1107. }
  1108. // Only initialize this state once per stream.
  1109. as.decompSet = true
  1110. }
  1111. err = recv(as.p, as.codec, as.s, as.dc, m, *as.callInfo.maxReceiveMessageSize, nil, as.decomp)
  1112. if err != nil {
  1113. if err == io.EOF {
  1114. if statusErr := as.s.Status().Err(); statusErr != nil {
  1115. return statusErr
  1116. }
  1117. return io.EOF // indicates successful end of stream.
  1118. }
  1119. return toRPCErr(err)
  1120. }
  1121. if channelz.IsOn() {
  1122. as.t.IncrMsgRecv()
  1123. }
  1124. if as.desc.ServerStreams {
  1125. // Subsequent messages should be received by subsequent RecvMsg calls.
  1126. return nil
  1127. }
  1128. // Special handling for non-server-stream rpcs.
  1129. // This recv expects EOF or errors, so we don't collect inPayload.
  1130. err = recv(as.p, as.codec, as.s, as.dc, m, *as.callInfo.maxReceiveMessageSize, nil, as.decomp)
  1131. if err == nil {
  1132. return toRPCErr(errors.New("grpc: client streaming protocol violation: get <nil>, want <EOF>"))
  1133. }
  1134. if err == io.EOF {
  1135. return as.s.Status().Err() // non-server streaming Recv returns nil on success
  1136. }
  1137. return toRPCErr(err)
  1138. }
  1139. func (as *addrConnStream) finish(err error) {
  1140. as.mu.Lock()
  1141. if as.finished {
  1142. as.mu.Unlock()
  1143. return
  1144. }
  1145. as.finished = true
  1146. if err == io.EOF {
  1147. // Ending a stream with EOF indicates a success.
  1148. err = nil
  1149. }
  1150. if as.s != nil {
  1151. as.t.CloseStream(as.s, err)
  1152. }
  1153. if err != nil {
  1154. as.ac.incrCallsFailed()
  1155. } else {
  1156. as.ac.incrCallsSucceeded()
  1157. }
  1158. as.cancel()
  1159. as.mu.Unlock()
  1160. }
  1161. // ServerStream defines the server-side behavior of a streaming RPC.
  1162. //
  1163. // All errors returned from ServerStream methods are compatible with the
  1164. // status package.
  1165. type ServerStream interface {
  1166. // SetHeader sets the header metadata. It may be called multiple times.
  1167. // When call multiple times, all the provided metadata will be merged.
  1168. // All the metadata will be sent out when one of the following happens:
  1169. // - ServerStream.SendHeader() is called;
  1170. // - The first response is sent out;
  1171. // - An RPC status is sent out (error or success).
  1172. SetHeader(metadata.MD) error
  1173. // SendHeader sends the header metadata.
  1174. // The provided md and headers set by SetHeader() will be sent.
  1175. // It fails if called multiple times.
  1176. SendHeader(metadata.MD) error
  1177. // SetTrailer sets the trailer metadata which will be sent with the RPC status.
  1178. // When called more than once, all the provided metadata will be merged.
  1179. SetTrailer(metadata.MD)
  1180. // Context returns the context for this stream.
  1181. Context() context.Context
  1182. // SendMsg sends a message. On error, SendMsg aborts the stream and the
  1183. // error is returned directly.
  1184. //
  1185. // SendMsg blocks until:
  1186. // - There is sufficient flow control to schedule m with the transport, or
  1187. // - The stream is done, or
  1188. // - The stream breaks.
  1189. //
  1190. // SendMsg does not wait until the message is received by the client. An
  1191. // untimely stream closure may result in lost messages.
  1192. //
  1193. // It is safe to have a goroutine calling SendMsg and another goroutine
  1194. // calling RecvMsg on the same stream at the same time, but it is not safe
  1195. // to call SendMsg on the same stream in different goroutines.
  1196. SendMsg(m interface{}) error
  1197. // RecvMsg blocks until it receives a message into m or the stream is
  1198. // done. It returns io.EOF when the client has performed a CloseSend. On
  1199. // any non-EOF error, the stream is aborted and the error contains the
  1200. // RPC status.
  1201. //
  1202. // It is safe to have a goroutine calling SendMsg and another goroutine
  1203. // calling RecvMsg on the same stream at the same time, but it is not
  1204. // safe to call RecvMsg on the same stream in different goroutines.
  1205. RecvMsg(m interface{}) error
  1206. }
  1207. // serverStream implements a server side Stream.
  1208. type serverStream struct {
  1209. ctx context.Context
  1210. t transport.ServerTransport
  1211. s *transport.Stream
  1212. p *parser
  1213. codec baseCodec
  1214. cp Compressor
  1215. dc Decompressor
  1216. comp encoding.Compressor
  1217. decomp encoding.Compressor
  1218. maxReceiveMessageSize int
  1219. maxSendMessageSize int
  1220. trInfo *traceInfo
  1221. statsHandler stats.Handler
  1222. binlog *binarylog.MethodLogger
  1223. // serverHeaderBinlogged indicates whether server header has been logged. It
  1224. // will happen when one of the following two happens: stream.SendHeader(),
  1225. // stream.Send().
  1226. //
  1227. // It's only checked in send and sendHeader, doesn't need to be
  1228. // synchronized.
  1229. serverHeaderBinlogged bool
  1230. mu sync.Mutex // protects trInfo.tr after the service handler runs.
  1231. }
  1232. func (ss *serverStream) Context() context.Context {
  1233. return ss.ctx
  1234. }
  1235. func (ss *serverStream) SetHeader(md metadata.MD) error {
  1236. if md.Len() == 0 {
  1237. return nil
  1238. }
  1239. return ss.s.SetHeader(md)
  1240. }
  1241. func (ss *serverStream) SendHeader(md metadata.MD) error {
  1242. err := ss.t.WriteHeader(ss.s, md)
  1243. if ss.binlog != nil && !ss.serverHeaderBinlogged {
  1244. h, _ := ss.s.Header()
  1245. ss.binlog.Log(&binarylog.ServerHeader{
  1246. Header: h,
  1247. })
  1248. ss.serverHeaderBinlogged = true
  1249. }
  1250. return err
  1251. }
  1252. func (ss *serverStream) SetTrailer(md metadata.MD) {
  1253. if md.Len() == 0 {
  1254. return
  1255. }
  1256. ss.s.SetTrailer(md)
  1257. }
  1258. func (ss *serverStream) SendMsg(m interface{}) (err error) {
  1259. defer func() {
  1260. if ss.trInfo != nil {
  1261. ss.mu.Lock()
  1262. if ss.trInfo.tr != nil {
  1263. if err == nil {
  1264. ss.trInfo.tr.LazyLog(&payload{sent: true, msg: m}, true)
  1265. } else {
  1266. ss.trInfo.tr.LazyLog(&fmtStringer{"%v", []interface{}{err}}, true)
  1267. ss.trInfo.tr.SetError()
  1268. }
  1269. }
  1270. ss.mu.Unlock()
  1271. }
  1272. if err != nil && err != io.EOF {
  1273. st, _ := status.FromError(toRPCErr(err))
  1274. ss.t.WriteStatus(ss.s, st)
  1275. // Non-user specified status was sent out. This should be an error
  1276. // case (as a server side Cancel maybe).
  1277. //
  1278. // This is not handled specifically now. User will return a final
  1279. // status from the service handler, we will log that error instead.
  1280. // This behavior is similar to an interceptor.
  1281. }
  1282. if channelz.IsOn() && err == nil {
  1283. ss.t.IncrMsgSent()
  1284. }
  1285. }()
  1286. data, err := encode(ss.codec, m)
  1287. if err != nil {
  1288. return err
  1289. }
  1290. compData, err := compress(data, ss.cp, ss.comp)
  1291. if err != nil {
  1292. return err
  1293. }
  1294. hdr, payload := msgHeader(data, compData)
  1295. // TODO(dfawley): should we be checking len(data) instead?
  1296. if len(payload) > ss.maxSendMessageSize {
  1297. return status.Errorf(codes.ResourceExhausted, "trying to send message larger than max (%d vs. %d)", len(payload), ss.maxSendMessageSize)
  1298. }
  1299. if err := ss.t.Write(ss.s, hdr, payload, &transport.Options{Last: false}); err != nil {
  1300. return toRPCErr(err)
  1301. }
  1302. if ss.binlog != nil {
  1303. if !ss.serverHeaderBinlogged {
  1304. h, _ := ss.s.Header()
  1305. ss.binlog.Log(&binarylog.ServerHeader{
  1306. Header: h,
  1307. })
  1308. ss.serverHeaderBinlogged = true
  1309. }
  1310. ss.binlog.Log(&binarylog.ServerMessage{
  1311. Message: data,
  1312. })
  1313. }
  1314. if ss.statsHandler != nil {
  1315. ss.statsHandler.HandleRPC(ss.s.Context(), outPayload(false, m, data, payload, time.Now()))
  1316. }
  1317. return nil
  1318. }
  1319. func (ss *serverStream) RecvMsg(m interface{}) (err error) {
  1320. defer func() {
  1321. if ss.trInfo != nil {
  1322. ss.mu.Lock()
  1323. if ss.trInfo.tr != nil {
  1324. if err == nil {
  1325. ss.trInfo.tr.LazyLog(&payload{sent: false, msg: m}, true)
  1326. } else if err != io.EOF {
  1327. ss.trInfo.tr.LazyLog(&fmtStringer{"%v", []interface{}{err}}, true)
  1328. ss.trInfo.tr.SetError()
  1329. }
  1330. }
  1331. ss.mu.Unlock()
  1332. }
  1333. if err != nil && err != io.EOF {
  1334. st, _ := status.FromError(toRPCErr(err))
  1335. ss.t.WriteStatus(ss.s, st)
  1336. // Non-user specified status was sent out. This should be an error
  1337. // case (as a server side Cancel maybe).
  1338. //
  1339. // This is not handled specifically now. User will return a final
  1340. // status from the service handler, we will log that error instead.
  1341. // This behavior is similar to an interceptor.
  1342. }
  1343. if channelz.IsOn() && err == nil {
  1344. ss.t.IncrMsgRecv()
  1345. }
  1346. }()
  1347. var payInfo *payloadInfo
  1348. if ss.statsHandler != nil || ss.binlog != nil {
  1349. payInfo = &payloadInfo{}
  1350. }
  1351. if err := recv(ss.p, ss.codec, ss.s, ss.dc, m, ss.maxReceiveMessageSize, payInfo, ss.decomp); err != nil {
  1352. if err == io.EOF {
  1353. if ss.binlog != nil {
  1354. ss.binlog.Log(&binarylog.ClientHalfClose{})
  1355. }
  1356. return err
  1357. }
  1358. if err == io.ErrUnexpectedEOF {
  1359. err = status.Errorf(codes.Internal, io.ErrUnexpectedEOF.Error())
  1360. }
  1361. return toRPCErr(err)
  1362. }
  1363. if ss.statsHandler != nil {
  1364. ss.statsHandler.HandleRPC(ss.s.Context(), &stats.InPayload{
  1365. RecvTime: time.Now(),
  1366. Payload: m,
  1367. // TODO truncate large payload.
  1368. Data: payInfo.uncompressedBytes,
  1369. Length: len(payInfo.uncompressedBytes),
  1370. })
  1371. }
  1372. if ss.binlog != nil {
  1373. ss.binlog.Log(&binarylog.ClientMessage{
  1374. Message: payInfo.uncompressedBytes,
  1375. })
  1376. }
  1377. return nil
  1378. }
  1379. // MethodFromServerStream returns the method string for the input stream.
  1380. // The returned string is in the format of "/service/method".
  1381. func MethodFromServerStream(stream ServerStream) (string, bool) {
  1382. return Method(stream.Context())
  1383. }