diff --git a/transport/http2_client.go b/transport/http2_client.go index 8fc6ed0d..258695b2 100644 --- a/transport/http2_client.go +++ b/transport/http2_client.go @@ -233,8 +233,10 @@ func (t *http2Client) NewStream(ctx context.Context, callHdr *CallHdr) (_ *Strea if dl, ok := ctx.Deadline(); ok { timeout = dl.Sub(time.Now()) } - if err := ctx.Err(); err != nil { - return nil, ContextErr(err) + select { + case <-ctx.Done(): + return nil, ContextErr(ctx.Err()) + default: } pr := &peer.Peer{ Addr: t.conn.RemoteAddr(), @@ -516,13 +518,15 @@ func (t *http2Client) Write(s *Stream, data []byte, opts *Options) error { } return err } - if s.ctx.Err() != nil { + select { + case <-s.ctx.Done(): t.sendQuotaPool.add(len(p)) if t.framer.adjustNumWriters(-1) == 0 { t.controlBuf.put(&flushIO{}) } t.writableChan <- 0 return ContextErr(s.ctx.Err()) + default: } if r.Len() == 0 && t.framer.adjustNumWriters(0) == 1 { // Do a force flush iff this is last frame for the entire gRPC message diff --git a/transport/http2_server.go b/transport/http2_server.go index df918d8d..6f233d9d 100644 --- a/transport/http2_server.go +++ b/transport/http2_server.go @@ -600,13 +600,15 @@ func (t *http2Server) Write(s *Stream, data []byte, opts *Options) error { } return err } - if s.ctx.Err() != nil { + select { + case <-s.ctx.Done(): t.sendQuotaPool.add(ps) if t.framer.adjustNumWriters(-1) == 0 { t.controlBuf.put(&flushIO{}) } t.writableChan <- 0 return ContextErr(s.ctx.Err()) + default: } var forceFlush bool if r.Len() == 0 && t.framer.adjustNumWriters(0) == 1 && !opts.Last {