perf: 优化ws模块协程资源消耗

src/modules/ws/service/ws.go

@@ -12,14 +12,21 @@ type IWS interface {
 	// UpgraderWs http升级ws请求
 	UpgraderWs(w http.ResponseWriter, r *http.Request) *websocket.Conn
 
-	// NewClient 新建客户端
+	// ClientCreate 客户端新建
 	//
 	// uid 登录用户ID
 	// groupIDs 用户订阅组
 	// conn ws连接实例
 	// childConn 子连接实例
-	NewClient(uid string, groupIDs []string, conn *websocket.Conn, childConn any) *model.WSClient
+	ClientCreate(uid string, groupIDs []string, conn *websocket.Conn, childConn any) *model.WSClient
 
-	// CloseClient 关闭客户端
-	CloseClient(clientID string)
+	// ClientClose 客户端关闭
+	ClientClose(clientID string)
+
+	// ClientReadListen 客户端读取消息监听
+	// receiveType 根据接收类型进行消息处理
+	ClientReadListen(wsClient *model.WSClient, receiveType int)
+
+	// ClientWriteListen 客户端写入消息监听
+	ClientWriteListen(wsClient *model.WSClient)
 }

src/modules/ws/service/ws.impl.go

@@ -14,15 +14,12 @@ import (
 )
 
 var (
-	// ws客户端 [clientId: client]
-	wsClients sync.Map
-	// ws用户对应的多个客户端id [uid:clientIds]
-	wsUsers sync.Map
-	// ws组对应的多个用户id [groupID:uids]
-	wsGroup sync.Map
+	wsClients sync.Map // ws客户端 [clientId: client]
+	wsUsers   sync.Map // ws用户对应的多个客户端id [uid:clientIds]
+	wsGroup   sync.Map // ws组对应的多个用户id [groupID:uids]
 )
 
-// 实例化服务层 WSImpl 结构体
+// NewWSImpl 实例化服务层 WSImpl 结构体
 var NewWSImpl = &WSImpl{}
 
 // WSImpl WebSocket通信 服务层处理
@@ -50,13 +47,13 @@ func (s *WSImpl) UpgraderWs(w http.ResponseWriter, r *http.Request) *websocket.C
 	return conn
 }
 
-// NewClient 新建客户端
+// ClientCreate 客户端新建
 //
 // uid 登录用户ID
 // groupIDs 用户订阅组
 // conn ws连接实例
 // childConn 子连接实例
-func (s *WSImpl) NewClient(uid string, groupIDs []string, conn *websocket.Conn, childConn any) *model.WSClient {
+func (s *WSImpl) ClientCreate(uid string, groupIDs []string, conn *websocket.Conn, childConn any) *model.WSClient {
 	// clientID也可以用其他方式生成，只要能保证在所有服务端中都能保证唯一即可
 	clientID := generate.Code(16)
 
@@ -106,76 +103,11 @@ func (s *WSImpl) NewClient(uid string, groupIDs []string, conn *websocket.Conn,
 		}
 	}
 
-	go s.clientRead(wsClient)
-	go s.clientWrite(wsClient)
-
-	// 发客户端id确认是否连接
-	msgByte, _ := json.Marshal(result.OkData(map[string]string{
-		"clientId": clientID,
-	}))
-	wsClient.MsgChan <- msgByte
-
 	return wsClient
 }
 
