From d84d792d100a02e885ea6b9b75207704553da0d9 Mon Sep 17 00:00:00 2001
From: Mathias Fredriksson <mafredri@gmail.com>
Date: Fri, 23 Sep 2022 11:07:38 +0300
Subject: [PATCH 1/2] fix: Close coordinator on context cancellation

---
 agent/agent.go | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/agent/agent.go b/agent/agent.go
index 18243ee788789..52d8c0db59c9f 100644
--- a/agent/agent.go
+++ b/agent/agent.go
@@ -288,6 +288,8 @@ func (a *agent) runCoordinator(ctx context.Context) {
 			a.logger.Warn(context.Background(), "failed to dial", slog.Error(err))
 			continue
 		}
+		//nolint:revive // Defer is ok because we're exiting this loop.
+		defer coordinator.Close()
 		a.logger.Info(context.Background(), "connected to coordination server")
 		break
 	}
@@ -296,7 +298,6 @@ func (a *agent) runCoordinator(ctx context.Context) {
 		return
 	default:
 	}
-	defer coordinator.Close()
 	sendNodes, errChan := tailnet.ServeCoordinator(coordinator, a.network.UpdateNodes)
 	a.network.SetNodeCallback(sendNodes)
 	select {

From 7e2ad50827cab32b054563f04efb320c552ad714 Mon Sep 17 00:00:00 2001
From: Mathias Fredriksson <mafredri@gmail.com>
Date: Fri, 23 Sep 2022 11:08:46 +0300
Subject: [PATCH 2/2] fix: Refactor runCoordinator so that previous is
 closed/freed

---
 agent/agent.go | 24 ++++++++++++++++--------
 1 file changed, 16 insertions(+), 8 deletions(-)

diff --git a/agent/agent.go b/agent/agent.go
index 52d8c0db59c9f..16badf1034d2c 100644
--- a/agent/agent.go
+++ b/agent/agent.go
@@ -272,6 +272,15 @@ func (a *agent) runTailnet(ctx context.Context, derpMap *tailcfg.DERPMap) {
 
 // runCoordinator listens for nodes and updates the self-node as it changes.
 func (a *agent) runCoordinator(ctx context.Context) {
+	for {
+		reconnect := a.runCoordinatorWithRetry(ctx)
+		if !reconnect {
+			return
+		}
+	}
+}
+
+func (a *agent) runCoordinatorWithRetry(ctx context.Context) (reconnect bool) {
 	var coordinator net.Conn
 	var err error
 	// An exponential back-off occurs when the connection is failing to dial.
@@ -280,10 +289,10 @@ func (a *agent) runCoordinator(ctx context.Context) {
 		coordinator, err = a.coordinatorDialer(ctx)
 		if err != nil {
 			if errors.Is(err, context.Canceled) {
-				return
+				return false
 			}
 			if a.isClosed() {
-				return
+				return false
 			}
 			a.logger.Warn(context.Background(), "failed to dial", slog.Error(err))
 			continue
@@ -295,24 +304,23 @@ func (a *agent) runCoordinator(ctx context.Context) {
 	}
 	select {
 	case <-ctx.Done():
-		return
+		return false
 	default:
 	}
 	sendNodes, errChan := tailnet.ServeCoordinator(coordinator, a.network.UpdateNodes)
 	a.network.SetNodeCallback(sendNodes)
 	select {
 	case <-ctx.Done():
-		return
+		return false
 	case err := <-errChan:
 		if a.isClosed() {
-			return
+			return false
 		}
 		if errors.Is(err, context.Canceled) {
-			return
+			return false
 		}
 		a.logger.Debug(ctx, "node broker accept exited; restarting connection", slog.Error(err))
-		a.runCoordinator(ctx)
-		return
+		return true
 	}
 }