From 9602024760c025b5b4f74e31fa0d3e90fd320572 Mon Sep 17 00:00:00 2001
From: Menghan Li <menghanl@google.com>
Date: Mon, 17 May 2021 15:15:44 -0700
Subject: [PATCH] [better_conn_close_error] c1

---
 internal/transport/http2_client.go   | 2 +-
 internal/transport/transport_test.go | 6 ++++++
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/internal/transport/http2_client.go b/internal/transport/http2_client.go
index 8999a65e5cf..452bfded0c5 100644
--- a/internal/transport/http2_client.go
+++ b/internal/transport/http2_client.go
@@ -59,7 +59,6 @@ type http2Client struct {
 	ctx        context.Context
 	cancel     context.CancelFunc
 	ctxDone    <-chan struct{} // Cache the ctx.Done() chan.
-	closeErr   error           // Reason the transport is closed. Only set if ctx is canceled.
 	userAgent  string
 	md         metadata.MD
 	conn       net.Conn // underlying communication channel
@@ -110,6 +109,7 @@ type http2Client struct {
 
 	mu            sync.Mutex // guard the following variables
 	state         transportState
+	closeErr      error // Reason the transport is closed. Only set if ctx is canceled.
 	activeStreams map[uint32]*Stream
 	// prevGoAway ID records the Last-Stream-ID in the previous GOAway frame.
 	prevGoAwayID uint32
diff --git a/internal/transport/transport_test.go b/internal/transport/transport_test.go
index 6fc6cba1a3e..ffe12f3d46d 100644
--- a/internal/transport/transport_test.go
+++ b/internal/transport/transport_test.go
@@ -781,6 +781,12 @@ func (s) TestGracefulClose(t *testing.T) {
 			defer wg.Done()
 			str, err := ct.NewStream(ctx, &CallHdr{})
 			if err != nil {
+				// The NewStream calls race with closing of the transport (after
+				// the last stream is closed).
+				// - if it is after the transport is closed (case <-ct.ctxDone),
+				//   we don't care about the error.
+				// - if it's before the transport is closed (default), we want
+				//   the NewStream call to succeed.
 				select {
 				case <-ct.ctxDone:
 				default: