From c226bdab7fcd010d333e01cd99be0bb8d09492c6 Mon Sep 17 00:00:00 2001
From: savinmax <savinmax@amazon.com>
Date: Sat, 13 Jun 2026 13:26:03 +0200
Subject: [PATCH] feat(hub): update graceful shutdown to iterate rooms for
 multi-room cleanup
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Refactor the stop case in Hub.Run() to iterate h.rooms directly
instead of h.connRoom. For each room, iterate all connections and
send CloseGoingAway frame before closing. After the loop, reset both
maps (h.rooms, h.connRoom) in one shot rather than deleting entries
incrementally. This is cleaner and avoids modifying a map during
iteration.

Add TestIntegration_GracefulShutdownMultiRoom to verify clients in
separate rooms all receive close frames during shutdown.

🤖 Assisted by the code-assist SOP
---
 internal/hub/hub.go                  | 17 +++----
 internal/hub/hub_integration_test.go | 69 ++++++++++++++++++++++++++++
 2 files changed, 76 insertions(+), 10 deletions(-)

diff --git a/internal/hub/hub.go b/internal/hub/hub.go
index 246b585..cc9f4df 100644
--- a/internal/hub/hub.go
+++ b/internal/hub/hub.go
@@ -53,18 +53,15 @@ func (h *Hub) Run() {
 		select {
 		case <-h.stop:
 			h.mu.Lock()
-			for conn, room := range h.connRoom {
-				conn.WriteMessage(websocket.CloseMessage,
-					websocket.FormatCloseMessage(websocket.CloseGoingAway, "server shutting down"))
-				conn.Close()
-				if clients, ok := h.rooms[room]; ok {
-					delete(clients, conn)
-					if len(clients) == 0 {
-						delete(h.rooms, room)
-					}
+			for _, clients := range h.rooms {
+				for conn := range clients {
+					conn.WriteMessage(websocket.CloseMessage,
+						websocket.FormatCloseMessage(websocket.CloseGoingAway, "server shutting down"))
+					conn.Close()
 				}
-				delete(h.connRoom, conn)
 			}
+			h.rooms = make(map[string]map[*websocket.Conn]bool)
+			h.connRoom = make(map[*websocket.Conn]string)
 			h.mu.Unlock()
 			metrics.ConnectedClients.Set(0)
 			h.logger.Info("Hub stopped, all clients disconnected")
diff --git a/internal/hub/hub_integration_test.go b/internal/hub/hub_integration_test.go
index b0bc8f0..077ef58 100644
--- a/internal/hub/hub_integration_test.go
+++ b/internal/hub/hub_integration_test.go
@@ -501,3 +501,72 @@ func TestIntegration_BroadcastToEmptyRoom(t *testing.T) {
 	}
 }
 
+func TestIntegration_GracefulShutdownMultiRoom(t *testing.T) {
+	logger := logging.NewLogger("debug", &bytes.Buffer{})
+	h := New(logger)
+	go h.Run()
+
+	s := httptest.NewServer(http.HandlerFunc(h.HandleWebSocket))
+	defer s.Close()
+
+	// Connect clients to different rooms
+	rooms := []string{"room-a", "room-b", "room-c"}
+	conns := make([]*websocket.Conn, 0, len(rooms))
+	for _, room := range rooms {
+		ws := dialWSWithRoom(t, s, room)
+		conns = append(conns, ws)
+	}
+
+	waitForClients(t, h, 3, time.Second)
+
+	// Verify all clients are connected to separate rooms
+	h.mu.RLock()
+	roomCount := len(h.rooms)
+	connCount := len(h.connRoom)
+	h.mu.RUnlock()
+	if roomCount != 3 {
+		t.Fatalf("expected 3 rooms, got %d", roomCount)
+	}
+	if connCount != 3 {
+		t.Fatalf("expected 3 connections, got %d", connCount)
+	}
+
+	// Shutdown the hub — all clients should receive close frame
+	h.Shutdown()
+
+	var wg sync.WaitGroup
+	for i, ws := range conns {
+		wg.Add(1)
+		go func(idx int, c *websocket.Conn) {
+			defer wg.Done()
+			c.SetReadDeadline(time.Now().Add(2 * time.Second))
+			_, _, err := c.ReadMessage()
+			if err == nil {
+				t.Errorf("client %d: expected error (close frame), got nil", idx)
+				return
+			}
+			closeErr, ok := err.(*websocket.CloseError)
+			if !ok {
+				t.Errorf("client %d: expected CloseError, got %T: %v", idx, err, err)
+				return
+			}
+			if closeErr.Code != websocket.CloseGoingAway {
+				t.Errorf("client %d: expected CloseGoingAway (%d), got %d",
+					idx, websocket.CloseGoingAway, closeErr.Code)
+			}
+		}(i, ws)
+	}
+	wg.Wait()
+
+	// Verify maps are cleared
+	h.mu.RLock()
+	roomsAfter := len(h.rooms)
+	connsAfter := len(h.connRoom)
+	h.mu.RUnlock()
+	if roomsAfter != 0 {
+		t.Errorf("expected 0 rooms after shutdown, got %d", roomsAfter)
+	}
+	if connsAfter != 0 {
+		t.Errorf("expected 0 connections after shutdown, got %d", connsAfter)
+	}
+}