chore: improve websocket handling

Author: MarcStdt

internal/core/domain/broadcast_channel.go

@@ -1,6 +1,8 @@
 package domain
 
 import (
+	"sync"
+
 	"github.com/highcard-dev/daemon/internal/utils/logger"
 )
 
@@ -10,31 +12,46 @@ type BroadcastChannel struct {
 
 	// Inbound messages from the clients.
 	Broadcast chan []byte
+
+	// Mutex to protect concurrent access to Clients map
+	mu sync.RWMutex
 }
 
 func NewHub() *BroadcastChannel {
 	return &BroadcastChannel{
-		Broadcast: make(chan []byte),
+		Broadcast: make(chan []byte, 100), // Buffered channel to handle bursts of file changes
 		Clients:   make(map[chan *[]byte]bool),
 	}
 }
 
 func (h *BroadcastChannel) Subscribe() chan *[]byte {
-	client := make(chan *[]byte, 10) // buffered channel to avoid blocking
+	h.mu.Lock()
+	defer h.mu.Unlock()
+
+	client := make(chan *[]byte, 50) // Increased buffer size to handle message bursts
 	h.Clients[client] = true
 	return client
 }
 
 func (h *BroadcastChannel) Unsubscribe(client chan *[]byte) {
-	delete(h.Clients, client)
-	close(client)
+	h.mu.Lock()
+	defer h.mu.Unlock()
+
+	if _, exists := h.Clients[client]; exists {
+		delete(h.Clients, client)
+		close(client)
+	}
 }
 
 func (h *BroadcastChannel) CloseChannel() {
+	h.mu.Lock()
+	defer h.mu.Unlock()
+
 	close(h.Broadcast)
 	for client := range h.Clients {
-		h.Unsubscribe(client)
+		close(client)
 	}
+	h.Clients = make(map[chan *[]byte]bool)
 }
 
 func (h *BroadcastChannel) Run() {
@@ -44,15 +61,39 @@ func (h *BroadcastChannel) Run() {
 			logger.Log().Debug("Broadcast channel closed")
 			return
 		}
+
+		h.mu.RLock()
+		clients := make([]chan *[]byte, 0, len(h.Clients))
 		for client := range h.Clients {
+			clients = append(clients, client)
+		}
+		h.mu.RUnlock()
+
+		// Track clients to remove (dead connections)
+		var deadClients []chan *[]byte
+
+		for _, client := range clients {
 			select {
-			case client <- &message: // Try to send the message.
+			case client <- &message:
+				// Successfully sent the message
 			default:
-				// logger.Log().Warn("Channel seems to be closed, removing")
-				logger.Log().Warn("Channel seems to be closed")
-				// delete(h.Clients, client)
-				// h.Unsubscribe(client)
+				// Client channel is blocked or closed, mark for removal
+				logger.Log().Debug("Client channel blocked or closed, marking for removal")
+				deadClients = append(deadClients, client)
+			}
+		}
+
+		// Remove dead clients
+		if len(deadClients) > 0 {
+			h.mu.Lock()
+			for _, deadClient := range deadClients {
+				if _, exists := h.Clients[deadClient]; exists {
+					delete(h.Clients, deadClient)
+					// Don't close the channel here as it might be closed elsewhere
+					logger.Log().Debug("Removed dead client from broadcast channel")
+				}
 			}
+			h.mu.Unlock()
 		}
 	}
 }

internal/core/services/ui_dev_service.go

@@ -252,12 +252,12 @@ func (uds *UiDevService) handleFileEvent(event fsnotify.Event) {
 		return
 	}
 
-	// Broadcast the event to all subscribers
+	// Broadcast the event to all subscribers with a timeout to prevent blocking
 	select {
 	case uds.broadcastChannel.Broadcast <- eventData:
 		logger.Log().Debug("File change event broadcasted", zap.String("path", event.Name), zap.String("op", event.Op.String()))
-	default:
-		logger.Log().Warn("Failed to broadcast file change event - channel full")
+	case <-time.After(100 * time.Millisecond): // Timeout after 100ms
+		logger.Log().Warn("Failed to broadcast file change event - channel full (timed out)")
 	}
 
 	// Handle directory creation - add new directories to watcher