-// clientRead 客户端读取消息
-func (s *WSImpl) clientRead(wsClient *model.WSClient) {
-	defer func() {
-		if err := recover(); err != nil {
-			logger.Errorf("ws ReadMessage Panic Error: %v", err)
-		}
-	}()
-	for {
-		// 读取消息
-		messageType, msg, err := wsClient.Conn.ReadMessage()
-		if err != nil {
-			logger.Warnf("ws ReadMessage UID %s err: %s", wsClient.BindUid, err.Error())
-			s.CloseClient(wsClient.ID)
-			return
-		}
-		// fmt.Println(messageType, string(msg))
-
-		// 文本和二进制类型，只处理文本json
-		if messageType == websocket.TextMessage {
-			var reqMsg model.WSRequest
-			err := json.Unmarshal(msg, &reqMsg)
-			if err != nil {
-				msgByte, _ := json.Marshal(result.ErrMsg("message format not supported"))
-				wsClient.MsgChan <- msgByte
-			} else {
-				// 协程异步处理
-				go NewWSReceiveImpl.AsyncReceive(wsClient, reqMsg)
-			}
-		}
-	}
-}
-
-// clientWrite 客户端写入消息
-func (s *WSImpl) clientWrite(wsClient *model.WSClient) {
-	defer func() {
-		if err := recover(); err != nil {
-			logger.Errorf("ws WriteMessage Panic Error: %v", err)
-		}
-	}()
-	for msg := range wsClient.MsgChan {
-		// 关闭句柄
-		if string(msg) == "ws:close" {
-			wsClient.Conn.WriteMessage(websocket.CloseMessage, []byte{})
-			return
-		}
-		// 发送消息
-		err := wsClient.Conn.WriteMessage(websocket.TextMessage, msg)
-		if err != nil {
-			logger.Warnf("ws WriteMessage UID %s err: %s", wsClient.BindUid, err.Error())
-			s.CloseClient(wsClient.ID)
-			return
-		}
-		wsClient.LastHeartbeat = time.Now().UnixMilli()
-	}
-}
-
-// CloseClient 客户端关闭
-func (s *WSImpl) CloseClient(clientID string) {
+// ClientClose 客户端关闭
+func (s *WSImpl) ClientClose(clientID string) {
 	v, ok := wsClients.Load(clientID)
 	if !ok {
 		return
@@ -226,3 +158,72 @@ func (s *WSImpl) CloseClient(clientID string) {
 		}
 	}
 }
+
+// ClientReadListen 客户端读取消息监听
+// receiveType 根据接收类型进行消息处理
+func (s *WSImpl) ClientReadListen(wsClient *model.WSClient, receiveType int) {
+	defer func() {
+		if err := recover(); err != nil {
+			logger.Errorf("ws ReadMessage Panic Error: %v", err)
+		}
+	}()
+	for {
+		// 读取消息
+		messageType, msg, err := wsClient.Conn.ReadMessage()
+		if err != nil {
+			logger.Warnf("ws ReadMessage UID %s err: %s", wsClient.BindUid, err.Error())
+			s.ClientClose(wsClient.ID)
+			return
+		}
+		// fmt.Println(messageType, string(msg))
+
+		// 文本 只处理文本json
+		if messageType == websocket.TextMessage {
+			var reqMsg model.WSRequest
+			if err := json.Unmarshal(msg, &reqMsg); err != nil {
+				msgByte, _ := json.Marshal(result.ErrMsg("message format not supported"))
+				wsClient.MsgChan <- msgByte
+				continue
+			}
+			// 接收器处理
+			switch receiveType {
+			case ReceiveCommont:
+				go NewWSReceiveImpl.Commont(wsClient, reqMsg)
+			case ReceiveShell:
+				go NewWSReceiveImpl.Shell(wsClient, reqMsg)
+			case ReceiveShellView:
+				go NewWSReceiveImpl.ShellView(wsClient, reqMsg)
+			}
+		}
+	}
+}
+
+// ClientWriteListen 客户端写入消息监听
+func (s *WSImpl) ClientWriteListen(wsClient *model.WSClient) {
+	defer func() {
+		if err := recover(); err != nil {
+			logger.Errorf("ws WriteMessage Panic Error: %v", err)
+		}
+	}()
+	// 发客户端id确认是否连接
+	msgByte, _ := json.Marshal(result.OkData(map[string]string{
+		"clientId": wsClient.ID,
+	}))
+	wsClient.MsgChan <- msgByte
+	// 消息发送监听
+	for msg := range wsClient.MsgChan {
+		// 关闭句柄
+		if string(msg) == "ws:close" {
+			wsClient.Conn.WriteMessage(websocket.CloseMessage, []byte{})
+			return
+		}
+		// 发送消息
+		err := wsClient.Conn.WriteMessage(websocket.TextMessage, msg)
+		if err != nil {
+			logger.Warnf("ws WriteMessage UID %s err: %s", wsClient.BindUid, err.Error())
+			s.ClientClose(wsClient.ID)
+			return
+		}
+		wsClient.LastHeartbeat = time.Now().UnixMilli()
+	}
+}

src/modules/ws/service/ws_receive.go

@@ -2,8 +2,20 @@ package service
 
 import "be.ems/src/modules/ws/model"
 
+const (
+	ReceiveCommont   = iota // Commont 接收通用业务处理
+	ReceiveShell            // Shell 接收终端交互业务处理
+	ReceiveShellView        // ShellView 接收查看文件终端交互业务处理
+)
+
 // IWSReceive WebSocket消息接收处理 服务层接口
 type IWSReceive interface {
-	// AsyncReceive 接收业务异步处理
-	AsyncReceive(client *model.WSClient, reqMsg model.WSRequest)
+	// Commont 接收通用业务处理
+	Commont(client *model.WSClient, reqMsg model.WSRequest)
+
+	// Shell 接收终端交互业务处理
+	Shell(client *model.WSClient, reqMsg model.WSRequest)
+
+	// ShellView 接收查看文件终端交互业务处理
+	ShellView(client *model.WSClient, reqMsg model.WSRequest)
 }

src/modules/ws/service/ws_receive.impl.go

@@ -20,12 +20,22 @@ var NewWSReceiveImpl = &WSReceiveImpl{}
 // WSReceiveImpl WebSocket消息接收处理 服务层处理
 type WSReceiveImpl struct{}
 
-// AsyncReceive 接收业务异步处理
-func (s *WSReceiveImpl) AsyncReceive(client *model.WSClient, reqMsg model.WSRequest) {
+// Commont 接收通用业务处理
+func (s *WSReceiveImpl) close(client *model.WSClient) {
+	// 主动关闭
+	resultByte, _ := json.Marshal(result.OkMsg("user initiated closure"))
+	client.MsgChan <- resultByte
+	// 等待1s后关闭连接
+	time.Sleep(1 * time.Second)
+	NewWSImpl.ClientClose(client.ID)
+}
+
+// Commont 接收通用业务处理
+func (s *WSReceiveImpl) Commont(client *model.WSClient, reqMsg model.WSRequest) {
 	// 必传requestId确认消息
 	if reqMsg.RequestID == "" {
 		msg := "message requestId is required"
-		logger.Infof("ws AsyncReceive UID %s err: %s", client.BindUid, msg)
+		logger.Infof("ws Commont UID %s err: %s", client.BindUid, msg)
 		msgByte, _ := json.Marshal(result.ErrMsg(msg))
 		client.MsgChan <- msgByte
 		return
@@ -36,34 +46,8 @@ func (s *WSReceiveImpl) AsyncReceive(client *model.WSClient, reqMsg model.WSRequ
 
 	switch reqMsg.Type {
 	case "close":
-		// 主动关闭
-		resultByte, _ := json.Marshal(result.OkMsg("user initiated closure"))
-		client.MsgChan <- resultByte
-		// 等待1s后关闭连接
-		time.Sleep(1 * time.Second)
-		client.StopChan <- struct{}{}
-	case "ssh":
-		// SSH会话消息接收直接写入会话
-		command := reqMsg.Data.(string)
-		sshClientSession := client.ChildConn.(*ssh.SSHClientSession)
-		_, err = sshClientSession.Write(command)
-	case "ssh_resize":
-		// SSH会话窗口重置
-		msgByte, _ := json.Marshal(reqMsg.Data)
-		var data struct {
-			Cols int `json:"cols"`
-			Rows int `json:"rows"`
-		}
-		err = json.Unmarshal(msgByte, &data)
-		if err == nil {
-			sshClientSession := client.ChildConn.(*ssh.SSHClientSession)
-			err = sshClientSession.Session.WindowChange(data.Rows, data.Cols)
-		}
-	case "telnet":
-		// Telnet会话消息接收直接写入会话
-		command := reqMsg.Data.(string)
-		telnetClientSession := client.ChildConn.(*telnet.TelnetClientSession)
-		_, err = telnetClientSession.Write(command)
+		s.close(client)
+		return
 	case "ps":
 		resByte, err = processor.GetProcessData(reqMsg.RequestID, reqMsg.Data)
 	case "net":
@@ -85,7 +69,146 @@ func (s *WSReceiveImpl) AsyncReceive(client *model.WSClient, reqMsg model.WSRequ
 	}
 
 	if err != nil {
-		logger.Warnf("ws AsyncReceive UID %s err: %s", client.BindUid, err.Error())
+		logger.Warnf("ws Commont UID %s err: %s", client.BindUid, err.Error())
+		msgByte, _ := json.Marshal(result.ErrMsg(err.Error()))
+		client.MsgChan <- msgByte
+		return
+	}
+	if len(resByte) > 0 {
+		client.MsgChan <- resByte
+	}
+}
+
+// Shell 接收终端交互业务处理
+func (s *WSReceiveImpl) Shell(client *model.WSClient, reqMsg model.WSRequest) {
+	// 必传requestId确认消息
+	if reqMsg.RequestID == "" {
+		msg := "message requestId is required"
+		logger.Infof("ws Shell UID %s err: %s", client.BindUid, msg)
+		msgByte, _ := json.Marshal(result.ErrMsg(msg))
+		client.MsgChan <- msgByte
+		return
+	}
+
+	var resByte []byte
+	var err error
+
+	switch reqMsg.Type {
+	case "close":
+		s.close(client)
+		return
+	case "ssh":
+		// SSH会话消息接收写入会话
+		if client.ChildConn == nil {
+			err = fmt.Errorf("message type ssh not connected")
+		} else {
+			command := reqMsg.Data.(string)
+			sshClientSession := client.ChildConn.(*ssh.SSHClientSession)
+			_, err = sshClientSession.Write(command)
+		}
+	case "ssh_resize":
+		// SSH会话窗口重置
+		if client.ChildConn == nil {
+			err = fmt.Errorf("message type ssh not connected")
+		} else {
+			msgByte, _ := json.Marshal(reqMsg.Data)
+			var data struct {
+				Cols int `json:"cols"`
+				Rows int `json:"rows"`
+			}
+			err = json.Unmarshal(msgByte, &data)
+			if err == nil {
+				sshClientSession := client.ChildConn.(*ssh.SSHClientSession)
+				err = sshClientSession.Session.WindowChange(data.Rows, data.Cols)
+			}
+		}
+	case "telnet":
+		// Telnet会话消息接收写入会话
+		if client.ChildConn == nil {
+			err = fmt.Errorf("message type telnet not connected")
+		} else {
+			command := reqMsg.Data.(string)
+			telnetClientSession := client.ChildConn.(*telnet.TelnetClientSession)
+			_, err = telnetClientSession.Write(command)
+		}
+	default:
+		err = fmt.Errorf("message type not supported")
+	}
+
+	if err != nil {
+		logger.Warnf("ws Shell UID %s err: %s", client.BindUid, err.Error())
+		msgByte, _ := json.Marshal(result.ErrMsg(err.Error()))
+		client.MsgChan <- msgByte
+		if err == io.EOF {
+			// 等待1s后关闭连接
+			time.Sleep(1 * time.Second)
+			client.StopChan <- struct{}{}
+		}
+		return
+	}
+	if len(resByte) > 0 {
+		client.MsgChan <- resByte
+	}
+}
+
+// ShellView 接收查看文件终端交互业务处理
+func (s *WSReceiveImpl) ShellView(client *model.WSClient, reqMsg model.WSRequest) {
+	// 必传requestId确认消息
+	if reqMsg.RequestID == "" {
+		msg := "message requestId is required"
+		logger.Infof("ws Shell UID %s err: %s", client.BindUid, msg)
+		msgByte, _ := json.Marshal(result.ErrMsg(msg))
+		client.MsgChan <- msgByte
+		return
+	}
+
+	var resByte []byte
+	var err error
+
+	switch reqMsg.Type {
+	case "close":
+		s.close(client)
+		return
+	case "ssh":
+		// SSH会话消息接收写入会话
+		if client.ChildConn == nil {
+			err = fmt.Errorf("message type ssh not connected")
+		} else {
+			command := reqMsg.Data.(string)
+			sshClientSession := client.ChildConn.(*ssh.SSHClientSession)
+			_, err = sshClientSession.Write(command)
+		}
+	case "ssh_resize":
+		// SSH会话窗口重置
+		if client.ChildConn == nil {
+			err = fmt.Errorf("message type ssh not connected")
+		} else {
+			msgByte, _ := json.Marshal(reqMsg.Data)
+			var data struct {
+				Cols int `json:"cols"`
+				Rows int `json:"rows"`
+			}
+			err = json.Unmarshal(msgByte, &data)
+			if err == nil {
+				sshClientSession := client.ChildConn.(*ssh.SSHClientSession)
+				err = sshClientSession.Session.WindowChange(data.Rows, data.Cols)
+			}
+		}
+	case "telnet":
+		// Telnet会话消息接收写入会话
+		if client.ChildConn == nil {
+			err = fmt.Errorf("message type telnet not connected")
+		} else {
+			command := reqMsg.Data.(string)
+			telnetClientSession := client.ChildConn.(*telnet.TelnetClientSession)
+			_, err = telnetClientSession.Write(command)
+		}
+	default:
+		err = fmt.Errorf("message type not supported")
+	}
+
+	if err != nil {
+		logger.Warnf("ws Shell UID %s err: %s", client.BindUid, err.Error())
 		msgByte, _ := json.Marshal(result.ErrMsg(err.Error()))
 		client.MsgChan <- msgByte
 		if err == io.EOF {

src/modules/ws/service/ws_send.impl.go

@@ -46,7 +46,7 @@ func (s *WSSendImpl) ByClientID(clientID string, data any) error {
 
 	client := v.(*model.WSClient)
 	if len(client.MsgChan) > 90 {
-		NewWSImpl.CloseClient(client.ID)
+		NewWSImpl.ClientClose(client.ID)
 		return fmt.Errorf("msg chan over 90 will close client ID: %s", clientID)
 	}
 	client.MsgChan <- dataByte