iotlcm/app/server/internal/library/websocket/client_manager.go

// Package websocket
// @Link  https://github.com/xiujiecn/xiu-admin
// @Copyright  Copyright (c) 2025 LiXiujie
// @Author  Lxj <li@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)
	}
}