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

This commit is contained in:
TsMask
2024-08-07 15:27:36 +08:00
parent 712a9fee0b
commit 0f98508169
7 changed files with 376 additions and 145 deletions

View File

@@ -20,14 +20,15 @@ import (
"github.com/gin-gonic/gin"
)
// 实例化控制层 WSController 结构体
// NewWSController 实例化控制层 WSController 结构体
var NewWSController = &WSController{
wsService: service.NewWSImpl,
wsSendService: service.NewWSSendImpl,
neHostService: neService.NewNeHostImpl,
neInfoService: neService.NewNeInfoImpl,
}
// WebSocket通信
// WSController WebSocket通信
//
// PATH /ws
type WSController struct {
@@ -37,9 +38,11 @@ type WSController struct {
wsSendService service.IWSSend
// 网元主机连接服务
neHostService neService.INeHost
// 网元信息服务
neInfoService neService.INeInfo
}
// 通用
// WS 通用
//
// GET /?subGroupIDs=0
func (s *WSController) WS(c *gin.Context) {
@@ -71,17 +74,19 @@ func (s *WSController) WS(c *gin.Context) {
}
defer conn.Close()
wsClient := s.wsService.NewClient(loginUser.UserID, subGroupIDs, conn, nil)
wsClient := s.wsService.ClientCreate(loginUser.UserID, subGroupIDs, conn, nil)
go s.wsService.ClientWriteListen(wsClient)
go s.wsService.ClientReadListen(wsClient, service.ReceiveCommont)
// 等待停止信号
for value := range wsClient.StopChan {
s.wsService.CloseClient(wsClient.ID)
s.wsService.ClientClose(wsClient.ID)
logger.Infof("ws Stop Client UID %s %s", wsClient.BindUid, value)
return
}
}
// 测试
// Test 测试
//
// GET /test?clientId=&groupID=
func (s *WSController) Test(c *gin.Context) {
@@ -115,7 +120,7 @@ func (s *WSController) Test(c *gin.Context) {
c.JSON(200, result.OkData(errMsgArr))
}
// SSH终端
// SSH 终端
//
// GET /ssh?hostId=1&cols=80&rows=40
func (s *WSController) SSH(c *gin.Context) {
@@ -185,13 +190,16 @@ func (s *WSController) SSH(c *gin.Context) {
}
defer wsConn.Close()
wsClient := s.wsService.NewClient(loginUser.UserID, nil, wsConn, clientSession)
wsClient := s.wsService.ClientCreate(loginUser.UserID, nil, wsConn, clientSession)
go s.wsService.ClientWriteListen(wsClient)
go s.wsService.ClientReadListen(wsClient, service.ReceiveShell)
// 实时读取SSH消息直接输出
msTicker := time.NewTicker(100 * time.Millisecond)
defer msTicker.Stop()
go func() {
for ms := range msTicker.C {
for {
select {
case ms := <-msTicker.C:
outputByte := clientSession.Read()
if len(outputByte) > 0 {
outputStr := string(outputByte)
@@ -208,18 +216,15 @@ func (s *WSController) SSH(c *gin.Context) {
// return
// }
}
case <-wsClient.StopChan: // 等待停止信号
s.wsService.ClientClose(wsClient.ID)
logger.Infof("ws Stop Client UID %s", wsClient.BindUid)
return
}
}()
// 等待停止信号
for value := range wsClient.StopChan {
s.wsService.CloseClient(wsClient.ID)
logger.Infof("ws Stop Client UID %s %s", wsClient.BindUid, value)
return
}
}
// Telnet终端
// Telnet 终端
//
// GET /telnet?hostId=1
func (s *WSController) Telnet(c *gin.Context) {
@@ -283,13 +288,16 @@ func (s *WSController) Telnet(c *gin.Context) {
}
defer wsConn.Close()
wsClient := s.wsService.NewClient(loginUser.UserID, nil, wsConn, clientSession)
wsClient := s.wsService.ClientCreate(loginUser.UserID, nil, wsConn, clientSession)
go s.wsService.ClientWriteListen(wsClient)
go s.wsService.ClientReadListen(wsClient, service.ReceiveShell)
// 实时读取Telnet消息直接输出
msTicker := time.NewTicker(100 * time.Millisecond)
defer msTicker.Stop()
go func() {
for ms := range msTicker.C {
for {
select {
case ms := <-msTicker.C:
outputByte := clientSession.Read()
if len(outputByte) > 0 {
outputStr := strings.TrimRight(string(outputByte), "\x00")
@@ -306,13 +314,88 @@ func (s *WSController) Telnet(c *gin.Context) {
// return
// }
}
case <-wsClient.StopChan: // 等待停止信号
s.wsService.ClientClose(wsClient.ID)
logger.Infof("ws Stop Client UID %s", wsClient.BindUid)
return
}
}
}
// ShellView 终端交互式文件内容查看
//
// GET /view
func (s *WSController) ShellView(c *gin.Context) {
language := ctx.AcceptLanguage(c)
var query struct {
NeType string `form:"neType" binding:"required"`
NeId string `form:"neId" binding:"required"`
Cols int `form:"cols"` // 终端单行字符数
Rows int `form:"rows"` // 终端显示行数
}
if err := c.ShouldBindQuery(&query); err != nil {
c.JSON(400, result.CodeMsg(400, i18n.TKey(language, "app.common.err400")))
return
}
// 登录用户信息
loginUser, err := ctx.LoginUser(c)
if err != nil {
c.JSON(401, result.CodeMsg(401, i18n.TKey(language, err.Error())))
return
}
if query.Cols < 120 || query.Cols > 400 {
query.Cols = 120
}
if query.Rows < 40 || query.Rows > 1200 {
query.Rows = 40
}
// 网元主机的SSH客户端
sshClient, err := s.neInfoService.NeRunSSHClient(query.NeType, query.NeId)
if err != nil {
c.JSON(200, result.ErrMsg(err.Error()))
return
}
defer sshClient.Close()
// ssh连接会话
clientSession, err := sshClient.NewClientSession(query.Cols, query.Rows)
if err != nil {
c.JSON(200, result.ErrMsg("neinfo ssh client session new err"))
return
}
defer clientSession.Close()
// 将 HTTP 连接升级为 WebSocket 连接
wsConn := s.wsService.UpgraderWs(c.Writer, c.Request)
if wsConn == nil {
return
}
defer wsConn.Close()
wsClient := s.wsService.ClientCreate(loginUser.UserID, nil, wsConn, clientSession)
go s.wsService.ClientWriteListen(wsClient)
go s.wsService.ClientReadListen(wsClient, service.ReceiveShellView)
// 实时读取SSH消息直接输出
msTicker := time.NewTicker(100 * time.Millisecond)
defer msTicker.Stop()
for {
select {
case ms := <-msTicker.C:
outputByte := clientSession.Read()
if len(outputByte) > 0 {
outputStr := string(outputByte)
msgByte, _ := json.Marshal(result.Ok(map[string]any{
"requestId": fmt.Sprintf("view_%d", ms.UnixMilli()),
"data": outputStr,
}))
wsClient.MsgChan <- msgByte
}
case <-wsClient.StopChan: // 等待停止信号
s.wsService.ClientClose(wsClient.ID)
logger.Infof("ws Stop Client UID %s", wsClient.BindUid)
return
}
}()
// 等待停止信号
for value := range wsClient.StopChan {
s.wsService.CloseClient(wsClient.ID)
logger.Infof("ws Stop Client UID %s %s", wsClient.BindUid, value)
return
}
}