From b93b4732ebf671e94d6a9e0e208c656709a2fc3a Mon Sep 17 00:00:00 2001
From: Erik Wilson <Erik.E.Wilson@gmail.com>
Date: Wed, 17 Jul 2019 17:13:40 -0700
Subject: [PATCH 1/2] Start endpoint tunnel watch before waiting

---
 pkg/agent/tunnel/tunnel.go | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/pkg/agent/tunnel/tunnel.go b/pkg/agent/tunnel/tunnel.go
index badc519ab9..ab02f2d513 100644
--- a/pkg/agent/tunnel/tunnel.go
+++ b/pkg/agent/tunnel/tunnel.go
@@ -83,7 +83,6 @@ func Setup(config *config.Node) error {
 			disconnect[address] = connect(wg, address, config, transportConfig)
 		}
 	}
-	wg.Wait()
 
 	go func() {
 	connect:
@@ -134,6 +133,8 @@ func Setup(config *config.Node) error {
 		}
 	}()
 
+	wg.Wait()
+
 	return nil
 }
 
@@ -178,6 +179,9 @@ func connect(waitGroup *sync.WaitGroup, address string, config *config.Node, tra
 			})
 
 			if ctx.Err() != nil {
+				if waitGroup != nil {
+					once.Do(waitGroup.Done)
+				}
 				logrus.Infof("Stopped tunnel to %s", wsURL)
 				return
 			}

From 23b079757846b028e42ba586ab24f97cc7c01365 Mon Sep 17 00:00:00 2001
From: Erik Wilson <Erik.E.Wilson@gmail.com>
Date: Wed, 17 Jul 2019 18:15:15 -0700
Subject: [PATCH 2/2] Add context to tunnel connect

---
 pkg/agent/run.go           |  2 +-
 pkg/agent/tunnel/tunnel.go | 15 +++++++++++++--
 2 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/pkg/agent/run.go b/pkg/agent/run.go
index 640e19c26a..9a4521149f 100644
--- a/pkg/agent/run.go
+++ b/pkg/agent/run.go
@@ -47,7 +47,7 @@ func run(ctx context.Context, cfg cmds.Agent) error {
 		return err
 	}
 
-	if err := tunnel.Setup(nodeConfig); err != nil {
+	if err := tunnel.Setup(ctx, nodeConfig); err != nil {
 		return err
 	}
 
diff --git a/pkg/agent/tunnel/tunnel.go b/pkg/agent/tunnel/tunnel.go
index ab02f2d513..f7884591c5 100644
--- a/pkg/agent/tunnel/tunnel.go
+++ b/pkg/agent/tunnel/tunnel.go
@@ -52,7 +52,7 @@ func getAddresses(endpoint *v1.Endpoints) []string {
 	return serverAddresses
 }
 
-func Setup(config *config.Node) error {
+func Setup(ctx context.Context, config *config.Node) error {
 	restConfig, err := clientcmd.BuildConfigFromFlags("", config.AgentConfig.KubeConfigNode)
 	if err != nil {
 		return err
@@ -133,7 +133,18 @@ func Setup(config *config.Node) error {
 		}
 	}()
 
-	wg.Wait()
+	wait := make(chan int, 1)
+	go func() {
+		wg.Wait()
+		wait <- 0
+	}()
+
+	select {
+	case <-ctx.Done():
+		logrus.Error("tunnel context canceled while waiting for connection")
+		return ctx.Err()
+	case <-wait:
+	}
 
 	return nil
 }