fix(coderd): add timeout to websocket waitgroup on shutdown (coder#12754)

coadler · web-flow · commit 5f28220eec60 · 2024-03-26T03:04:15.000Z
diff --git a/coderd/coderd.go b/coderd/coderd.go
@@ -1285,9 +1285,23 @@ func (api *API) Close() error {
 		api.derpCloseFunc()
 	}
 
-	api.WebsocketWaitMutex.Lock()
-	api.WebsocketWaitGroup.Wait()
-	api.WebsocketWaitMutex.Unlock()
+	wsDone := make(chan struct{})
+	timer := time.NewTimer(10 * time.Second)
+	defer timer.Stop()
+	go func() {
+		api.WebsocketWaitMutex.Lock()
+		defer api.WebsocketWaitMutex.Unlock()
+		api.WebsocketWaitGroup.Wait()
+		close(wsDone)
+	}()
+	// This will technically leak the above func if the timer fires, but this is
+	// maintly a last ditch effort to un-stuck coderd on shutdown. This
+	// shouldn't affect tests at all.
+	select {
+	case <-wsDone:
+	case <-timer.C:
+		api.Logger.Warn(api.ctx, "websocket shutdown timed out after 10 seconds")
+	}
 
 	api.dbRolluper.Close()
 	api.metricsCache.Close()
