stream.go 22 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780
  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. "errors"
  21. "io"
  22. "sync"
  23. "time"
  24. "golang.org/x/net/context"
  25. "golang.org/x/net/trace"
  26. "google.golang.org/grpc/balancer"
  27. "google.golang.org/grpc/codes"
  28. "google.golang.org/grpc/encoding"
  29. "google.golang.org/grpc/internal/channelz"
  30. "google.golang.org/grpc/metadata"
  31. "google.golang.org/grpc/stats"
  32. "google.golang.org/grpc/status"
  33. "google.golang.org/grpc/transport"
  34. )
  35. // StreamHandler defines the handler called by gRPC server to complete the
  36. // execution of a streaming RPC. If a StreamHandler returns an error, it
  37. // should be produced by the status package, or else gRPC will use
  38. // codes.Unknown as the status code and err.Error() as the status message
  39. // of the RPC.
  40. type StreamHandler func(srv interface{}, stream ServerStream) error
  41. // StreamDesc represents a streaming RPC service's method specification.
  42. type StreamDesc struct {
  43. StreamName string
  44. Handler StreamHandler
  45. // At least one of these is true.
  46. ServerStreams bool
  47. ClientStreams bool
  48. }
  49. // Stream defines the common interface a client or server stream has to satisfy.
  50. //
  51. // All errors returned from Stream are compatible with the status package.
  52. type Stream interface {
  53. // Context returns the context for this stream.
  54. Context() context.Context
  55. // SendMsg blocks until it sends m, the stream is done or the stream
  56. // breaks.
  57. // On error, it aborts the stream and returns an RPC status on client
  58. // side. On server side, it simply returns the error to the caller.
  59. // SendMsg is called by generated code. Also Users can call SendMsg
  60. // directly when it is really needed in their use cases.
  61. // It's safe to have a goroutine calling SendMsg and another goroutine calling
  62. // recvMsg on the same stream at the same time.
  63. // But it is not safe to call SendMsg on the same stream in different goroutines.
  64. SendMsg(m interface{}) error
  65. // RecvMsg blocks until it receives a message or the stream is
  66. // done. On client side, it returns io.EOF when the stream is done. On
  67. // any other error, it aborts the stream and returns an RPC status. On
  68. // server side, it simply returns the error to the caller.
  69. // It's safe to have a goroutine calling SendMsg and another goroutine calling
  70. // recvMsg on the same stream at the same time.
  71. // But it is not safe to call RecvMsg on the same stream in different goroutines.
  72. RecvMsg(m interface{}) error
  73. }
  74. // ClientStream defines the interface a client stream has to satisfy.
  75. type ClientStream interface {
  76. // Header returns the header metadata received from the server if there
  77. // is any. It blocks if the metadata is not ready to read.
  78. Header() (metadata.MD, error)
  79. // Trailer returns the trailer metadata from the server, if there is any.
  80. // It must only be called after stream.CloseAndRecv has returned, or
  81. // stream.Recv has returned a non-nil error (including io.EOF).
  82. Trailer() metadata.MD
  83. // CloseSend closes the send direction of the stream. It closes the stream
  84. // when non-nil error is met.
  85. CloseSend() error
  86. // Stream.SendMsg() may return a non-nil error when something wrong happens sending
  87. // the request. The returned error indicates the status of this sending, not the final
  88. // status of the RPC.
  89. //
  90. // Always call Stream.RecvMsg() to drain the stream and get the final
  91. // status, otherwise there could be leaked resources.
  92. Stream
  93. }
  94. // NewStream creates a new Stream for the client side. This is typically
  95. // called by generated code. ctx is used for the lifetime of the stream.
  96. //
  97. // To ensure resources are not leaked due to the stream returned, one of the following
  98. // actions must be performed:
  99. //
  100. // 1. Call Close on the ClientConn.
  101. // 2. Cancel the context provided.
  102. // 3. Call RecvMsg until a non-nil error is returned. A protobuf-generated
  103. // client-streaming RPC, for instance, might use the helper function
  104. // CloseAndRecv (note that CloseSend does not Recv, therefore is not
  105. // guaranteed to release all resources).
  106. // 4. Receive a non-nil, non-io.EOF error from Header or SendMsg.
  107. //
  108. // If none of the above happen, a goroutine and a context will be leaked, and grpc
  109. // will not call the optionally-configured stats handler with a stats.End message.
  110. func (cc *ClientConn) NewStream(ctx context.Context, desc *StreamDesc, method string, opts ...CallOption) (ClientStream, error) {
  111. // allow interceptor to see all applicable call options, which means those
  112. // configured as defaults from dial option as well as per-call options
  113. opts = combine(cc.dopts.callOptions, opts)
  114. if cc.dopts.streamInt != nil {
  115. return cc.dopts.streamInt(ctx, desc, cc, method, newClientStream, opts...)
  116. }
  117. return newClientStream(ctx, desc, cc, method, opts...)
  118. }
  119. // NewClientStream is a wrapper for ClientConn.NewStream.
  120. //
  121. // DEPRECATED: Use ClientConn.NewStream instead.
  122. func NewClientStream(ctx context.Context, desc *StreamDesc, cc *ClientConn, method string, opts ...CallOption) (ClientStream, error) {
  123. return cc.NewStream(ctx, desc, method, opts...)
  124. }
  125. func newClientStream(ctx context.Context, desc *StreamDesc, cc *ClientConn, method string, opts ...CallOption) (_ ClientStream, err error) {
  126. if channelz.IsOn() {
  127. cc.incrCallsStarted()
  128. defer func() {
  129. if err != nil {
  130. cc.incrCallsFailed()
  131. }
  132. }()
  133. }
  134. c := defaultCallInfo()
  135. mc := cc.GetMethodConfig(method)
  136. if mc.WaitForReady != nil {
  137. c.failFast = !*mc.WaitForReady
  138. }
  139. // Possible context leak:
  140. // The cancel function for the child context we create will only be called
  141. // when RecvMsg returns a non-nil error, if the ClientConn is closed, or if
  142. // an error is generated by SendMsg.
  143. // https://github.com/grpc/grpc-go/issues/1818.
  144. var cancel context.CancelFunc
  145. if mc.Timeout != nil && *mc.Timeout >= 0 {
  146. ctx, cancel = context.WithTimeout(ctx, *mc.Timeout)
  147. } else {
  148. ctx, cancel = context.WithCancel(ctx)
  149. }
  150. defer func() {
  151. if err != nil {
  152. cancel()
  153. }
  154. }()
  155. for _, o := range opts {
  156. if err := o.before(c); err != nil {
  157. return nil, toRPCErr(err)
  158. }
  159. }
  160. c.maxSendMessageSize = getMaxSize(mc.MaxReqSize, c.maxSendMessageSize, defaultClientMaxSendMessageSize)
  161. c.maxReceiveMessageSize = getMaxSize(mc.MaxRespSize, c.maxReceiveMessageSize, defaultClientMaxReceiveMessageSize)
  162. if err := setCallInfoCodec(c); err != nil {
  163. return nil, err
  164. }
  165. callHdr := &transport.CallHdr{
  166. Host: cc.authority,
  167. Method: method,
  168. // If it's not client streaming, we should already have the request to be sent,
  169. // so we don't flush the header.
  170. // If it's client streaming, the user may never send a request or send it any
  171. // time soon, so we ask the transport to flush the header.
  172. Flush: desc.ClientStreams,
  173. ContentSubtype: c.contentSubtype,
  174. }
  175. // Set our outgoing compression according to the UseCompressor CallOption, if
  176. // set. In that case, also find the compressor from the encoding package.
  177. // Otherwise, use the compressor configured by the WithCompressor DialOption,
  178. // if set.
  179. var cp Compressor
  180. var comp encoding.Compressor
  181. if ct := c.compressorType; ct != "" {
  182. callHdr.SendCompress = ct
  183. if ct != encoding.Identity {
  184. comp = encoding.GetCompressor(ct)
  185. if comp == nil {
  186. return nil, status.Errorf(codes.Internal, "grpc: Compressor is not installed for requested grpc-encoding %q", ct)
  187. }
  188. }
  189. } else if cc.dopts.cp != nil {
  190. callHdr.SendCompress = cc.dopts.cp.Type()
  191. cp = cc.dopts.cp
  192. }
  193. if c.creds != nil {
  194. callHdr.Creds = c.creds
  195. }
  196. var trInfo traceInfo
  197. if EnableTracing {
  198. trInfo.tr = trace.New("grpc.Sent."+methodFamily(method), method)
  199. trInfo.firstLine.client = true
  200. if deadline, ok := ctx.Deadline(); ok {
  201. trInfo.firstLine.deadline = deadline.Sub(time.Now())
  202. }
  203. trInfo.tr.LazyLog(&trInfo.firstLine, false)
  204. ctx = trace.NewContext(ctx, trInfo.tr)
  205. defer func() {
  206. if err != nil {
  207. // Need to call tr.finish() if error is returned.
  208. // Because tr will not be returned to caller.
  209. trInfo.tr.LazyPrintf("RPC: [%v]", err)
  210. trInfo.tr.SetError()
  211. trInfo.tr.Finish()
  212. }
  213. }()
  214. }
  215. ctx = newContextWithRPCInfo(ctx, c.failFast)
  216. sh := cc.dopts.copts.StatsHandler
  217. var beginTime time.Time
  218. if sh != nil {
  219. ctx = sh.TagRPC(ctx, &stats.RPCTagInfo{FullMethodName: method, FailFast: c.failFast})
  220. beginTime = time.Now()
  221. begin := &stats.Begin{
  222. Client: true,
  223. BeginTime: beginTime,
  224. FailFast: c.failFast,
  225. }
  226. sh.HandleRPC(ctx, begin)
  227. defer func() {
  228. if err != nil {
  229. // Only handle end stats if err != nil.
  230. end := &stats.End{
  231. Client: true,
  232. Error: err,
  233. BeginTime: beginTime,
  234. EndTime: time.Now(),
  235. }
  236. sh.HandleRPC(ctx, end)
  237. }
  238. }()
  239. }
  240. var (
  241. t transport.ClientTransport
  242. s *transport.Stream
  243. done func(balancer.DoneInfo)
  244. )
  245. for {
  246. // Check to make sure the context has expired. This will prevent us from
  247. // looping forever if an error occurs for wait-for-ready RPCs where no data
  248. // is sent on the wire.
  249. select {
  250. case <-ctx.Done():
  251. return nil, toRPCErr(ctx.Err())
  252. default:
  253. }
  254. t, done, err = cc.getTransport(ctx, c.failFast)
  255. if err != nil {
  256. return nil, err
  257. }
  258. s, err = t.NewStream(ctx, callHdr)
  259. if err != nil {
  260. if done != nil {
  261. done(balancer.DoneInfo{Err: err})
  262. done = nil
  263. }
  264. // In the event of any error from NewStream, we never attempted to write
  265. // anything to the wire, so we can retry indefinitely for non-fail-fast
  266. // RPCs.
  267. if !c.failFast {
  268. continue
  269. }
  270. return nil, toRPCErr(err)
  271. }
  272. break
  273. }
  274. cs := &clientStream{
  275. opts: opts,
  276. c: c,
  277. cc: cc,
  278. desc: desc,
  279. codec: c.codec,
  280. cp: cp,
  281. comp: comp,
  282. cancel: cancel,
  283. attempt: &csAttempt{
  284. t: t,
  285. s: s,
  286. p: &parser{r: s},
  287. done: done,
  288. dc: cc.dopts.dc,
  289. ctx: ctx,
  290. trInfo: trInfo,
  291. statsHandler: sh,
  292. beginTime: beginTime,
  293. },
  294. }
  295. cs.c.stream = cs
  296. cs.attempt.cs = cs
  297. if desc != unaryStreamDesc {
  298. // Listen on cc and stream contexts to cleanup when the user closes the
  299. // ClientConn or cancels the stream context. In all other cases, an error
  300. // should already be injected into the recv buffer by the transport, which
  301. // the client will eventually receive, and then we will cancel the stream's
  302. // context in clientStream.finish.
  303. go func() {
  304. select {
  305. case <-cc.ctx.Done():
  306. cs.finish(ErrClientConnClosing)
  307. case <-ctx.Done():
  308. cs.finish(toRPCErr(ctx.Err()))
  309. }
  310. }()
  311. }
  312. return cs, nil
  313. }
  314. // clientStream implements a client side Stream.
  315. type clientStream struct {
  316. opts []CallOption
  317. c *callInfo
  318. cc *ClientConn
  319. desc *StreamDesc
  320. codec baseCodec
  321. cp Compressor
  322. comp encoding.Compressor
  323. cancel context.CancelFunc // cancels all attempts
  324. sentLast bool // sent an end stream
  325. mu sync.Mutex // guards finished
  326. finished bool // TODO: replace with atomic cmpxchg or sync.Once?
  327. attempt *csAttempt // the active client stream attempt
  328. // TODO(hedging): hedging will have multiple attempts simultaneously.
  329. }
  330. // csAttempt implements a single transport stream attempt within a
  331. // clientStream.
  332. type csAttempt struct {
  333. cs *clientStream
  334. t transport.ClientTransport
  335. s *transport.Stream
  336. p *parser
  337. done func(balancer.DoneInfo)
  338. dc Decompressor
  339. decomp encoding.Compressor
  340. decompSet bool
  341. ctx context.Context // the application's context, wrapped by stats/tracing
  342. mu sync.Mutex // guards trInfo.tr
  343. // trInfo.tr is set when created (if EnableTracing is true),
  344. // and cleared when the finish method is called.
  345. trInfo traceInfo
  346. statsHandler stats.Handler
  347. beginTime time.Time
  348. }
  349. func (cs *clientStream) Context() context.Context {
  350. // TODO(retry): commit the current attempt (the context has peer-aware data).
  351. return cs.attempt.context()
  352. }
  353. func (cs *clientStream) Header() (metadata.MD, error) {
  354. m, err := cs.attempt.header()
  355. if err != nil {
  356. // TODO(retry): maybe retry on error or commit attempt on success.
  357. err = toRPCErr(err)
  358. cs.finish(err)
  359. }
  360. return m, err
  361. }
  362. func (cs *clientStream) Trailer() metadata.MD {
  363. // TODO(retry): on error, maybe retry (trailers-only).
  364. return cs.attempt.trailer()
  365. }
  366. func (cs *clientStream) SendMsg(m interface{}) (err error) {
  367. // TODO(retry): buffer message for replaying if not committed.
  368. return cs.attempt.sendMsg(m)
  369. }
  370. func (cs *clientStream) RecvMsg(m interface{}) (err error) {
  371. // TODO(retry): maybe retry on error or commit attempt on success.
  372. return cs.attempt.recvMsg(m)
  373. }
  374. func (cs *clientStream) CloseSend() error {
  375. cs.attempt.closeSend()
  376. return nil
  377. }
  378. func (cs *clientStream) finish(err error) {
  379. if err == io.EOF {
  380. // Ending a stream with EOF indicates a success.
  381. err = nil
  382. }
  383. cs.mu.Lock()
  384. if cs.finished {
  385. cs.mu.Unlock()
  386. return
  387. }
  388. cs.finished = true
  389. cs.mu.Unlock()
  390. if channelz.IsOn() {
  391. if err != nil {
  392. cs.cc.incrCallsFailed()
  393. } else {
  394. cs.cc.incrCallsSucceeded()
  395. }
  396. }
  397. // TODO(retry): commit current attempt if necessary.
  398. cs.attempt.finish(err)
  399. for _, o := range cs.opts {
  400. o.after(cs.c)
  401. }
  402. cs.cancel()
  403. }
  404. func (a *csAttempt) context() context.Context {
  405. return a.s.Context()
  406. }
  407. func (a *csAttempt) header() (metadata.MD, error) {
  408. return a.s.Header()
  409. }
  410. func (a *csAttempt) trailer() metadata.MD {
  411. return a.s.Trailer()
  412. }
  413. func (a *csAttempt) sendMsg(m interface{}) (err error) {
  414. // TODO Investigate how to signal the stats handling party.
  415. // generate error stats if err != nil && err != io.EOF?
  416. cs := a.cs
  417. defer func() {
  418. // For non-client-streaming RPCs, we return nil instead of EOF on success
  419. // because the generated code requires it. finish is not called; RecvMsg()
  420. // will call it with the stream's status independently.
  421. if err == io.EOF && !cs.desc.ClientStreams {
  422. err = nil
  423. }
  424. if err != nil && err != io.EOF {
  425. // Call finish on the client stream for errors generated by this SendMsg
  426. // call, as these indicate problems created by this client. (Transport
  427. // errors are converted to an io.EOF error below; the real error will be
  428. // returned from RecvMsg eventually in that case, or be retried.)
  429. cs.finish(err)
  430. }
  431. }()
  432. // TODO: Check cs.sentLast and error if we already ended the stream.
  433. if EnableTracing {
  434. a.mu.Lock()
  435. if a.trInfo.tr != nil {
  436. a.trInfo.tr.LazyLog(&payload{sent: true, msg: m}, true)
  437. }
  438. a.mu.Unlock()
  439. }
  440. data, err := encode(cs.codec, m)
  441. if err != nil {
  442. return err
  443. }
  444. compData, err := compress(data, cs.cp, cs.comp)
  445. if err != nil {
  446. return err
  447. }
  448. hdr, payload := msgHeader(data, compData)
  449. // TODO(dfawley): should we be checking len(data) instead?
  450. if len(payload) > *cs.c.maxSendMessageSize {
  451. return status.Errorf(codes.ResourceExhausted, "trying to send message larger than max (%d vs. %d)", len(payload), *cs.c.maxSendMessageSize)
  452. }
  453. if !cs.desc.ClientStreams {
  454. cs.sentLast = true
  455. }
  456. err = a.t.Write(a.s, hdr, payload, &transport.Options{Last: !cs.desc.ClientStreams})
  457. if err == nil {
  458. if a.statsHandler != nil {
  459. a.statsHandler.HandleRPC(a.ctx, outPayload(true, m, data, payload, time.Now()))
  460. }
  461. if channelz.IsOn() {
  462. a.t.IncrMsgSent()
  463. }
  464. return nil
  465. }
  466. return io.EOF
  467. }
  468. func (a *csAttempt) recvMsg(m interface{}) (err error) {
  469. cs := a.cs
  470. defer func() {
  471. if err != nil || !cs.desc.ServerStreams {
  472. // err != nil or non-server-streaming indicates end of stream.
  473. cs.finish(err)
  474. }
  475. }()
  476. var inPayload *stats.InPayload
  477. if a.statsHandler != nil {
  478. inPayload = &stats.InPayload{
  479. Client: true,
  480. }
  481. }
  482. if !a.decompSet {
  483. // Block until we receive headers containing received message encoding.
  484. if ct := a.s.RecvCompress(); ct != "" && ct != encoding.Identity {
  485. if a.dc == nil || a.dc.Type() != ct {
  486. // No configured decompressor, or it does not match the incoming
  487. // message encoding; attempt to find a registered compressor that does.
  488. a.dc = nil
  489. a.decomp = encoding.GetCompressor(ct)
  490. }
  491. } else {
  492. // No compression is used; disable our decompressor.
  493. a.dc = nil
  494. }
  495. // Only initialize this state once per stream.
  496. a.decompSet = true
  497. }
  498. err = recv(a.p, cs.codec, a.s, a.dc, m, *cs.c.maxReceiveMessageSize, inPayload, a.decomp)
  499. if err != nil {
  500. if err == io.EOF {
  501. if statusErr := a.s.Status().Err(); statusErr != nil {
  502. return statusErr
  503. }
  504. return io.EOF // indicates successful end of stream.
  505. }
  506. return toRPCErr(err)
  507. }
  508. if EnableTracing {
  509. a.mu.Lock()
  510. if a.trInfo.tr != nil {
  511. a.trInfo.tr.LazyLog(&payload{sent: false, msg: m}, true)
  512. }
  513. a.mu.Unlock()
  514. }
  515. if inPayload != nil {
  516. a.statsHandler.HandleRPC(a.ctx, inPayload)
  517. }
  518. if channelz.IsOn() {
  519. a.t.IncrMsgRecv()
  520. }
  521. if cs.desc.ServerStreams {
  522. // Subsequent messages should be received by subsequent RecvMsg calls.
  523. return nil
  524. }
  525. // Special handling for non-server-stream rpcs.
  526. // This recv expects EOF or errors, so we don't collect inPayload.
  527. err = recv(a.p, cs.codec, a.s, a.dc, m, *cs.c.maxReceiveMessageSize, nil, a.decomp)
  528. if err == nil {
  529. return toRPCErr(errors.New("grpc: client streaming protocol violation: get <nil>, want <EOF>"))
  530. }
  531. if err == io.EOF {
  532. return a.s.Status().Err() // non-server streaming Recv returns nil on success
  533. }
  534. return toRPCErr(err)
  535. }
  536. func (a *csAttempt) closeSend() {
  537. cs := a.cs
  538. if cs.sentLast {
  539. return
  540. }
  541. cs.sentLast = true
  542. cs.attempt.t.Write(cs.attempt.s, nil, nil, &transport.Options{Last: true})
  543. // We ignore errors from Write. Any error it would return would also be
  544. // returned by a subsequent RecvMsg call, and the user is supposed to always
  545. // finish the stream by calling RecvMsg until it returns err != nil.
  546. }
  547. func (a *csAttempt) finish(err error) {
  548. a.mu.Lock()
  549. a.t.CloseStream(a.s, err)
  550. if a.done != nil {
  551. a.done(balancer.DoneInfo{
  552. Err: err,
  553. BytesSent: true,
  554. BytesReceived: a.s.BytesReceived(),
  555. })
  556. }
  557. if a.statsHandler != nil {
  558. end := &stats.End{
  559. Client: true,
  560. BeginTime: a.beginTime,
  561. EndTime: time.Now(),
  562. Error: err,
  563. }
  564. a.statsHandler.HandleRPC(a.ctx, end)
  565. }
  566. if a.trInfo.tr != nil {
  567. if err == nil {
  568. a.trInfo.tr.LazyPrintf("RPC: [OK]")
  569. } else {
  570. a.trInfo.tr.LazyPrintf("RPC: [%v]", err)
  571. a.trInfo.tr.SetError()
  572. }
  573. a.trInfo.tr.Finish()
  574. a.trInfo.tr = nil
  575. }
  576. a.mu.Unlock()
  577. }
  578. // ServerStream defines the interface a server stream has to satisfy.
  579. type ServerStream interface {
  580. // SetHeader sets the header metadata. It may be called multiple times.
  581. // When call multiple times, all the provided metadata will be merged.
  582. // All the metadata will be sent out when one of the following happens:
  583. // - ServerStream.SendHeader() is called;
  584. // - The first response is sent out;
  585. // - An RPC status is sent out (error or success).
  586. SetHeader(metadata.MD) error
  587. // SendHeader sends the header metadata.
  588. // The provided md and headers set by SetHeader() will be sent.
  589. // It fails if called multiple times.
  590. SendHeader(metadata.MD) error
  591. // SetTrailer sets the trailer metadata which will be sent with the RPC status.
  592. // When called more than once, all the provided metadata will be merged.
  593. SetTrailer(metadata.MD)
  594. Stream
  595. }
  596. // serverStream implements a server side Stream.
  597. type serverStream struct {
  598. ctx context.Context
  599. t transport.ServerTransport
  600. s *transport.Stream
  601. p *parser
  602. codec baseCodec
  603. cp Compressor
  604. dc Decompressor
  605. comp encoding.Compressor
  606. decomp encoding.Compressor
  607. maxReceiveMessageSize int
  608. maxSendMessageSize int
  609. trInfo *traceInfo
  610. statsHandler stats.Handler
  611. mu sync.Mutex // protects trInfo.tr after the service handler runs.
  612. }
  613. func (ss *serverStream) Context() context.Context {
  614. return ss.ctx
  615. }
  616. func (ss *serverStream) SetHeader(md metadata.MD) error {
  617. if md.Len() == 0 {
  618. return nil
  619. }
  620. return ss.s.SetHeader(md)
  621. }
  622. func (ss *serverStream) SendHeader(md metadata.MD) error {
  623. return ss.t.WriteHeader(ss.s, md)
  624. }
  625. func (ss *serverStream) SetTrailer(md metadata.MD) {
  626. if md.Len() == 0 {
  627. return
  628. }
  629. ss.s.SetTrailer(md)
  630. }
  631. func (ss *serverStream) SendMsg(m interface{}) (err error) {
  632. defer func() {
  633. if ss.trInfo != nil {
  634. ss.mu.Lock()
  635. if ss.trInfo.tr != nil {
  636. if err == nil {
  637. ss.trInfo.tr.LazyLog(&payload{sent: true, msg: m}, true)
  638. } else {
  639. ss.trInfo.tr.LazyLog(&fmtStringer{"%v", []interface{}{err}}, true)
  640. ss.trInfo.tr.SetError()
  641. }
  642. }
  643. ss.mu.Unlock()
  644. }
  645. if err != nil && err != io.EOF {
  646. st, _ := status.FromError(toRPCErr(err))
  647. ss.t.WriteStatus(ss.s, st)
  648. }
  649. if channelz.IsOn() && err == nil {
  650. ss.t.IncrMsgSent()
  651. }
  652. }()
  653. data, err := encode(ss.codec, m)
  654. if err != nil {
  655. return err
  656. }
  657. compData, err := compress(data, ss.cp, ss.comp)
  658. if err != nil {
  659. return err
  660. }
  661. hdr, payload := msgHeader(data, compData)
  662. // TODO(dfawley): should we be checking len(data) instead?
  663. if len(payload) > ss.maxSendMessageSize {
  664. return status.Errorf(codes.ResourceExhausted, "trying to send message larger than max (%d vs. %d)", len(payload), ss.maxSendMessageSize)
  665. }
  666. if err := ss.t.Write(ss.s, hdr, payload, &transport.Options{Last: false}); err != nil {
  667. return toRPCErr(err)
  668. }
  669. if ss.statsHandler != nil {
  670. ss.statsHandler.HandleRPC(ss.s.Context(), outPayload(false, m, data, payload, time.Now()))
  671. }
  672. return nil
  673. }
  674. func (ss *serverStream) RecvMsg(m interface{}) (err error) {
  675. defer func() {
  676. if ss.trInfo != nil {
  677. ss.mu.Lock()
  678. if ss.trInfo.tr != nil {
  679. if err == nil {
  680. ss.trInfo.tr.LazyLog(&payload{sent: false, msg: m}, true)
  681. } else if err != io.EOF {
  682. ss.trInfo.tr.LazyLog(&fmtStringer{"%v", []interface{}{err}}, true)
  683. ss.trInfo.tr.SetError()
  684. }
  685. }
  686. ss.mu.Unlock()
  687. }
  688. if err != nil && err != io.EOF {
  689. st, _ := status.FromError(toRPCErr(err))
  690. ss.t.WriteStatus(ss.s, st)
  691. }
  692. if channelz.IsOn() && err == nil {
  693. ss.t.IncrMsgRecv()
  694. }
  695. }()
  696. var inPayload *stats.InPayload
  697. if ss.statsHandler != nil {
  698. inPayload = &stats.InPayload{}
  699. }
  700. if err := recv(ss.p, ss.codec, ss.s, ss.dc, m, ss.maxReceiveMessageSize, inPayload, ss.decomp); err != nil {
  701. if err == io.EOF {
  702. return err
  703. }
  704. if err == io.ErrUnexpectedEOF {
  705. err = status.Errorf(codes.Internal, io.ErrUnexpectedEOF.Error())
  706. }
  707. return toRPCErr(err)
  708. }
  709. if inPayload != nil {
  710. ss.statsHandler.HandleRPC(ss.s.Context(), inPayload)
  711. }
  712. return nil
  713. }
  714. // MethodFromServerStream returns the method string for the input stream.
  715. // The returned string is in the format of "/service/method".
  716. func MethodFromServerStream(stream ServerStream) (string, bool) {
  717. return Method(stream.Context())
  718. }