From b2f30e082e413423186459dd9677bd00f11c2ef1 Mon Sep 17 00:00:00 2001
From: iamqizhao <toqizhao@gmail.com>
Date: Thu, 28 Jul 2016 16:32:18 -0700
Subject: [PATCH] fix some goaway related test flakiness

---
 transport/http2_client.go | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/transport/http2_client.go b/transport/http2_client.go
index 62a6fa28..b6b744d5 100644
--- a/transport/http2_client.go
+++ b/transport/http2_client.go
@@ -308,6 +308,14 @@ func (t *http2Client) NewStream(ctx context.Context, callHdr *CallHdr) (_ *Strea
 		return nil, err
 	}
 	t.mu.Lock()
+	if t.state == draining {
+		t.mu.Unlock()
+		// t has been drained. Return the quota.
+		t.streamsQuota.add(1)
+		// Need to make t writable again so that the rpc in flight can still proceed.
+		t.writableChan <- 0
+		return nil, ErrStreamDrain
+	}
 	if t.state != reachable {
 		t.mu.Unlock()
 		return nil, ErrConnClosing