From 4da1223a807e50dd9f307c61bcddd4315dd2ad73 Mon Sep 17 00:00:00 2001
From: Colin Adler <colin1adler@gmail.com>
Date: Fri, 29 Sep 2023 13:37:17 -0500
Subject: [PATCH] fix: pass `OnSubscribe` to HA MultiAgent (#9947)

Fixes https://github.com/coder/coder/issues/9929
---
 enterprise/tailnet/coordinator.go | 41 +++++++++++++++++++++----------
 1 file changed, 28 insertions(+), 13 deletions(-)

diff --git a/enterprise/tailnet/coordinator.go b/enterprise/tailnet/coordinator.go
index 70ad50687b..5a26cdc92a 100644
--- a/enterprise/tailnet/coordinator.go
+++ b/enterprise/tailnet/coordinator.go
@@ -57,8 +57,9 @@ func (c *haCoordinator) ServeMultiAgent(id uuid.UUID) agpl.MultiAgentConn {
 		ID:                id,
 		AgentIsLegacyFunc: c.agentIsLegacy,
 		OnSubscribe:       c.clientSubscribeToAgent,
+		OnUnsubscribe:     c.clientUnsubscribeFromAgent,
 		OnNodeUpdate:      c.clientNodeUpdate,
-		OnRemove:          func(enq agpl.Queue) { c.clientDisconnected(enq.UniqueID()) },
+		OnRemove:          c.clientDisconnected,
 	}).Init()
 	c.addClient(id, m)
 	return m
@@ -101,6 +102,22 @@ func (c *haCoordinator) clientSubscribeToAgent(enq agpl.Queue, agentID uuid.UUID
 	return nil, nil
 }
 
+func (c *haCoordinator) clientUnsubscribeFromAgent(enq agpl.Queue, agentID uuid.UUID) error {
+	c.mutex.Lock()
+	defer c.mutex.Unlock()
+
+	connectionSockets, ok := c.agentToConnectionSockets[agentID]
+	if !ok {
+		return nil
+	}
+	delete(connectionSockets, enq.UniqueID())
+	if len(connectionSockets) == 0 {
+		delete(c.agentToConnectionSockets, agentID)
+	}
+
+	return nil
+}
+
 type haCoordinator struct {
 	id        uuid.UUID
 	log       slog.Logger
@@ -161,7 +178,7 @@ func (c *haCoordinator) ServeClient(conn net.Conn, id, agentID uuid.UUID) error
 	defer tc.Close()
 
 	c.addClient(id, tc)
-	defer c.clientDisconnected(id)
+	defer c.clientDisconnected(tc)
 
 	agentNode, err := c.clientSubscribeToAgent(tc, agentID)
 	if err != nil {
@@ -200,26 +217,24 @@ func (c *haCoordinator) initOrSetAgentConnectionSocketLocked(agentID uuid.UUID,
 	c.clientsToAgents[enq.UniqueID()][agentID] = c.agentSockets[agentID]
 }
 
-func (c *haCoordinator) clientDisconnected(id uuid.UUID) {
+func (c *haCoordinator) clientDisconnected(enq agpl.Queue) {
 	c.mutex.Lock()
 	defer c.mutex.Unlock()
 
-	for agentID := range c.clientsToAgents[id] {
-		// Clean all traces of this connection from the map.
-		delete(c.nodes, id)
+	for agentID := range c.clientsToAgents[enq.UniqueID()] {
 		connectionSockets, ok := c.agentToConnectionSockets[agentID]
 		if !ok {
-			return
+			continue
 		}
-		delete(connectionSockets, id)
-		if len(connectionSockets) != 0 {
-			return
+		delete(connectionSockets, enq.UniqueID())
+		if len(connectionSockets) == 0 {
+			delete(c.agentToConnectionSockets, agentID)
 		}
-		delete(c.agentToConnectionSockets, agentID)
 	}
 
-	delete(c.clients, id)
-	delete(c.clientsToAgents, id)
+	delete(c.nodes, enq.UniqueID())
+	delete(c.clients, enq.UniqueID())
+	delete(c.clientsToAgents, enq.UniqueID())
 }
 
 func (c *haCoordinator) handleNextClientMessage(id uuid.UUID, decoder *json.Decoder) error {