chore: 并入 iotlcm-admin-svr 历史到 admin-svr/（原提交 49818c6）
This commit is contained in:
MiFaZhan
@@ -0,0 +1,384 @@
|
||||
// Package websocket
|
||||
// @Link https://github.com/xiujiecn/xiu-admin
|
||||
// @Copyright Copyright (c) 2025 LiXiujie
|
||||
// @Author Lxj <xxx@example.com>
|
||||
// @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)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user