internal/handler/ui_dev_handler.go

@@ -2,7 +2,10 @@ package handler
 
 import (
 	"encoding/json"
+	"net"
 	"path/filepath"
+	"strings"
+	"syscall"
 	"time"
 
 	"github.com/gofiber/contrib/websocket"
@@ -42,6 +45,36 @@ func NewUiDevHandler(uiDevService ports.UiDevServiceInterface, scrollService por
 	}
 }
 
+// isConnectionError checks if the error is related to a broken connection
+func isConnectionError(err error) bool {
+	if err == nil {
+		return false
+	}
+
+	// Check for common connection error patterns
+	errStr := strings.ToLower(err.Error())
+	if strings.Contains(errStr, "broken pipe") ||
+		strings.Contains(errStr, "connection reset") ||
+		strings.Contains(errStr, "connection refused") ||
+		strings.Contains(errStr, "use of closed network connection") {
+		return true
+	}
+
+	// Check for specific error types
+	if netErr, ok := err.(*net.OpError); ok {
+		if netErr.Op == "write" {
+			return true
+		}
+	}
+
+	// Check for syscall errors
+	if errno, ok := err.(syscall.Errno); ok {
+		return errno == syscall.EPIPE || errno == syscall.ECONNRESET
+	}
+
+	return false
+}
+
 // @Summary Enable development mode
 // @ID enableDev
 // @Tags ui, dev, druid, daemon
