feat(hub): update graceful shutdown to iterate rooms for multi-room cleanup

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
This commit is contained in:
savinmax 2026-06-13 13:26:03 +02:00
parent 5bd08409dc
commit c226bdab7f
2 changed files with 76 additions and 10 deletions

View File

@ -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")

View File

@ -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)
}
}