/** * package manager * 作用:管理全量的 WebSocket 连接,提供线程安全的注册、注销、绑定和消息发送功能。 * 核心:使用 sync.Map 替代传统 map + RWMutex 以提高高并发下的读写性能。 */ package manager import ( "log" "sync" "time" "github.com/gorilla/websocket" ) /** * Client * 结构体:代表一个 WebSocket 客户端连接会话。 */ type Client struct { ID string // 客户端唯一标识 (格式: NodeID-Timestamp) UserID string // 绑定的用户ID (未绑定时为空) RemoteIP string // WebSocket 连接来源 IP Conn *websocket.Conn // 底层 WebSocket 连接对象 SendQueue chan []byte // 发送缓冲队列 (防止网络阻塞导致写协程卡死) closeOnce sync.Once // 保证 SendQueue/Conn 只被关闭一次,防止并发注销时 double-close panic } /** * ClientManager * 结构体:连接管理器,维护 ClientID 映射和 UserID 反向索引。 */ type ClientManager struct { // 客户端列表: map[string]*Client (Key: ClientID) clients sync.Map // 用户反向索引: map[string][]string (Key: UserID, Value: [ClientID1, ClientID2]) userClients sync.Map // userMu 保护 userClients 中切片的读-改-写复合操作。 // 为什么需要:sync.Map 只保证单次 Load/Store 原子,Bind/Remove 是"读出切片→修改→写回", // 并发时会互相覆盖丢失设备,必须用互斥锁串行化。 userMu sync.Mutex } // Manager 全局单例 var Manager = &ClientManager{} /** * Register * 功能:注册新建立的 WebSocket 连接。 * @param c *Client 客户端对象 */ func (m *ClientManager) Register(c *Client) { m.clients.Store(c.ID, c) log.Printf("✅ [ClientManager] 客户端注册: %s", c.ID) // 启动该客户端独享的写协程 go c.writePump() } /** * Unregister * 功能:注销连接,清理资源和索引。 * @param c *Client 客户端对象 */ func (m *ClientManager) Unregister(c *Client) { // LoadAndDelete 原子操作:并发调用时只有一个 goroutine 拿到 ok=true, // 再配合 closeOnce 双保险,彻底避免重复 close(chan) 导致的 panic if _, ok := m.clients.LoadAndDelete(c.ID); ok { c.closeOnce.Do(func() { close(c.SendQueue) // 关闭通道,退出 writePump c.Conn.Close() }) // 读取 c.UserID 与清理反向索引都放在同一把锁内, // 与 BindUser 对 c.UserID 的写入互斥,避免数据竞争 m.userMu.Lock() if c.UserID != "" { m.removeUserClientLocked(c.UserID, c.ID) } m.userMu.Unlock() log.Printf("🔌 [ClientManager] 客户端注销: %s", c.ID) } } /** * GetClientUserID * 功能:在锁保护下读取某连接当前绑定的 UserID。 * 为什么:c.UserID 会被 BindUser 并发写入,裸读构成数据竞争,需与写入同锁。 */ func (m *ClientManager) GetClientUserID(c *Client) string { m.userMu.Lock() defer m.userMu.Unlock() return c.UserID } /** * BindUser * 功能:将 ClientID 与 UserID 进行绑定,支持多端登录。 * @param clientID 客户端ID * @param userID 用户ID */ func (m *ClientManager) BindUser(clientID, userID string) { clientInterface, ok := m.clients.Load(clientID) if !ok { return } client := clientInterface.(*Client) // 加锁保护:既保护 userClients 的"读出切片→追加→写回"复合操作, // 也保护 client.UserID 的写入(Unregister/GetClientUserID 会并发读取该字段) m.userMu.Lock() defer m.userMu.Unlock() // 仅在实际变更时写入 client.UserID: // 1) 重绑到不同用户时,先把 clientID 从旧用户的反向索引里摘掉, // 否则旧用户的 SendToUser 仍会把消息发到这条连接,注销时也只按新 UserID 清理、 // 在旧用户列表里永久残留悬空 clientID; // 2) 重复 self-bind(值相同)时跳过写入,避免冗余写与业务路径的读构成竞争 if client.UserID != userID { if client.UserID != "" { m.removeUserClientLocked(client.UserID, clientID) } client.UserID = userID } actual, _ := m.userClients.LoadOrStore(userID, make([]string, 0)) ids := actual.([]string) // 去重逻辑:防止重复添加同一个 ClientID for _, id := range ids { if id == clientID { return } } // 拷贝后追加,避免直接修改共享切片底层数组引起数据竞争 newIds := make([]string, len(ids), len(ids)+1) copy(newIds, ids) newIds = append(newIds, clientID) m.userClients.Store(userID, newIds) log.Printf("🔗 [ClientManager] 绑定成功: %s -> %s", clientID, userID) } /** * removeUserClientLocked * 功能:从用户的客户端列表中移除指定 ClientID(无锁内核,调用方须已持有 m.userMu)。 * 为什么是"Locked"内核:BindUser/Unregister 已在锁内需要清理旧索引, * 若再走一个自己加锁的版本会重复加锁导致死锁,故统一用这个不加锁的内核。 */ func (m *ClientManager) removeUserClientLocked(userID, clientID string) { val, ok := m.userClients.Load(userID) if !ok { return } ids := val.([]string) newIds := make([]string, 0, len(ids)) for _, id := range ids { if id != clientID { newIds = append(newIds, id) } } if len(newIds) == 0 { m.userClients.Delete(userID) // 如果没有设备了,删除 Key } else { m.userClients.Store(userID, newIds) // 更新列表 } } /** * UserHasClients * 功能:判断某用户在本节点是否仍有在线连接。 * 场景:客户端下线时,用于决定是否需要清理 Redis 中该用户的节点路由 key。 */ func (m *ClientManager) UserHasClients(userID string) bool { if val, ok := m.userClients.Load(userID); ok { return len(val.([]string)) > 0 } return false } /** * SendToClient * 功能:向指定客户端发送消息 (非阻塞模式)。 * @param clientID 目标客户端ID * @param message 消息字节数据 */ func (m *ClientManager) SendToClient(clientID string, message []byte) { if val, ok := m.clients.Load(clientID); ok { client := val.(*Client) select { case client.SendQueue <- message: // 成功入队 default: log.Printf("⚠️ [ClientManager] 发送队列已满,丢弃消息: %s", clientID) } } } /** * SendToUser * 功能:向指定用户的所有在线设备广播消息。 * @param userID 目标用户ID * @param message 消息字节数据 */ func (m *ClientManager) SendToUser(userID string, message []byte) { if val, ok := m.userClients.Load(userID); ok { ids := val.([]string) for _, clientID := range ids { m.SendToClient(clientID, message) } } } /** * GetOtherClients * 功能:获取该用户除 excludeClientID 以外的其他所有在线客户端ID。 * 场景:用于多端同步(如一端接听,通知其他端关闭振铃)。 */ func (m *ClientManager) GetOtherClients(userID, excludeClientID string) []string { var results []string if val, ok := m.userClients.Load(userID); ok { ids := val.([]string) for _, id := range ids { if id != excludeClientID { results = append(results, id) } } } return results } /** * writePump * 功能:每个客户端独享的写协程,确保 WebSocket 并发写安全,并处理心跳 Ping。 */ func (c *Client) writePump() { ticker := time.NewTicker(50 * time.Second) // 心跳间隔 defer func() { ticker.Stop() c.Conn.Close() }() for { select { case message, ok := <-c.SendQueue: c.Conn.SetWriteDeadline(time.Now().Add(10 * time.Second)) if !ok { // 队列关闭,发送 Close 帧 c.Conn.WriteMessage(websocket.CloseMessage, []byte{}) return } w, err := c.Conn.NextWriter(websocket.TextMessage) if err != nil { return } w.Write(message) // 优化:如果队列里堆积了多条消息,一次性写完,减少系统调用次数 n := len(c.SendQueue) for i := 0; i < n; i++ { w.Write([]byte{'\n'}) // JSON消息分隔符 w.Write(<-c.SendQueue) } if err := w.Close(); err != nil { return } case <-ticker.C: // 发送心跳 Ping c.Conn.SetWriteDeadline(time.Now().Add(10 * time.Second)) if err := c.Conn.WriteMessage(websocket.PingMessage, nil); err != nil { return } } } }