From d30d76bb928d93848b8d035cf2dccc0e6cb1a746 Mon Sep 17 00:00:00 2001 From: rcourtman Date: Fri, 7 Nov 2025 10:20:26 +0000 Subject: [PATCH] Fix P1: Add shutdown mechanism to WebSocket Hub Fixed goroutine leaks in WebSocket hub from missing shutdown mechanism: Problem: 1. Hub.Run() has infinite loop with no exit condition 2. runBroadcastSequencer() reads from channel forever 3. No way to cleanly shutdown hub during restarts or tests Solution: - Added stopChan chan struct{} field to Hub - Initialize stopChan in NewHub() - Added Stop() method that closes stopChan - Modified Run() main loop to select on stopChan - On shutdown: close all client connections and return - Modified runBroadcastSequencer() from 'for range' to select - Changed from: for msg := range h.broadcastSeq - Changed to: for { select { case msg := <-h.broadcastSeq: ... case <-h.stopChan: ... }} - On shutdown: stop coalesce timer and return Shutdown sequence: 1. Call hub.Stop() to close stopChan 2. Both Run() and runBroadcastSequencer() exit their loops 3. All client send channels are closed 4. Clients map is cleared 5. Pending coalesce timer is stopped Impact: - Enables graceful shutdown during service restarts - Prevents goroutine leaks in tests - Allows proper cleanup of WebSocket connections - No more orphaned broadcast sequencer goroutines --- internal/websocket/hub.go | 115 ++++++++++++++++++++++++-------------- 1 file changed, 73 insertions(+), 42 deletions(-) diff --git a/internal/websocket/hub.go b/internal/websocket/hub.go index a538d15..bde48b8 100644 --- a/internal/websocket/hub.go +++ b/internal/websocket/hub.go @@ -272,6 +272,7 @@ type Hub struct { broadcastSeq chan Message // Sequenced broadcast channel for ordering register chan *Client unregister chan *Client + stopChan chan struct{} // Signals shutdown mu sync.RWMutex getState func() interface{} // Function to get current state allowedOrigins []string // Allowed origins for CORS @@ -305,6 +306,7 @@ func NewHub(getState func() interface{}) *Hub { broadcastSeq: make(chan Message, 256), // Buffered sequenced channel register: make(chan *Client), unregister: make(chan *Client), + stopChan: make(chan struct{}), getState: getState, allowedOrigins: []string{}, // Default to empty (will be set based on actual host) coalesceWindow: 100 * time.Millisecond, // Coalesce rapid updates within 100ms @@ -430,10 +432,26 @@ func (h *Hub) Run() { case <-pingTicker.C: h.sendPing() + + case <-h.stopChan: + log.Info().Msg("WebSocket hub shutting down") + // Close all client connections + h.mu.Lock() + for client := range h.clients { + close(client.send) + } + h.clients = make(map[*Client]bool) + h.mu.Unlock() + return } } } +// Stop gracefully shuts down the hub +func (h *Hub) Stop() { + close(h.stopChan) +} + // HandleWebSocket handles WebSocket upgrade requests func (h *Hub) HandleWebSocket(w http.ResponseWriter, r *http.Request) { log.Info(). @@ -475,54 +493,67 @@ func (h *Hub) HandleWebSocket(w http.ResponseWriter, r *http.Request) { // runBroadcastSequencer handles sequenced broadcasts with coalescing for rapid state updates func (h *Hub) runBroadcastSequencer() { - for msg := range h.broadcastSeq { - // Handle raw data (state) messages with coalescing - if msg.Type == "rawData" { - h.coalesceMutex.Lock() + for { + select { + case msg := <-h.broadcastSeq: + // Handle raw data (state) messages with coalescing + if msg.Type == "rawData" { + h.coalesceMutex.Lock() + // Cancel pending timer if exists + if h.coalesceTimer != nil { + h.coalesceTimer.Stop() + } + + // Update pending message + h.coalescePending = &msg + + // Set timer to send after coalesce window + h.coalesceTimer = time.AfterFunc(h.coalesceWindow, func() { + h.coalesceMutex.Lock() + if h.coalescePending != nil { + // Send the coalesced message + if data, err := json.Marshal(*h.coalescePending); err == nil { + h.mu.RLock() + for client := range h.clients { + select { + case client.send <- data: + default: + log.Warn().Str("client", client.id).Msg("Client send channel full, skipping coalesced message") + } + } + h.mu.RUnlock() + } + h.coalescePending = nil + } + h.coalesceMutex.Unlock() + }) + + h.coalesceMutex.Unlock() + } else { + // Non-state messages (alerts, etc.) - send immediately + if data, err := json.Marshal(msg); err == nil { + h.mu.RLock() + for client := range h.clients { + select { + case client.send <- data: + default: + log.Warn().Str("client", client.id).Msg("Client send channel full, skipping message") + } + } + h.mu.RUnlock() + } + } + + case <-h.stopChan: + log.Debug().Msg("Broadcast sequencer shutting down") // Cancel pending timer if exists + h.coalesceMutex.Lock() if h.coalesceTimer != nil { h.coalesceTimer.Stop() } - - // Update pending message - h.coalescePending = &msg - - // Set timer to send after coalesce window - h.coalesceTimer = time.AfterFunc(h.coalesceWindow, func() { - h.coalesceMutex.Lock() - if h.coalescePending != nil { - // Send the coalesced message - if data, err := json.Marshal(*h.coalescePending); err == nil { - h.mu.RLock() - for client := range h.clients { - select { - case client.send <- data: - default: - log.Warn().Str("client", client.id).Msg("Client send channel full, skipping coalesced message") - } - } - h.mu.RUnlock() - } - h.coalescePending = nil - } - h.coalesceMutex.Unlock() - }) - h.coalesceMutex.Unlock() - } else { - // Non-state messages (alerts, etc.) - send immediately - if data, err := json.Marshal(msg); err == nil { - h.mu.RLock() - for client := range h.clients { - select { - case client.send <- data: - default: - log.Warn().Str("client", client.id).Msg("Client send channel full, skipping message") - } - } - h.mu.RUnlock() - } + return } } }