From d25e31e741ddfb45f4126cd20e357185751e42c2 Mon Sep 17 00:00:00 2001
From: Doug Fawley <dfawley@google.com>
Date: Fri, 10 Sep 2021 14:12:13 -0700
Subject: [PATCH] client: fix case where GOAWAY would leak connections and
 memory (#4755)

---
 clientconn.go        |  7 ++++++-
 test/end2end_test.go | 10 +++++-----
 2 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/clientconn.go b/clientconn.go
index 8012b3af..e933eae3 100644
--- a/clientconn.go
+++ b/clientconn.go
@@ -1287,9 +1287,14 @@ func (ac *addrConn) createTransport(addr resolver.Address, copts transport.Conne
 		ac.mu.Lock()
 		defer ac.mu.Unlock()
 		defer connClosed.Fire()
-		if !hcStarted {
+		if !hcStarted || hctx.Err() != nil {
 			// We didn't start the health check or set the state to READY, so
 			// no need to do anything else here.
+			//
+			// OR, we have already cancelled the health check context, meaning
+			// we have already called onClose once for this transport.  In this
+			// case it would be dangerous to clear the transport and update the
+			// state, since there may be a new transport in this addrConn.
 			return
 		}
 		hcancel()
diff --git a/test/end2end_test.go b/test/end2end_test.go
index 5702f18b..18e4ffc1 100644
--- a/test/end2end_test.go
+++ b/test/end2end_test.go
@@ -7074,11 +7074,11 @@ func (s) TestGoAwayThenClose(t *testing.T) {
 	s2 := grpc.NewServer()
 	defer s2.Stop()
 	testpb.RegisterTestServiceServer(s2, ts)
-	go s2.Serve(lis2)
 
 	r := manual.NewBuilderWithScheme("whatever")
 	r.InitialState(resolver.State{Addresses: []resolver.Address{
 		{Addr: lis1.Addr().String()},
+		{Addr: lis2.Addr().String()},
 	}})
 	cc, err := grpc.DialContext(ctx, r.Scheme()+":///", grpc.WithResolvers(r), grpc.WithInsecure())
 	if err != nil {
@@ -7102,10 +7102,7 @@ func (s) TestGoAwayThenClose(t *testing.T) {
 		t.Fatalf("unexpected error from first recv: %v", err)
 	}
 
-	r.UpdateState(resolver.State{Addresses: []resolver.Address{
-		{Addr: lis1.Addr().String()},
-		{Addr: lis2.Addr().String()},
-	}})
+	go s2.Serve(lis2)
 
 	// Send GO_AWAY to connection 1.
 	go s1.GracefulStop()
@@ -7126,6 +7123,9 @@ func (s) TestGoAwayThenClose(t *testing.T) {
 	// Assert that connection 2 has been established.
 	<-conn2Established.Done()
 
+	// Close the listener for server2 to prevent it from allowing new connections.
+	lis2.Close()
+
 	// Close connection 1.
 	s1.Stop()