增加关闭hub

internal/websocket/hub.go

@@ -46,6 +46,9 @@ type Hub struct {
 
 	// deviceRepo 设备仓库
 	deviceRepo repository.DeviceRepo
+
+	// 关闭消息
+	close chan struct{}
 }
 
 // Client WebSocket客户端结构
@@ -78,6 +81,7 @@ func NewHub(deviceRepo repository.DeviceRepo) *Hub {
 		deviceClients: make(map[string]*Client),
 		logger:        logs.NewLogger(),
 		deviceRepo:    deviceRepo,
+		close:         make(chan struct{}),
 	}
 }
 
@@ -101,10 +105,20 @@ func (h *Hub) Run() {
 			h.unregisterClient(client)
 		case message := <-h.broadcast:
 			h.broadcastMessage(message)
+		case <-h.close:
+			return
 		}
 	}
 }
 
+func (h *Hub) Close() {
+	// 关闭时清理所有资源
+	for client := range h.clients {
+		h.unregisterClient(client)
+	}
+	close(h.close)
+}
+
 // registerClient 注册客户端
 func (h *Hub) registerClient(client *Client) {
 	h.mutex.Lock()

internal/websocket/server.go

@@ -75,6 +75,10 @@ func (s *Server) Start() {
 	go s.hub.Run()
 }
 
+func (s *Server) Stop() {
+	s.hub.Close()
+}
+
 // readPump 从WebSocket连接读取消息
 func (c *Client) readPump() {
 	defer func() {