From dda56f0bcb7d101d5a8fae8706ee16f92c7ae155 Mon Sep 17 00:00:00 2001
From: Spike Curtis <spike@coder.com>
Date: Fri, 4 Apr 2025 10:12:30 +0000
Subject: [PATCH] fix: fix closeMutex unlock bug

---
 agent/agent.go | 20 ++++++++++++++------
 1 file changed, 14 insertions(+), 6 deletions(-)

diff --git a/agent/agent.go b/agent/agent.go
index 1e18290d84b6a..1d81fe3209e25 100644
--- a/agent/agent.go
+++ b/agent/agent.go
@@ -1518,14 +1518,11 @@ func (a *agent) runCoordinator(ctx context.Context, tClient tailnetproto.DRPCTai
 	a.logger.Info(ctx, "connected to coordination RPC")
 
 	// This allows the Close() routine to wait for the coordinator to gracefully disconnect.
-	a.closeMutex.Lock()
-	if a.isClosed() {
-		return nil
+	disconnected := a.setCoordDisconnected()
+	if disconnected == nil {
+		return nil // already closed by something else
 	}
-	disconnected := make(chan struct{})
-	a.coordDisconnected = disconnected
 	defer close(disconnected)
-	a.closeMutex.Unlock()
 
 	ctrl := tailnet.NewAgentCoordinationController(a.logger, network)
 	coordination := ctrl.New(coordinate)
@@ -1547,6 +1544,17 @@ func (a *agent) runCoordinator(ctx context.Context, tClient tailnetproto.DRPCTai
 	return <-errCh
 }
 
+func (a *agent) setCoordDisconnected() chan struct{} {
+	a.closeMutex.Lock()
+	defer a.closeMutex.Unlock()
+	if a.isClosed() {
+		return nil
+	}
+	disconnected := make(chan struct{})
+	a.coordDisconnected = disconnected
+	return disconnected
+}
+
 // runDERPMapSubscriber runs a coordinator and returns if a reconnect should occur.
 func (a *agent) runDERPMapSubscriber(ctx context.Context, tClient tailnetproto.DRPCTailnetClient24, network *tailnet.Conn) error {
 	defer a.logger.Debug(ctx, "disconnected from derp map RPC")