fix: adding monitoring

internal/hub.go

@@ -1,6 +1,7 @@
 package internal
 
 import (
+	"context"
 	"fmt"
 	"log"
 	"sync"
@@ -17,6 +18,11 @@ const (
 	maxMessageSize = 512
 )
 
+const (
+	minHighWaterMarkMapRebuild = 100
+	mapRebuildThreshold        = 4
+)
+
 type Client struct {
 	ID             string
 	Conn           *websocket.Conn
@@ -36,34 +42,86 @@ func NewClient(conn *websocket.Conn, subscribedPath string) *Client {
 }
 
 type Hub struct {
-	Clients    map[*Client]bool
-	Broadcast  chan []byte
-	Register   chan *Client
-	Unregister chan *Client
+	path          string
+	maxClients    int
+	Clients       map[*Client]bool
+	Broadcast     chan []byte
+	Register      chan *Client
+	Unregister    chan *Client
+	monitor       *MemoryMonitor
+	highWaterMark int
 }
 
-func NewHub() *Hub {
+func NewHub(path string, maxClients int) *Hub {
+	log.Printf("[%s] Hub created with max clients: %d", path, maxClients)
 	return &Hub{
+		path:       path,
+		maxClients: maxClients,
+		monitor:    NewMemoryMonitor(),
 		Broadcast:  make(chan []byte, 256),
-		Register:   make(chan *Client, 10),
-		Unregister: make(chan *Client, 10),
+		Register:   make(chan *Client, maxClients),
+		Unregister: make(chan *Client, maxClients),
 		Clients:    make(map[*Client]bool),
 	}
 }
 
-func (h *Hub) Run() {
+// Run starts the hub event loop. It exits when ctx is cancelled.
+func (h *Hub) Run(ctx context.Context) {
+	monitorTicker := time.NewTicker(MonitorInterval)
 	go func() {
+		defer func() {
+			monitorTicker.Stop()
+			// On shutdown, close every client's Send channel so WritePump sends
+			// a WebSocket close frame and exits. Also expire the read deadline so
+			// blocked ReadMessage() calls in ReadPump return immediately instead
+			// of waiting up to pongWait (60 s).
+			for client := range h.Clients {
+				close(client.Send)
+				client.Conn.SetReadDeadline(time.Now())
+			}
+			log.Printf("[%s] Hub stopped\n", h.path)
+		}()
+
 		for {
 			select {
+			case <-ctx.Done():
+				log.Printf("[%s] Hub shutting down\n", h.path)
+				return
+
 			case client := <-h.Register:
+				if len(h.Clients) >= h.maxClients {
+					close(client.Send)
+					client.Conn.Close()
+					log.Printf("[%s] Rejected client %s (max %d reached)\n", h.path, client.ID, h.maxClients)
+					break
+				}
 				h.Clients[client] = true
-				log.Printf("Client registered %s\n", client.ID)
+				if len(h.Clients) > h.highWaterMark {
+					h.highWaterMark = len(h.Clients)
+				}
+				log.Printf("[%s] Client registered %s\n", h.path, client.ID)
+
 			case client := <-h.Unregister:
 				if _, ok := h.Clients[client]; ok {
 					delete(h.Clients, client)
 					close(client.Send)
+
+					// Rebuild the map when the live set has dropped to less than
+					// 1/mapRebuildThreshold of the peak, so the old backing buckets
+					// are released to the GC.
+					if h.highWaterMark >= minHighWaterMarkMapRebuild &&
+						len(h.Clients) < h.highWaterMark/mapRebuildThreshold {
+						rebuilt := make(map[*Client]bool, len(h.Clients))
+						for c := range h.Clients {
+							rebuilt[c] = true
+						}
+						h.Clients = rebuilt
+						h.highWaterMark = len(h.Clients)
+						log.Printf("[%s] Clients map rebuilt: %d active clients\n", h.path, len(h.Clients))
+					}
+					log.Printf("[%s] Client Unregistered %s\n", h.path, client.ID)
 				}
-				log.Printf("Client Unregistered %s\n", client.ID)
+
 			case message := <-h.Broadcast:
 				for client := range h.Clients {
 					select {
@@ -71,9 +129,18 @@ func (h *Hub) Run() {
 					default:
 						close(client.Send)
 						delete(h.Clients, client)
-						log.Printf("Client %s removed (slow/disconnected)", client.ID)
+						log.Printf("[%s] Client %s removed (slow/disconnected)\n", h.path, client.ID)
 					}
 				}
+
+			case <-monitorTicker.C:
+				current, peak := h.monitor.Snapshot()
+				clientLength := len(h.Clients)
+				if clientLength > 0 {
+					log.Printf("[%s] connected clients: %d | heap alloc: %s | peak heap alloc: %s",
+						h.path, clientLength, FormatBytes(current), FormatBytes(peak),
+					)
+				}
 			}
 		}
 	}()
@@ -82,7 +149,10 @@ func (h *Hub) Run() {
 func WritePump(c *Client, h *Hub) {
 	pingTicker := time.NewTicker(pingPeriod)
 	defer func() {
-		h.Unregister <- c
+		select {
+		case h.Unregister <- c:
+		default:
+		}
 		pingTicker.Stop()
 		c.Conn.Close()
 	}()
@@ -103,7 +173,7 @@ func WritePump(c *Client, h *Hub) {
 			}
 			w.Write(message)
 
-			// Queue queued messages in the same buffer (optional optimization)
+			// Flush any messages that queued up while we were writing.
 			n := len(c.Send)
 			for i := 0; i < n; i++ {
 				w.Write(<-c.Send)
@@ -120,13 +190,15 @@ func WritePump(c *Client, h *Hub) {
 				return
 			}
 		}
-
 	}
 }
 
 func ReadPump(c *Client, h *Hub) {
 	defer func() {
-		h.Unregister <- c
+		select {
+		case h.Unregister <- c:
+		default:
+		}
 		c.Conn.Close()
 	}()