cmd/server/main.go

@@ -54,7 +54,6 @@ func Start() {
 	r := http.NewServeMux()
 	r.HandleFunc("/", websocket.HandleConnection)
 	go websocket.SendPermanentBroadCastMessage()
-	go websocket.RemoveStaleClients()
 
 	helper.Logger.Info("GoTomato started", "version", metadata.GoTomatoVersion)
 	helper.Logger.Info("Websocket listening", "address", listen)

internal/websocket/broadcast.go

@@ -21,15 +21,18 @@ func SendPermanentBroadCastMessage() {
 		}
 
 		// Iterate over all connected clients and broadcast the message
-		mu.Lock()
 		for _, client := range Clients {
+			// Remove unresponsive client
+			if client.IsStale() {
+				client.Conn.Close()
+				continue
+			}
 			// Send message to client
 			err := client.SendMessage(websocket.TextMessage, jsonMessage)
 			if err != nil {
 				helper.Logger.Info("Error broadcasting to client:", "msg", err, "host", client.RealIP, "clients", len(Clients))
 			}
 		}
-		mu.Unlock()
 
 		<-tick.C
 	}

internal/websocket/client_commands.go

@@ -22,9 +22,7 @@ func handleClientCommands(c models.WebsocketClient) {
 		_, message, err := ws.ReadMessage()
 		if err != nil {
 			// remove client on error/disconnect
-			mu.Lock()
 			delete(Clients, ws.LocalAddr())
-			mu.Unlock()
 			helper.Logger.Info("Client disconnected:", "msg", err, "host", c.RealIP, "clients", len(Clients))
 			break
 		}

internal/websocket/handle_connection.go

@@ -36,7 +36,8 @@ func HandleConnection(w http.ResponseWriter, r *http.Request) {
 		LastPong: time.Now(),
 		RealIP:   r.RemoteAddr,
 	}
-	client.Conn.SetPongHandler(func(s string) error {
+
+	client.Conn.SetPongHandler(func(appData string) error {
 		client.LastPong = time.Now()
 		return nil
 	})

internal/websocket/staleClients.go

@@ -1,21 +0,0 @@
-package websocket
-
-import (
-	"time"
-)
-
-// Check and remove stale clients
-func RemoveStaleClients() {
-	ticker := time.NewTicker(30 * time.Second)
-	defer ticker.Stop()
-
-	for range ticker.C {
-		mu.Lock()
-		for _, client := range Clients {
-			if client.IsStale() {
-				client.Conn.Close()
-			}
-		}
-		mu.Unlock()
-	}
-}

pkg/models/client.go

@@ -39,7 +39,7 @@ func (c *WebsocketClient) SendMessage(messageType int, data []byte) error {
 	return nil
 }
 
-// Checks if the websockets last Pong is recent
+// Check if websockets last Pong is recent
 func (c *WebsocketClient) IsStale() bool {
 	return time.Since(c.LastPong) > (TIMEOUT * time.Second)
 }