init

ws/hub.go

@@ -0,0 +1,88 @@
+package ws
+
+import (
+	"log"
+	"sync"
+)
+
+// Hub 维护活动客户端的集合并广播消息
+type Hub struct {
+	// 按 IMEI 分组的客户端连接
+	Clients map[string]map[*Client]bool
+
+	// 广播消息的通道（添加 IMEI）
+	Broadcast chan *WsMessage
+
+	// 注册请求
+	Register chan *Client
+
+	// 注销请求
+	Unregister chan *Client
+
+	// 互斥锁保护 clients map
+	mu sync.RWMutex
+}
+
+// Message WebSocket 消息结构
+type WsMessage struct {
+	IMEI string `json:"imei"`
+	Data string `json:"data"`
+}
+
+func NewHub() *Hub {
+	return &Hub{
+		Broadcast:  make(chan *WsMessage),
+		Register:   make(chan *Client),
+		Unregister: make(chan *Client),
+		Clients:    make(map[string]map[*Client]bool),
+	}
+}
+
+func (h *Hub) Run() {
+	for {
+		select {
+		case client := <-h.Register:
+			h.mu.Lock()
+			if _, ok := h.Clients[client.Imei]; !ok {
+				h.Clients[client.Imei] = make(map[*Client]bool)
+			}
+			h.Clients[client.Imei][client] = true
+			log.Printf("客户端注册: IMEI=%s", client.Imei)
+			h.mu.Unlock()
+
+		case client := <-h.Unregister:
+			h.mu.Lock()
+			if _, ok := h.Clients[client.Imei]; ok {
+				if _, ok := h.Clients[client.Imei][client]; ok {
+					delete(h.Clients[client.Imei], client)
+					close(client.Send)
+					if len(h.Clients[client.Imei]) == 0 {
+						delete(h.Clients, client.Imei)
+					}
+					log.Printf("客户端注销: IMEI=%s", client.Imei)
+				}
+			}
+			h.mu.Unlock()
+
+		case broadcast := <-h.Broadcast:
+			h.mu.RLock()
+			log.Printf("收到广播消息: IMEI=%s, Data=%s", broadcast.IMEI, broadcast.Data)
+			if clients, ok := h.Clients[broadcast.IMEI]; ok {
+				log.Printf("找到目标客户端组: IMEI=%s, 客户端数=%d", broadcast.IMEI, len(clients))
+				for client := range clients {
+					select {
+					case client.Send <- broadcast:
+						log.Printf("消息已发送到客户端: IMEI=%s", client.Imei)
+					default:
+						log.Printf("发送失败，关闭客户端: IMEI=%s", client.Imei)
+						close(client.Send)
+						delete(clients, client)
+					}
+				}
+			} else {
+				log.Printf("未找到目标客户端组: IMEI=%s", broadcast.IMEI)
+			}
+			h.mu.RUnlock()
+		}
+	}
+}