From b7c270ccb31e8ed7ebed9fa31268259e72f64072 Mon Sep 17 00:00:00 2001
From: Ev Kontsevoy <ev@kontsevoy.com>
Date: Wed, 29 Jun 2016 16:50:48 -0700
Subject: [PATCH] Goroutine leak when closing reverse tunnels

Fixes #468

Sasha, let me know if you want to discuss. It's a small change but if
dont incorrectly can be hard to debug later.
---
 lib/reversetunnel/agent.go | 17 +++++++++++++----
 1 file changed, 13 insertions(+), 4 deletions(-)

diff --git a/lib/reversetunnel/agent.go b/lib/reversetunnel/agent.go
index ebb9a1577ecc4..3a8d5a2643d02 100644
--- a/lib/reversetunnel/agent.go
+++ b/lib/reversetunnel/agent.go
@@ -243,6 +243,9 @@ func (a *Agent) proxyTransport(ch ssh.Channel, reqC <-chan *ssh.Request) {
 // to the given SSH connection.
 //
 func (a *Agent) runHeartbeat(conn *ssh.Client) {
+	ticker := time.NewTicker(defaults.ReverseTunnelAgentHeartbeatPeriod)
+	defer ticker.Stop()
+
 	heartbeatLoop := func() error {
 		if conn == nil {
 			return trace.Errorf("heartbeat cannot ping: need to reconnect")
@@ -258,8 +261,6 @@ func (a *Agent) runHeartbeat(conn *ssh.Client) {
 
 		// send first ping right away, then start a ping timer:
 		hb.SendRequest("ping", false, nil)
-		ticker := time.NewTicker(defaults.ReverseTunnelAgentHeartbeatPeriod)
-		defer ticker.Stop()
 
 		for {
 			select {
@@ -307,10 +308,18 @@ func (a *Agent) runHeartbeat(conn *ssh.Client) {
 		}
 	}
 
+	// run heartbeat loop, and when it fails (probably means that a tunnel got disconnected)
+	// keep repeating to reconnect until we're asked to stop
 	err := heartbeatLoop()
 	if err != nil || conn == nil {
-		time.Sleep(defaults.ReverseTunnelAgentHeartbeatPeriod)
-		a.Start()
+		select {
+		// abort if asked to stop:
+		case <-a.broadcastClose.C:
+			return
+			// reconnect
+		case <-ticker.C:
+			a.Start()
+		}
 	}
 }