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:
parent
5bd08409dc
commit
c226bdab7f
@ -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")
|
||||
|
||||
@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user