From f20b730eff2de42ec7350459a4197826448cc596 Mon Sep 17 00:00:00 2001
From: ivynya <ivy@ivy.direct>
Date: Tue, 26 Dec 2023 14:33:44 -0800
Subject: [PATCH] fix: broadcast stats more effectively

---
 server/broadcast.go | 29 +++++++++++++++++++----------
 server/server.go    | 10 +++++-----
 2 files changed, 24 insertions(+), 15 deletions(-)

diff --git a/server/broadcast.go b/server/broadcast.go
index f816e2b..69c035a 100644
--- a/server/broadcast.go
+++ b/server/broadcast.go
@@ -48,25 +48,34 @@ func broadcastToClient(c map[string]*websocket.Conn, req *internal.Request) erro
 	return nil
 }
 
-func broadcastAll(c map[string]*websocket.Conn, req *internal.Request) error {
+// broadcast to all connections and return false if >= 1 failure
+func broadcastAll(c map[string]*websocket.Conn, req *internal.Request) bool {
 	data, _ := json.Marshal(req)
 
+	ok := true
 	for tag, conn := range c {
 		err := conn.WriteMessage(websocket.TextMessage, data)
 		if err != nil {
 			delete(c, tag)
-			return err
+			ok = false
 		}
 	}
-	return nil
+	return ok
 }
 
-func updateConnCount(clientType string, c *map[string]*websocket.Conn) {
-	err := broadcastAll(*c, &internal.Request{
-		Action: clientType,
-		Data:   strconv.Itoa(len(*c)),
-	})
-	if err != nil {
-		updateConnCount(clientType, c)
+// broadcast number of clients and providers to all clients
+func broadcastConnectionStats(clients map[string]*websocket.Conn, providers map[string]*websocket.Conn) {
+	retry_remaining := 3
+	ok := true
+	for ok != true && retry_remaining > 0 {
+		ok = ok && broadcastAll(clients, &internal.Request{
+			Action: "clients",
+			Data:   strconv.Itoa(len(clients)),
+		})
+		ok = ok && broadcastAll(clients, &internal.Request{
+			Action: "providers",
+			Data:   strconv.Itoa(len(providers)),
+		})
+		retry_remaining--
 	}
 }
diff --git a/server/server.go b/server/server.go
index 699791b..f5be1d0 100644
--- a/server/server.go
+++ b/server/server.go
@@ -35,7 +35,7 @@ func main() {
 		// Log join message
 		fmt.Println("Provider joined from " + c.RemoteAddr().String())
 		fmt.Println("Total providers:", len(providers))
-		updateConnCount("providers", &providers)
+		broadcastConnectionStats(clients, providers)
 
 		for {
 			// Read message from provider
@@ -69,7 +69,7 @@ func main() {
 
 		// Unregister provider
 		delete(providers, tag)
-		updateConnCount("providers", &providers)
+		broadcastConnectionStats(clients, providers)
 	}))
 
 	// WebSocket endpoint
@@ -81,10 +81,10 @@ func main() {
 		}
 		clients[tag] = c
 
-		// Log join message
+		// Log join message and broadcast counts
 		fmt.Println("Client joined from " + c.RemoteAddr().String())
 		fmt.Println("Total clients:", len(clients))
-		updateConnCount("clients", &clients)
+		broadcastConnectionStats(clients, providers)
 
 		for {
 			// Read message from client
@@ -118,7 +118,7 @@ func main() {
 
 		// Unregister client
 		delete(clients, tag)
-		updateConnCount("clients", &clients)
+		broadcastConnectionStats(clients, providers)
 	}))
 
 	// Start the server