From b4b99fa9790c57d6ddfd0b1e92e2655a7a11b109 Mon Sep 17 00:00:00 2001
From: Natalie Villasana <navillasa@gmail.com>
Date: Wed, 20 Mar 2019 06:11:07 -0400
Subject: [PATCH] adds context timeout to transport.DialAddress (#1528)

---
 pkg/transport/transport.go | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/pkg/transport/transport.go b/pkg/transport/transport.go
index 8d36c0964..c1a2c5897 100644
--- a/pkg/transport/transport.go
+++ b/pkg/transport/transport.go
@@ -67,8 +67,8 @@ func (transport *Transport) DialNode(ctx context.Context, node *pb.Node, opts ..
 		grpc.FailOnNonTempDialError(true),
 	}, opts...)
 
-	timedCtx, cf := context.WithTimeout(ctx, connWaitTimeout)
-	defer cf()
+	timedCtx, cancel := context.WithTimeout(ctx, connWaitTimeout)
+	defer cancel()
 
 	conn, err = grpc.DialContext(timedCtx, node.GetAddress().Address, options...)
 	if err != nil {
@@ -98,7 +98,10 @@ func (transport *Transport) DialAddress(ctx context.Context, address string, opt
 		grpc.FailOnNonTempDialError(true),
 	}, opts...)
 
-	conn, err = grpc.DialContext(ctx, address, options...)
+	timedCtx, cancel := context.WithTimeout(ctx, connWaitTimeout)
+	defer cancel()
+
+	conn, err = grpc.DialContext(timedCtx, address, options...)
 	if err == context.Canceled {
 		return nil, err
 	}