eod commit

ws/util.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/json"
 	"errors"
+	"homestead/homestead_gateway/util/cache"
 	"log/slog"
 	"net/http"
 	"time"
@@ -41,6 +42,11 @@ func (wsg *WebsocketGateway) sendHttpError(w http.ResponseWriter, message string
 	_ = json.NewEncoder(w).Encode(map[string]interface{}{"message": message, "code": code})
 }
 
+func (wsg *WebsocketGateway) sendWebsocketPing(conn *websocket.Conn) {
+	_ = conn.SetWriteDeadline(time.Now().Add(5 * time.Second))
+	_ = conn.WriteMessage(websocket.PingMessage, nil)
+}
+
 func (wsg *WebsocketGateway) sendWebsocketError(conn *websocket.Conn, message string, code int) {
 	_ = conn.SetWriteDeadline(time.Now().Add(5 * time.Second))
 	_ = conn.WriteJSON(map[string]interface{}{"message": message, "code": code})
@@ -105,26 +111,20 @@ func loggingMiddleware(logger *slog.Logger, next http.Handler) http.Handler {
 
 // connections
 
-func (wsg *WebsocketGateway) registerConn(c *websocket.Conn, meta connectionMetaData) {
-	wsg.connsMu.Lock()
-	wsg.conns[c] = meta
-	wsg.connsMu.Unlock()
+func (wsg *WebsocketGateway) registerConn(conn *websocket.Conn, meta cache.ConnectionMetaData) {
+	wsg.conns.Set(meta.ID, conn, &meta)
 }
 
-func (wsg *WebsocketGateway) unregisterConn(c *websocket.Conn) {
-	wsg.connsMu.Lock()
-	delete(wsg.conns, c)
-	wsg.connsMu.Unlock()
+func (wsg *WebsocketGateway) unregisterConn(conn *websocket.Conn, meta cache.ConnectionMetaData) {
+	wsg.conns.RemoveById(meta.ID)
+	_ = conn.Close()
 }
 
 func (wsg *WebsocketGateway) closeAll() {
-	wsg.connsMu.Lock()
-	defer wsg.connsMu.Unlock()
-
 	wsg.logger.Info("Closing all websocket connections.")
 
-	for c := range wsg.conns {
+	wsg.conns.Range(func(id string, c *websocket.Conn, meta *cache.ConnectionMetaData) {
 		_ = c.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, "Shutting down."), time.Now().Add(time.Second))
 		_ = c.Close()
-	}
+	})
 }