From b38541aeb0a214ca654e964850800972b5d3b1a0 Mon Sep 17 00:00:00 2001
From: Spencer Kimball <spencer.kimball@gmail.com>
Date: Mon, 15 Aug 2016 13:17:09 -0400
Subject: [PATCH] Implement DialContext to afford caller option of managing
 cancelation

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

diff --git a/clientconn.go b/clientconn.go
index abe82cb97..8a7ec7842 100644
--- a/clientconn.go
+++ b/clientconn.go
@@ -214,9 +214,14 @@ func WithUserAgent(s string) DialOption {
 	}
 }
 
-// Dial creates a client connection the given target.
+// Dial creates a client connection to the given target.
 func Dial(target string, opts ...DialOption) (*ClientConn, error) {
-	ctx := context.Background()
+	return DialContext(context.Background(), target, opts...)
+}
+
+// DialContext creates a client connection to the given target
+// using the supplied context.
+func DialContext(ctx context.Context, target string, opts ...DialOption) (*ClientConn, error) {
 	cc := &ClientConn{
 		target: target,
 		conns:  make(map[Address]*addrConn),
diff --git a/clientconn_test.go b/clientconn_test.go
index 71ba45e65..59c0cc85e 100644
--- a/clientconn_test.go
+++ b/clientconn_test.go
@@ -37,6 +37,8 @@ import (
 	"testing"
 	"time"
 
+	"golang.org/x/net/context"
+
 	"google.golang.org/grpc/credentials"
 	"google.golang.org/grpc/credentials/oauth"
 )
@@ -67,6 +69,18 @@ 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)
+	}
+}
+
 func TestCredentialsMisuse(t *testing.T) {
 	tlsCreds, err := credentials.NewClientTLSFromFile(tlsDir+"ca.pem", "x.test.youtube.com")
 	if err != nil {