From ba2f175db49d7966390e88e135f01752bfa972e0 Mon Sep 17 00:00:00 2001
From: CalebRose <ctrose17@gmail.com>
Date: Mon, 28 Oct 2024 14:41:17 -0700
Subject: [PATCH] ws improvements

---
 main.go         |  5 ++++-
 ws/broadcast.go |  5 ++++-
 ws/websocket.go | 20 ++++++++++++++++----
 3 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/main.go b/main.go
index 1cc425c..51d9656 100644
--- a/main.go
+++ b/main.go
@@ -38,7 +38,10 @@ func monitorDBForUpdates() {
 		currentTS := controller.GetUpdatedTimestamp()
 		if currentTS.UpdatedAt.After(ts.UpdatedAt) {
 			ts = currentTS
-			ws.BroadcastTSUpdate(ts)
+			err := ws.BroadcastTSUpdate(ts)
+			if err != nil {
+				log.Printf("Error broadcasting timestamp: %v", err)
+			}
 		}
 
 		time.Sleep(60 * time.Second)
diff --git a/ws/broadcast.go b/ws/broadcast.go
index bbd2235..3676816 100644
--- a/ws/broadcast.go
+++ b/ws/broadcast.go
@@ -7,7 +7,9 @@ import (
 )
 
 // BroadcastTSUpdate sends the updated timestamp to all connected WebSocket clients
-func BroadcastTSUpdate(ts structs.Timestamp) {
+func BroadcastTSUpdate(ts structs.Timestamp) error {
+	mu.Lock()
+	defer mu.Unlock()
 	for conn := range clients {
 		err := conn.WriteJSON(ts)
 		if err != nil {
@@ -16,4 +18,5 @@ func BroadcastTSUpdate(ts structs.Timestamp) {
 			delete(clients, conn) // Remove client on error
 		}
 	}
+	return nil
 }
diff --git a/ws/websocket.go b/ws/websocket.go
index c860479..6decef4 100644
--- a/ws/websocket.go
+++ b/ws/websocket.go
@@ -3,6 +3,7 @@ package ws
 import (
 	"log"
 	"net/http"
+	"sync"
 
 	"github.com/CalebRose/SimFBA/managers"
 	"github.com/gorilla/websocket"
@@ -19,7 +20,10 @@ var upgrader = websocket.Upgrader{
 }
 
 // Global map to keep track of connected WebSocket clients
-var clients = make(map[*websocket.Conn]bool)
+var (
+	clients = make(map[*websocket.Conn]bool)
+	mu      sync.Mutex
+)
 
 // WebSocketHandler handles WebSocket connection requests
 func WebSocketHandler(w http.ResponseWriter, r *http.Request) {
@@ -32,9 +36,19 @@ func WebSocketHandler(w http.ResponseWriter, r *http.Request) {
 	defer conn.Close()
 
 	// Add new WebSocket connection to the clients map
+	mu.Lock()
 	clients[conn] = true
+	mu.Unlock()
 	log.Println("New WebSocket client connected")
 
+	defer func() {
+		conn.Close()
+		mu.Lock()
+		delete(clients, conn)
+		mu.Unlock()
+		log.Println("WebSocket client disconnected")
+	}()
+
 	ts := managers.GetTimestamp()
 	// Send the latest timestamp to the user immediately upon connection
 	err = conn.WriteJSON(ts)
@@ -47,9 +61,7 @@ func WebSocketHandler(w http.ResponseWriter, r *http.Request) {
 	for {
 		_, msg, err := conn.ReadMessage()
 		if err != nil {
-			log.Println("WebSocket client disconnected:", err)
-			conn.Close()
-			delete(clients, conn)
+			log.Println("WebSocket client error:", err)
 			break
 		}