// Package websocket // @Link https://github.com/xiujiecn/xiu-admin // @Copyright Copyright (c) 2025 LiXiujie // @Author Lxj // @License https://github.com/xiujiecn/xiu-admin/blob/master/LICENSE package websocket import ( "context" "fmt" "runtime/debug" "sync" "github.com/gogf/gf/v2/frame/g" "github.com/gogf/gf/v2/os/gcron" "github.com/gogf/gf/v2/os/gtime" ) // ClientManager 客户端管理 type ClientManager struct { Clients map[*Client]bool // 全部的连接 ClientsLock sync.RWMutex // 读写锁 Users map[string]*Client // 登录的用户 // uuid UserLock sync.RWMutex // 读写锁 Register chan *Client // 连接连接处理 Login chan *login // 用户登录处理 Unregister chan *Client // 断开连接处理程序 Broadcast chan *WResponse // 广播 向全部成员发送数据 ClientBroadcast chan *ClientWResponse // 广播 向某个客户端发送数据 TagBroadcast chan *TagWResponse // 广播 向某个标签成员发送数据 UserBroadcast chan *UserWResponse // 广播 向某个用户的所有链接发送数据 TagCallbackMap map[string]func(client *Client, tag string) // 注册标签回调 closeSignal chan struct{} // 关闭信号 TagCountMap map[string]int // 标签map TagLock sync.RWMutex // 标签读写锁 } func NewClientManager() (clientManager *ClientManager) { clientManager = &ClientManager{ Clients: make(map[*Client]bool), Users: make(map[string]*Client), Register: make(chan *Client, 1000), Unregister: make(chan *Client, 1000), Broadcast: make(chan *WResponse, 1000), TagBroadcast: make(chan *TagWResponse, 1000), UserBroadcast: make(chan *UserWResponse, 1000), TagCallbackMap: make(map[string]func(client *Client, tag string)), closeSignal: make(chan struct{}, 1), TagCountMap: make(map[string]int), TagLock: sync.RWMutex{}, } return } // GetUserKey 获取用户key func GetUserKey(userId uint64) (key string) { key = fmt.Sprintf("%s_%d", "ws", userId) return } // InClient 客户端是否存在 func (manager *ClientManager) InClient(client *Client) (ok bool) { manager.ClientsLock.RLock() defer manager.ClientsLock.RUnlock() _, ok = manager.Clients[client] return } // GetClients 获取所有客户端 func (manager *ClientManager) GetClients() (clients map[*Client]bool) { clients = make(map[*Client]bool) manager.ClientsRange(func(client *Client, value bool) (result bool) { clients[client] = value return true }) return } // ClientsRange 遍历 func (manager *ClientManager) ClientsRange(f func(client *Client, value bool) bool) { manager.ClientsLock.RLock() defer func() { manager.ClientsLock.RUnlock() }() for key, value := range manager.Clients { r := f(key, value) if !r { break } } } // GetClientsLen 获取客户端总数 func (manager *ClientManager) GetClientsLen() (clientsLen int) { clientsLen = len(manager.Clients) return } // AddClients 添加客户端 func (manager *ClientManager) AddClients(client *Client) { manager.ClientsLock.Lock() defer manager.ClientsLock.Unlock() manager.Clients[client] = true } // DelClients 删除客户端 func (manager *ClientManager) DelClients(client *Client) { manager.ClientsLock.Lock() defer manager.ClientsLock.Unlock() if _, ok := manager.Clients[client]; ok { // 删除标签 client.tags.Iterator(func(k int, name string) bool { DelTagCount(name) return true }) delete(manager.Clients, client) } } // GetUserClient 获取用户的连接 func (manager *ClientManager) GetUserClient(userId uint64) (client *Client) { manager.UserLock.RLock() defer manager.UserLock.RUnlock() userKey := GetUserKey(userId) if value, ok := manager.Users[userKey]; ok { client = value } return } // AddUsers 添加用户 func (manager *ClientManager) AddUsers(key string, client *Client) { manager.UserLock.Lock() defer manager.UserLock.Unlock() manager.Users[key] = client } // DelUsers 删除用户 func (manager *ClientManager) DelUsers(client *Client) (result bool) { manager.UserLock.Lock() defer manager.UserLock.Unlock() key := GetUserKey(client.UserId) if value, ok := manager.Users[key]; ok { // 判断是否为相同的用户 if value.Addr != client.Addr { return } delete(manager.Users, key) result = true } return } // GetUsersLen 已登录用户数 func (manager *ClientManager) GetUsersLen() (userLen int) { userLen = len(manager.Users) return } // EventRegister 用户建立连接事件 func (manager *ClientManager) EventRegister(client *Client) { manager.AddClients(client) //发送当前客户端标识 client.SendMsg(&WResponse{Event: "connected", Data: g.Map{ "ID": client.ID, }}) if client.UserId > 0 { manager.AddUsers(GetUserKey(client.UserId), client) } g.Log().Info(client.Ctx, "websocket EventRegister.", client.UserId) } // EventLogin 用户登录事件 func (manager *ClientManager) EventLogin(login *login) { client := login.Client if manager.InClient(client) { userKey := login.GetKey() manager.AddUsers(userKey, login.Client) } } // EventUnregister 用户断开连接事件 func (manager *ClientManager) EventUnregister(client *Client) { g.Log().Info(client.Ctx, "websocket EventUnregister.", client.UserId) manager.DelClients(client) // 删除用户连接 deleteResult := manager.DelUsers(client) if !deleteResult { // 不是当前连接的客户端 return } // 关闭 chan // close(client.Send) client.close() } // ClearTimeoutConnections 定时清理超时连接 func (manager *ClientManager) clearTimeoutConnections() { currentTime := uint64(gtime.Now().Unix()) clients := clientManager.GetClients() for client := range clients { if client.IsHeartbeatTimeout(currentTime) { //fmt.Println("心跳时间超时 关闭连接", client.Addr, client.UserId, client.LoginTime, client.HeartbeatTime) _ = client.Socket.Close() } } } // WebsocketPing 心跳处理 func (manager *ClientManager) ping(ctx context.Context) { //定时任务,发送心跳包 _, _ = gcron.Add(ctx, "0 */1 * * * *", func(ctx context.Context) { res := &WResponse{ Event: Ping, Data: g.Map{}, } SendToAll(res) }) // 定时任务,清理超时连接 _, _ = gcron.Add(ctx, "*/30 * * * * *", func(ctx context.Context) { manager.clearTimeoutConnections() }) } // 管道处理程序 func (manager *ClientManager) start() { defer func() { if r := recover(); r != nil { g.Log().Warningf(mctx, "websocket start recover:%+v, stack:%+v", r, string(debug.Stack())) return } }() for { select { case conn := <-manager.Register: // 建立连接事件 manager.EventRegister(conn) case login := <-manager.Login: // 用户登录 manager.EventLogin(login) case conn := <-manager.Unregister: // 断开连接事件 manager.EventUnregister(conn) case message := <-manager.Broadcast: // 全部客户端广播事件 clients := manager.GetClients() for conn := range clients { conn.SendMsg(message) } case message := <-manager.TagBroadcast: // 标签广播事件 clients := manager.GetClients() // find := false for conn := range clients { if conn.tags.Contains(message.Tag) { conn.SendMsg(message.WResponse) // find = true } } // g.Log().Debugf(context.TODO(), "client_manager.go 标签广播事件 find:%v msg:%v", find, message) case message := <-manager.UserBroadcast: // 用户广播事件 clients := manager.GetClients() for conn := range clients { if conn.UserId == message.UserID { conn.SendMsg(message.WResponse) } } case message := <-manager.ClientBroadcast: // 单个客户端广播事件 clients := manager.GetClients() for conn := range clients { if conn.ID == message.ID { conn.SendMsg(message.WResponse) } } case <-manager.closeSignal: g.Log().Debug(mctx, "websocket closeSignal exit.") return } } } // SendToAll 发送全部客户端 func SendToAll(response *WResponse) { clientManager.Broadcast <- response } // SendToClientID 发送单个客户端 func SendToClientID(id string, response *WResponse) { clientRes := &ClientWResponse{ ID: id, WResponse: response, } clientManager.ClientBroadcast <- clientRes } // SendToUser 发送单个用户 func SendToUser(userID uint64, response *WResponse) { userRes := &UserWResponse{ UserID: userID, WResponse: response, } clientManager.UserBroadcast <- userRes // g.Log().Debugf(context.Background(), "websocket SendToUser userID:%v response:%+v", userID, response) } // SendToTag 发送某个标签 func SendToTag(tag string, response *WResponse) { count := GetTagCount(tag) if count == 0 { return } tagRes := &TagWResponse{ Tag: tag, WResponse: response, } clientManager.TagBroadcast <- tagRes } // RegisterTagCallback 注册标签回调 func RegisterTagCallback(tag string, callback func(client *Client, tag string)) { clientManager.TagLock.Lock() defer clientManager.TagLock.Unlock() clientManager.TagCallbackMap[tag] = callback } // UnregisterTagCallback 注销标签回调 func UnregisterTagCallback(tag string) { clientManager.TagLock.Lock() defer clientManager.TagLock.Unlock() delete(clientManager.TagCallbackMap, tag) } // 添加标签订阅计数 func AddTagCount(tag string) { clientManager.TagLock.Lock() defer clientManager.TagLock.Unlock() if _, ok := clientManager.TagCountMap[tag]; !ok { clientManager.TagCountMap[tag] = 1 } else { clientManager.TagCountMap[tag]++ } } // 删除标签订阅计数 func DelTagCount(tag string) { clientManager.TagLock.Lock() defer clientManager.TagLock.Unlock() count, ok := clientManager.TagCountMap[tag] if ok { if count > 0 { clientManager.TagCountMap[tag]-- } else { delete(clientManager.TagCountMap, tag) } } } // GetTagCount 获取标签订阅计数 func GetTagCount(tag string) (count int) { clientManager.TagLock.RLock() defer clientManager.TagLock.RUnlock() count, ok := clientManager.TagCountMap[tag] if ok { return count } return 0 } // 变量Tag Count func RangeTagCount(callback func(tag string, count int)) { clientManager.TagLock.RLock() defer clientManager.TagLock.RUnlock() for tag, count := range clientManager.TagCountMap { callback(tag, count) } }