From 7eae19acb7801d043976b181b1ae4047716a19c2 Mon Sep 17 00:00:00 2001
From: iamqizhao <toqizhao@gmail.com>
Date: Tue, 23 Aug 2016 19:23:04 -0700
Subject: [PATCH] fix the issue

---
 clientconn.go      | 14 ++++++++++++--
 clientconn_test.go | 10 +++-------
 2 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/clientconn.go b/clientconn.go
index 215829ff..1c37cfc4 100644
--- a/clientconn.go
+++ b/clientconn.go
@@ -221,19 +221,29 @@ func Dial(target string, opts ...DialOption) (*ClientConn, error) {
 }
 
 // DialContext creates a client connection to the given target. ctx can be used to
-// cancel or expire the pending connecting. Once the initial connection is done,
-// the cancellation and expiration of ctx will not affect the connection.
+// cancel or expire the pending connecting. Once this function returns, the
+// cancellation and expiration of ctx will be noop. Users should call ClientConn.Close
+// to terminate all the pending operations after this function returns.
 func DialContext(ctx context.Context, target string, opts ...DialOption) (*ClientConn, error) {
 	cc := &ClientConn{
 		target: target,
 		conns:  make(map[Address]*addrConn),
 	}
 	cc.ctx, cc.cancel = context.WithCancel(context.Background())
+	done := make(chan struct{})
+	defer close(done)
 	go func() {
 		select {
 		case <-ctx.Done():
 			cc.Close()
+		case <-done:
+			select {
+			case <-ctx.Done():
+				cc.Close()
+			default:
+			}
 		case <-cc.ctx.Done():
+			// ClientConn.Close has been called.
 		}
 	}()
 
diff --git a/clientconn_test.go b/clientconn_test.go
index 59c0cc85..d6d7c6be 100644
--- a/clientconn_test.go
+++ b/clientconn_test.go
@@ -71,13 +71,9 @@ func TestTLSDialTimeout(t *testing.T) {
 
 func TestDialContextCancel(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
-	go cancel()
-	conn, err := DialContext(ctx, "Non-Existent.Server:80", WithBlock(), WithInsecure())
-	if err == nil {
-		conn.Close()
-	}
-	if err != context.Canceled {
-		t.Fatalf("DialContext(_, _) = %v, %v, want %v", conn, err, context.Canceled)
+	cancel()
+	if _, err := DialContext(ctx, "Non-Existent.Server:80", WithBlock(), WithInsecure()); err != context.Canceled && err != ErrClientConnClosing {
+		t.Fatalf("grpc.DialContext(%v, _) = _, %v, want _, %v or %v", ctx, err, context.Canceled, ErrClientConnClosing)
 	}
 }