fix: ensure wsproxy MultiAgent is closed when websocket dies

coadler · coadler · commit 4b6777d328bd · 2024-01-04T19:26:21.000Z
The `SingleTailnet` behavior only checked to see if the `MultiAgent` was
closed, but the websocket error was not being propogated into the
`MultiAgent`, causing it to never be swapped for a new working one.
diff --git a/coderd/httpapi/websocket.go b/coderd/httpapi/websocket.go
@@ -26,10 +26,10 @@ func Heartbeat(ctx context.Context, conn *websocket.Conn) {
 	}
 }
 
-// Heartbeat loops to ping a WebSocket to keep it alive. It kills the connection
-// on ping failure.
+// Heartbeat loops to ping a WebSocket to keep it alive. It calls `exit` on ping
+// failure.
 func HeartbeatClose(ctx context.Context, exit func(), conn *websocket.Conn) {
-	ticker := time.NewTicker(30 * time.Second)
+	ticker := time.NewTicker(15 * time.Second)
 	defer ticker.Stop()
 
 	for {
diff --git a/coderd/tailnet.go b/coderd/tailnet.go
@@ -224,6 +224,7 @@ func (s *ServerTailnet) watchAgentUpdates() {
 		nodes, ok := conn.NextUpdate(s.ctx)
 		if !ok {
 			if conn.IsClosed() && s.ctx.Err() == nil {
+				s.logger.Warn(s.ctx, "multiagent closed, reinitializing")
 				s.reinitCoordinator()
 				continue
 			}
diff --git a/enterprise/wsproxy/wsproxysdk/wsproxysdk.go b/enterprise/wsproxy/wsproxysdk/wsproxysdk.go
@@ -472,6 +472,11 @@ func (c *Client) DialCoordinator(ctx context.Context) (agpl.MultiAgentConn, erro
 		OnRemove:          func(agpl.Queue) { conn.Close(websocket.StatusGoingAway, "closed") },
 	}).Init()
 
+	go func() {
+		<-ctx.Done()
+		ma.Close()
+	}()
+
 	go func() {
 		defer cancel()
 		dec := json.NewDecoder(nc)

Original file line number	Diff line number	Diff line change
`@@ -26,10 +26,10 @@ func Heartbeat(ctx context.Context, conn *websocket.Conn) {`
`26`	`26`	`}`
`27`	`27`	`}`
`28`	`28`
`29`		`-// Heartbeat loops to ping a WebSocket to keep it alive. It kills the connection`
`30`		`-// on ping failure.`
	`29`	+// Heartbeat loops to ping a WebSocket to keep it alive. It calls `exit` on ping
	`30`	`+// failure.`
`31`	`31`	`func HeartbeatClose(ctx context.Context, exit func(), conn *websocket.Conn) {`
`32`		`- ticker := time.NewTicker(30 * time.Second)`
	`32`	`+ ticker := time.NewTicker(15 * time.Second)`
`33`	`33`	`defer ticker.Stop()`
`34`	`34`
`35`	`35`	`for {`
Original file line number	Diff line number	Diff line change
`@@ -224,6 +224,7 @@ func (s *ServerTailnet) watchAgentUpdates() {`
`224`	`224`	`nodes, ok := conn.NextUpdate(s.ctx)`
`225`	`225`	`if !ok {`
`226`	`226`	`if conn.IsClosed() && s.ctx.Err() == nil {`
	`227`	`+ s.logger.Warn(s.ctx, "multiagent closed, reinitializing")`
`227`	`228`	`s.reinitCoordinator()`
`228`	`229`	`continue`
`229`	`230`	`}`