@@ -148,19 +181,23 @@ func (udh *UiDevHandler) Status(ctx *fiber.Ctx) error {
 
 // NotifyChange handles WebSocket connections for real-time file change notifications
 func (udh *UiDevHandler) NotifyChange(c *websocket.Conn) {
-	defer func() {
-		if err := c.Close(); err != nil {
-			logger.Log().Error("Error closing WebSocket connection", zap.Error(err))
-		}
-	}()
-
 	// Set connection timeouts
 	const (
 		writeWait  = 10 * time.Second
 		pongWait   = 60 * time.Second
 		pingPeriod = (pongWait * 9) / 10
 	)
 
+	// Create a done channel to signal when the connection should be closed
+	done := make(chan struct{})
+
+	defer func() {
+		close(done)
+		if err := c.Close(); err != nil {
+			logger.Log().Debug("Error closing WebSocket connection", zap.Error(err))
+		}
+	}()
+
 	c.SetReadDeadline(time.Now().Add(pongWait))
 	c.SetPongHandler(func(string) error {
 		c.SetReadDeadline(time.Now().Add(pongWait))
@@ -199,17 +236,50 @@ func (udh *UiDevHandler) NotifyChange(c *websocket.Conn) {
 		"watchedPaths": udh.uiDevService.GetWatchedPaths(),
 		"timestamp":    time.Now(),
 	}
-	c.WriteJSON(connectMsg)
+
+	c.SetWriteDeadline(time.Now().Add(writeWait))
+	if err := c.WriteJSON(connectMsg); err != nil {
+		logger.Log().Error("Failed to send initial connection message", zap.Error(err))
+		return
+	}
 
 	logger.Log().Info("WebSocket client connected for file change notifications")
 
 	// Start ping ticker
 	ticker := time.NewTicker(pingPeriod)
 	defer ticker.Stop()
 
+	// Start a goroutine to read messages (to handle pong responses and detect broken connections)
+	go func() {
+		defer func() {
+			select {
+			case <-done:
+				// Connection is already being closed
+			default:
+				close(done)
+			}
+		}()
+
+		for {
+			_, _, err := c.ReadMessage()
+			if err != nil {
+				if isConnectionError(err) {
+					logger.Log().Debug("WebSocket client disconnected", zap.Error(err))
+				} else {
+					logger.Log().Warn("WebSocket read error", zap.Error(err))
+				}
+				return
+			}
+		}
+	}()
+
 	// Handle messages and file changes
 	for {
 		select {
+		case <-done:
+			logger.Log().Debug("WebSocket connection done signal received")
+			return
+
 		case data := <-changesChan:
 			if data == nil {
 				logger.Log().Info("File change channel closed")
@@ -232,7 +302,11 @@ func (udh *UiDevHandler) NotifyChange(c *websocket.Conn) {
 
 			c.SetWriteDeadline(time.Now().Add(writeWait))
 			if err := c.WriteJSON(changeMessage); err != nil {
-				logger.Log().Error("Failed to write file change to WebSocket", zap.Error(err))
+				if isConnectionError(err) {
+					logger.Log().Debug("WebSocket client disconnected while sending file change", zap.Error(err))
+				} else {
+					logger.Log().Error("Failed to write file change to WebSocket", zap.Error(err))
+				}
 				return
 			}
 
@@ -243,7 +317,11 @@ func (udh *UiDevHandler) NotifyChange(c *websocket.Conn) {
 		case <-ticker.C:
 			c.SetWriteDeadline(time.Now().Add(writeWait))
 			if err := c.WriteMessage(websocket.PingMessage, nil); err != nil {
-				logger.Log().Error("Failed to send ping", zap.Error(err))
+				if isConnectionError(err) {
+					logger.Log().Debug("WebSocket client disconnected during ping", zap.Error(err))
+				} else {
+					logger.Log().Error("Failed to send ping", zap.Error(err))
+				}
 				return
 			}
 		}

internal/handler/websocket_handler.go

@@ -7,6 +7,8 @@ import (
 	"github.com/gofiber/fiber/v2"
 	"github.com/highcard-dev/daemon/internal/core/domain"
 	"github.com/highcard-dev/daemon/internal/core/ports"
+	"github.com/highcard-dev/daemon/internal/utils/logger"
+	"go.uber.org/zap"
 )
 
 type WebsocketHandler struct {
@@ -82,8 +84,12 @@ func (ah WebsocketHandler) Consoles(c *fiber.Ctx) error {
 func (wh WebsocketHandler) HandleProcess(c *websocket.Conn) {
 	param := c.Params("console")
 
+	// Create a done channel to signal when the connection should be closed
+	done := make(chan struct{})
+
 	ticker := time.NewTicker(pingPeriod)
 	defer func() {
+		close(done)
 		ticker.Stop()
 		c.Close()
 	}()
@@ -94,31 +100,55 @@ func (wh WebsocketHandler) HandleProcess(c *websocket.Conn) {
 	}
 
 	subscriptionChannel := channel.Channel.Subscribe()
-
 	defer channel.Channel.Unsubscribe(subscriptionChannel)
 
 	c.SetReadLimit(maxMessageSize)
 	c.SetReadDeadline(time.Now().Add(pongWait))
 	c.SetPongHandler(func(string) error { c.SetReadDeadline(time.Now().Add(pongWait)); return nil })
 
+	// Start a goroutine to read messages (to handle pong responses and detect broken connections)
+	go func() {
+		defer func() {
+			select {
+			case <-done:
+				// Connection is already being closed
+			default:
+				close(done)
+			}
+		}()
+
+		for {
+			_, _, err := c.ReadMessage()
+			if err != nil {
+				logger.Log().Debug("WebSocket client disconnected", zap.Error(err))
+				return
+			}
+		}
+	}()
+
 	//fetch channel and send to websocket
 	for {
 		select {
+		case <-done:
+			logger.Log().Debug("WebSocket connection done signal received")
+			return
+
 		//send 1024 bytes at a time
 		case buffer, ok := <-subscriptionChannel:
-
 			c.SetWriteDeadline(time.Now().Add(writeWait))
 			//if nil is send, assume the channel is closed
 			if buffer == nil || !ok {
 				return
 			}
 			err := c.WriteMessage(websocket.TextMessage, *buffer)
 			if err != nil {
+				logger.Log().Debug("WebSocket client disconnected while sending message", zap.Error(err))
 				return
 			}
 		case <-ticker.C:
 			c.SetWriteDeadline(time.Now().Add(writeWait))
 			if err := c.WriteMessage(websocket.PingMessage, nil); err != nil {
+				logger.Log().Debug("WebSocket client disconnected during ping", zap.Error(err))
 				return
 			}
 		}