refactor: websocket支持单用户多连接

server/pkg/ws/client_manager.go

@@ -9,10 +9,29 @@ import (
 // 心跳间隔
 const heartbeatInterval = 25 * time.Second
 
+// 单个用户的全部的连接, key->clientId, value->Client
+type UserClients map[string]*Client
+
+func (ucs UserClients) GetByCid(clientId string) *Client {
+	return ucs[clientId]
+}
+
+func (ucs UserClients) AddClient(client *Client) {
+	ucs[client.ClientId] = client
+}
+
+func (ucs UserClients) DeleteByCid(clientId string) {
+	delete(ucs, clientId)
+}
+
+func (ucs UserClients) Count() int {
+	return len(ucs)
+}
+
 // 连接管理
 type ClientManager struct {
-	ClientMap map[string]*Client // 全部的连接, key->token, value->&client
-	RwLock    sync.RWMutex       // 读写锁
+	UserClientsMap map[UserId]UserClients // 全部的用户连接, key->userid, value->UserClients
+	RwLock         sync.RWMutex           // 读写锁
 
 	ConnectChan    chan *Client // 连接处理
 	DisConnectChan chan *Client // 断开连接处理
@@ -21,7 +40,7 @@ type ClientManager struct {
 
 func NewClientManager() (clientManager *ClientManager) {
 	return &ClientManager{
-		ClientMap:      make(map[string]*Client),
+		UserClientsMap: make(map[UserId]UserClients),
 		ConnectChan:    make(chan *Client, 10),
 		DisConnectChan: make(chan *Client, 10),
 		MsgChan:        make(chan *Msg, 100),
@@ -58,47 +77,45 @@ func (manager *ClientManager) CloseClient(client *Client) {
 }
 
 // 根据用户id关闭客户端连接
-func (manager *ClientManager) CloseByClientUuid(clientUuid string) {
-	manager.CloseClient(manager.GetByClientUuid(clientUuid))
+func (manager *ClientManager) CloseByUid(userId UserId) {
+	for _, client := range manager.GetByUid(userId) {
+		manager.CloseClient(client)
+	}
 }
 
 // 获取所有的客户端
-func (manager *ClientManager) AllClient() map[string]*Client {
+func (manager *ClientManager) AllUserClient() map[UserId]UserClients {
 	manager.RwLock.RLock()
 	defer manager.RwLock.RUnlock()
 
-	return manager.ClientMap
+	return manager.UserClientsMap
 }
 
-// 通过userId获取
-func (manager *ClientManager) GetByUid(userId UserId) *Client {
+// 通过userId获取用户所有客户端信息
+func (manager *ClientManager) GetByUid(userId UserId) UserClients {
 	manager.RwLock.RLock()
 	defer manager.RwLock.RUnlock()
-	for _, client := range manager.ClientMap {
-		if userId == client.UserId {
-			return client
-		}
+	return manager.UserClientsMap[userId]
+}
+
+// 通过userId和clientId获取客户端信息
+func (manager *ClientManager) GetByUidAndCid(uid UserId, clientId string) *Client {
+	if clients := manager.GetByUid(uid); clients != nil {
+		return clients.GetByCid(clientId)
 	}
 	return nil
 }
 
-// 通过userId获取
-func (manager *ClientManager) GetByClientUuid(uuid string) *Client {
-	manager.RwLock.RLock()
-	defer manager.RwLock.RUnlock()
-	return manager.ClientMap[uuid]
-}
-
 // 客户端数量
 func (manager *ClientManager) Count() int {
 	manager.RwLock.RLock()
 	defer manager.RwLock.RUnlock()
-	return len(manager.ClientMap)
+	return len(manager.UserClientsMap)
 }
 
 // 发送json数据给指定用户
-func (manager *ClientManager) SendJsonMsg(clientUuid string, data any) {
-	manager.MsgChan <- &Msg{ToClientUuid: clientUuid, Data: data, Type: JsonMsg}
+func (manager *ClientManager) SendJsonMsg(userId UserId, clientId string, data any) {
+	manager.MsgChan <- &Msg{ToUserId: userId, ToClientId: clientId, Data: data, Type: JsonMsg}
 }
 
 // 监听并发送给客户端信息
@@ -106,10 +123,22 @@ func (manager *ClientManager) WriteMessage() {
 	go func() {
 		for {
 			msg := <-manager.MsgChan
-			if cli := manager.GetByClientUuid(msg.ToClientUuid); cli != nil {
-				if err := cli.WriteMsg(msg); err != nil {
-					manager.CloseClient(cli)
+			uid := msg.ToUserId
+			cid := msg.ToClientId
+			// 客户端id不为空，则向指定客户端发送消息即可
+			if cid != "" {
+				cli := manager.GetByUidAndCid(uid, cid)
+				if cli != nil {
+					cli.WriteMsg(msg)
+				} else {
+					logx.Warnf("[uid=%v, cid=%s]的ws连接不存在", uid, cid)
 				}
+				continue
+			}
+
+			// cid为空，则向该用户所有客户端发送该消息
+			for _, cli := range manager.GetByUid(uid) {
+				cli.WriteMsg(msg)
 			}
 		}
 	}()
@@ -123,15 +152,17 @@ func (manager *ClientManager) HeartbeatTimer() {
 		for {
 			<-ticker.C
 			//发送心跳
-			for userId, cli := range manager.AllClient() {
-				if cli == nil || cli.WsConn == nil {
-					continue
-				}
-				if err := cli.Ping(); err != nil {
-					manager.CloseClient(cli)
-					logx.Debugf("WS发送心跳失败: %v 总连接数：%d", userId, Manager.Count())
-				} else {
-					logx.Debugf("WS发送心跳成功: uid=%v", userId)
+			for userId, clis := range manager.AllUserClient() {
+				for _, cli := range clis {
+					if cli == nil || cli.WsConn == nil {
+						continue
+					}
+					if err := cli.Ping(); err != nil {
+						manager.CloseClient(cli)
+						logx.Debugf("WS发送心跳失败: uid=%v, cid=%s, usercount=%d", userId, cli.ClientId, Manager.Count())
+					} else {
+						logx.Debugf("WS发送心跳成功: uid=%v, cid=%s", userId, cli.ClientId)
+					}
 				}
 			}
 		}
@@ -141,12 +172,12 @@ func (manager *ClientManager) HeartbeatTimer() {
 
 // 处理建立连接
 func (manager *ClientManager) doConnect(client *Client) {
-	cli := manager.GetByClientUuid(client.ClientUuid)
+	cli := manager.GetByUidAndCid(client.UserId, client.ClientId)
 	if cli != nil {
 		manager.doDisconnect(cli)
 	}
-	manager.addClient2Map(client)
-	logx.Debugf("WS客户端已连接: uid=%d, count=%d", client.UserId, manager.Count())
+	manager.addUserClient2Map(client)
+	logx.Debugf("WS客户端已连接: uid=%d, cid=%s, usercount=%d", client.UserId, client.ClientId, manager.Count())
 }
 
 // 处理断开连接
@@ -156,18 +187,32 @@ func (manager *ClientManager) doDisconnect(client *Client) {
 		_ = client.WsConn.Close()
 		client.WsConn = nil
 	}
-	manager.delClient4Map(client)
-	logx.Debugf("WS客户端已断开: uid=%d, count=%d", client.UserId, Manager.Count())
+	manager.delUserClient4Map(client)
+	logx.Debugf("WS客户端已断开: uid=%d, cid=%s, usercount=%d", client.UserId, client.ClientId, Manager.Count())
 }
 
-func (manager *ClientManager) addClient2Map(client *Client) {
+func (manager *ClientManager) addUserClient2Map(client *Client) {
 	manager.RwLock.Lock()
 	defer manager.RwLock.Unlock()
-	manager.ClientMap[client.ClientUuid] = client
+
+	userClients := manager.UserClientsMap[client.UserId]
+	if userClients == nil {
+		userClients = make(UserClients)
+		manager.UserClientsMap[client.UserId] = userClients
+	}
+	userClients.AddClient(client)
 }
 
-func (manager *ClientManager) delClient4Map(client *Client) {
+func (manager *ClientManager) delUserClient4Map(client *Client) {
 	manager.RwLock.Lock()
 	defer manager.RwLock.Unlock()
-	delete(manager.ClientMap, client.ClientUuid)
+
+	userClients := manager.UserClientsMap[client.UserId]
+	if userClients != nil {
+		userClients.DeleteByCid(client.ClientId)
+		// 如果用户所有客户端都关闭，则移除manager中的UserClientsMap值
+		if userClients.Count() == 0 {
+			delete(manager.UserClientsMap, client.UserId)
+		}
+	}
 }