opensource/coder
3
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-09-21 20:51:01 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix: pass OnSubscribe to HA MultiAgent (#9947)

Browse Source 
Fixes https://github.com/coder/coder/issues/9929
This commit is contained in:
Colin Adler
2023-09-29 13:37:17 -05:00
committed by GitHub
parent 61154a6bb5
commit 4da1223a80
1 changed files with 28 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
enterprise/tailnet/coordinator.go
+28 -13
View File
@@ -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 {