280 lines
8.0 KiB
Go
280 lines
8.0 KiB
Go
/**
|
||
* 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
|
||
}
|
||
}
|
||
}
|
||
}
|