89 lines
2.1 KiB
Go
89 lines
2.1 KiB
Go
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()
|
||
}
|
||
}
|
||
}
|