TRKJ/mayfly-go
2
0
Fork 0
mirror of https://gitee.com/dromara/mayfly-go synced 2025-11-04 08:20:25 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor: 系统websocket消息重构

Browse Source
This commit is contained in:
meilin.huang
2023-09-12 20:54:07 +08:00
parent e4e68d02bc
commit 4b973b22a4
10 changed files with 303 additions and 75 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
mayfly_go_web/src/views/ops/machine/file/MachineFile.vue
View File

@@ -190,7 +190,7 @@
<template #reference>
<SvgIcon name="QuestionFilled" :size="18" class="pointer-icon mr10" />
</template>
<div>rename: 双击文件名单元格后回车</div>
<div>rename: 双击文件名单元格修改后回车</div>
</el-popover>
操作
</template>

2
server/internal/auth/api/account_login.go
View File

@@ -17,6 +17,7 @@ import (
"mayfly-go/pkg/req"
"mayfly-go/pkg/utils/cryptox"
"mayfly-go/pkg/utils/jsonx"
"mayfly-go/pkg/ws"
"strconv"
"time"
)
@@ -118,4 +119,5 @@ func (a *AccountLogin) OtpVerify(rc *req.Ctx) {
func (a *AccountLogin) Logout(rc *req.Ctx) {
req.GetPermissionCodeRegistery().Remove(rc.LoginAccount.Id)
ws.CloseClient(rc.LoginAccount.Id)
}

10
server/internal/db/api/db.go
View File

@@ -182,7 +182,7 @@ func (d *Db) ExecSqlFile(rc *req.Ctx) {
errInfo = t.Error()
}
if len(errInfo) > 0 {
d.MsgApp.CreateAndSend(rc.LoginAccount, ws.ErrMsg("sql脚本执行失败", fmt.Sprintf("[%s]%s执行失败: [%s]", filename, dbConn.Info.GetLogDesc(), errInfo)))
d.MsgApp.CreateAndSend(rc.LoginAccount, ws.ErrSysMsg("sql脚本执行失败", fmt.Sprintf("[%s]%s执行失败: [%s]", filename, dbConn.Info.GetLogDesc(), errInfo)))
}
}
}()
@@ -202,7 +202,7 @@ func (d *Db) ExecSqlFile(rc *req.Ctx) {
break
}
if err != nil {
d.MsgApp.CreateAndSend(rc.LoginAccount, ws.ErrMsg("sql脚本执行失败", fmt.Sprintf("[%s][%s] 解析SQL失败: [%s]", filename, dbConn.Info.GetLogDesc(), err.Error())))
d.MsgApp.CreateAndSend(rc.LoginAccount, ws.ErrSysMsg("sql脚本执行失败", fmt.Sprintf("[%s][%s] 解析SQL失败: [%s]", filename, dbConn.Info.GetLogDesc(), err.Error())))
return
}
sql := sqlparser.String(stmt)
@@ -215,11 +215,11 @@ func (d *Db) ExecSqlFile(rc *req.Ctx) {
}
if err != nil {
d.MsgApp.CreateAndSend(rc.LoginAccount, ws.ErrMsg("sql脚本执行失败", fmt.Sprintf("[%s][%s] -> sql=[%s] 执行失败: [%s]", filename, dbConn.Info.GetLogDesc(), sql, err.Error())))
d.MsgApp.CreateAndSend(rc.LoginAccount, ws.ErrSysMsg("sql脚本执行失败", fmt.Sprintf("[%s][%s] -> sql=[%s] 执行失败: [%s]", filename, dbConn.Info.GetLogDesc(), sql, err.Error())))
return
}
}
d.MsgApp.CreateAndSend(rc.LoginAccount, ws.SuccessMsg("sql脚本执行成功", fmt.Sprintf("[%s]执行完成 -> %s", filename, dbConn.Info.GetLogDesc())))
d.MsgApp.CreateAndSend(rc.LoginAccount, ws.SuccessSysMsg("sql脚本执行成功", fmt.Sprintf("[%s]执行完成 -> %s", filename, dbConn.Info.GetLogDesc())))
}()
}
@@ -275,7 +275,7 @@ func (d *Db) DumpSql(rc *req.Ctx) {
if len(msg) > 0 {
msg = "数据库导出失败: " + msg
writer.WriteString(msg)
d.MsgApp.CreateAndSend(rc.LoginAccount, ws.ErrMsg("数据库导出失败", msg))
d.MsgApp.CreateAndSend(rc.LoginAccount, ws.ErrSysMsg("数据库导出失败", msg))
}
writer.Close()
}()

8
server/internal/machine/api/machine_file.go
View File

@@ -187,7 +187,7 @@ func (m *MachineFile) UploadFile(rc *req.Ctx) {
logx.Errorf("文件上传失败: %s", err)
switch t := err.(type) {
case biz.BizError:
m.MsgApp.CreateAndSend(la, ws.ErrMsg("文件上传失败", fmt.Sprintf("执行文件上传失败：\n<-e errCode: %d, errMsg: %s", t.Code(), t.Error())))
m.MsgApp.CreateAndSend(la, ws.ErrSysMsg("文件上传失败", fmt.Sprintf("执行文件上传失败：\n<-e errCode: %d, errMsg: %s", t.Code(), t.Error())))
}
}
}()
@@ -196,7 +196,7 @@ func (m *MachineFile) UploadFile(rc *req.Ctx) {
rc.ReqParam = jsonx.Kvs("machine", mi, "path", fmt.Sprintf("%s/%s", path, fileheader.Filename))
biz.ErrIsNilAppendErr(err, "创建文件失败: %s")
// 保存消息并发送文件上传成功通知
m.MsgApp.CreateAndSend(la, ws.SuccessMsg("文件上传成功", fmt.Sprintf("[%s]文件已成功上传至 %s[%s:%s]", fileheader.Filename, mi.Name, mi.Ip, path)))
m.MsgApp.CreateAndSend(la, ws.SuccessSysMsg("文件上传成功", fmt.Sprintf("[%s]文件已成功上传至 %s[%s:%s]", fileheader.Filename, mi.Name, mi.Ip, path)))
}
type FolderFile struct {
@@ -262,7 +262,7 @@ func (m *MachineFile) UploadFolder(rc *req.Ctx) {
logx.Errorf("文件上传失败: %s", err)
switch t := err.(type) {
case biz.BizError:
m.MsgApp.CreateAndSend(la, ws.ErrMsg("文件上传失败", fmt.Sprintf("执行文件上传失败：\n<-e errCode: %d, errMsg: %s", t.Code(), t.Error())))
m.MsgApp.CreateAndSend(la, ws.ErrSysMsg("文件上传失败", fmt.Sprintf("执行文件上传失败：\n<-e errCode: %d, errMsg: %s", t.Code(), t.Error())))
}
}
}()
@@ -286,7 +286,7 @@ func (m *MachineFile) UploadFolder(rc *req.Ctx) {
// 等待所有协程执行完成
wg.Wait()
// 保存消息并发送文件上传成功通知
m.MsgApp.CreateAndSend(rc.LoginAccount, ws.SuccessMsg("文件上传成功", fmt.Sprintf("[%s]文件夹已成功上传至 %s[%s:%s]", folderName, mi.Name, mi.Ip, basePath)))
m.MsgApp.CreateAndSend(rc.LoginAccount, ws.SuccessSysMsg("文件上传成功", fmt.Sprintf("[%s]文件夹已成功上传至 %s[%s:%s]", folderName, mi.Name, mi.Ip, basePath)))
}
func (m *MachineFile) RemoveFile(rc *req.Ctx) {

6
server/internal/msg/application/msg.go
View File

@@ -14,7 +14,7 @@ type Msg interface {
Create(msg *entity.Msg)
// 创建消息，并通过ws发送
CreateAndSend(la *model.LoginAccount, msg *ws.Msg)
CreateAndSend(la *model.LoginAccount, msg *ws.SysMsg)
}
func newMsgApp(msgRepo repository.Msg) Msg {
@@ -35,9 +35,9 @@ func (a *msgAppImpl) Create(msg *entity.Msg) {
a.msgRepo.Insert(msg)
}
func (a *msgAppImpl) CreateAndSend(la *model.LoginAccount, wmsg *ws.Msg) {
func (a *msgAppImpl) CreateAndSend(la *model.LoginAccount, wmsg *ws.SysMsg) {
now := time.Now()
msg := &entity.Msg{Type: 2, Msg: wmsg.Msg, RecipientId: int64(la.Id), CreateTime: &now, CreatorId: la.Id, Creator: la.Username}
msg := &entity.Msg{Type: 2, Msg: wmsg.SysMsg, RecipientId: int64(la.Id), CreateTime: &now, CreatorId: la.Id, Creator: la.Username}
a.msgRepo.Insert(msg)
ws.SendMsg(la.Id, wmsg)
}

2
server/internal/sys/api/system.go
View File

@@ -37,5 +37,5 @@ func (s *System) ConnectWs(g *gin.Context) {
// 登录账号信息
la := rc.LoginAccount
ws.Put(la.Id, wsConn)
ws.AddClient(la.Id, wsConn)
}

84
server/pkg/ws/client.go Normal file
View File

@@ -0,0 +1,84 @@
package ws
import (
"encoding/json"
"errors"
"mayfly-go/pkg/utils/stringx"
"time"
"github.com/gorilla/websocket"
)
type UserId uint64
// 客户端读取消息处理函数
// @param msg
type ReadMsgHandlerFunc func([]byte)
type Client struct {
ClientId string // 标识ID
UserId UserId // 用户ID
WsConn *websocket.Conn // 用户连接
ReadMsgHander ReadMsgHandlerFunc // 读取消息处理函数
}
func (c *Client) Read() {
go func() {
for {
messageType, data, err := c.WsConn.ReadMessage()
if err != nil {
if messageType == -1 && websocket.IsCloseError(err, websocket.CloseGoingAway, websocket.CloseNormalClosure, websocket.CloseNoStatusReceived) {
Manager.CloseClient(c)
return
}
if messageType != websocket.PingMessage {
return
}
}
if c.ReadMsgHander != nil {
c.ReadMsgHander(data)
}
}
}()
}
// 向客户端写入消息
func (c *Client) WriteMsg(msg *Msg) error {
if msg.Type == JsonMsg {
bytes, _ := json.Marshal(msg.Data)
return c.WsConn.WriteMessage(websocket.TextMessage, bytes)
}
if msg.Type == BinaryMsg {
if byteData, ok := msg.Data.([]byte); ok {
return c.WsConn.WriteMessage(websocket.BinaryMessage, byteData)
} else {
return errors.New("该数据不为数组类型")
}
}
if msg.Type == TextMsg {
if strData, ok := msg.Data.(string); ok {
return c.WsConn.WriteMessage(websocket.TextMessage, []byte(strData))
} else {
return errors.New("该数据类型不为字符串")
}
}
return errors.New("不存在该消息类型, 无法发送")
}
// 向客户写入ping消息
func (c *Client) Ping() error {
return c.WsConn.WriteControl(websocket.PingMessage, []byte{}, time.Now().Add(time.Second))
}
func NewClient(userId UserId, socket *websocket.Conn) *Client {
cli := &Client{
ClientId: stringx.Rand(16),
UserId: userId,
WsConn: socket,
}
return cli
}

160
server/pkg/ws/client_manager.go Normal file
View File

@@ -0,0 +1,160 @@
package ws
import (
"mayfly-go/pkg/logx"
"sync"
"time"
)
// 心跳间隔
var heartbeatInterval = 25 * time.Second
// 连接管理
type ClientManager struct {
ClientMap map[UserId]*Client // 全部的连接, key->userid, value->&client
RwLock sync.RWMutex // 读写锁
ConnectChan chan *Client // 连接处理
DisConnectChan chan *Client // 断开连接处理
MsgChan chan *Msg // 消息信息channel通道
}
func NewClientManager() (clientManager *ClientManager) {
return &ClientManager{
ClientMap: make(map[UserId]*Client),
ConnectChan: make(chan *Client, 10),
DisConnectChan: make(chan *Client, 10),
MsgChan: make(chan *Msg, 100),
}
}
// 管道处理程序
func (manager *ClientManager) Start() {
manager.HeartbeatTimer()
go manager.WriteMessage()
for {
select {
case client := <-manager.ConnectChan:
// 建立连接
manager.doConnect(client)
case conn := <-manager.DisConnectChan:
// 断开连接
manager.doDisconnect(conn)
}
}
}
// 添加客户端
func (manager *ClientManager) AddClient(client *Client) {
manager.ConnectChan <- client
}
// 关闭客户端
func (manager *ClientManager) CloseClient(client *Client) {
if client == nil {
return
}
manager.DisConnectChan <- client
}
// 根据用户id关闭客户端连接
func (manager *ClientManager) CloseByUid(uid UserId) {
manager.CloseClient(manager.GetByUid(UserId(uid)))
}
// 获取所有的客户端
func (manager *ClientManager) AllClient() map[UserId]*Client {
manager.RwLock.RLock()
defer manager.RwLock.RUnlock()
return manager.ClientMap
}
// 通过userId获取
func (manager *ClientManager) GetByUid(userId UserId) *Client {
manager.RwLock.RLock()
defer manager.RwLock.RUnlock()
return manager.ClientMap[userId]
}
// 客户端数量
func (manager *ClientManager) Count() int {
manager.RwLock.RLock()
defer manager.RwLock.RUnlock()
return len(manager.ClientMap)
}
// 发送json数据给指定用户
func (manager *ClientManager) SendJsonMsg(userId UserId, data any) {
logx.Debugf("发送消息: toUid=%v, data=%v", userId, data)
manager.MsgChan <- &Msg{ToUserId: userId, Data: data, Type: JsonMsg}
}
// 监听并发送给客户端信息
func (manager *ClientManager) WriteMessage() {
go func() {
for {
msg := <-manager.MsgChan
if cli := manager.GetByUid(msg.ToUserId); cli != nil {
if err := cli.WriteMsg(msg); err != nil {
manager.CloseClient(cli)
}
}
}
}()
}
// 启动定时器进行心跳检测
func (manager *ClientManager) HeartbeatTimer() {
go func() {
ticker := time.NewTicker(heartbeatInterval)
defer ticker.Stop()
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.Errorf("WS发送心跳失败: %v 总连接数：%d", userId, Manager.Count())
} else {
logx.Debugf("WS发送心跳成功: uid=%v", userId)
}
}
}
}()
}
// 处理建立连接
func (manager *ClientManager) doConnect(client *Client) {
cli := manager.GetByUid(client.UserId)
if cli != nil {
manager.doDisconnect(cli)
}
manager.addClient2Map(client)
logx.Debugf("WS客户端已连接: uid=%d, count=%d", client.UserId, manager.Count())
}
// 处理断开连接
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())
}
func (manager *ClientManager) addClient2Map(client *Client) {
manager.RwLock.Lock()
defer manager.RwLock.Unlock()
manager.ClientMap[client.UserId] = client
}
func (manager *ClientManager) delClient4Map(client *Client) {
manager.RwLock.Lock()
defer manager.RwLock.Unlock()
delete(manager.ClientMap, client.UserId)
}

45
server/pkg/ws/msg.go
View File

@@ -1,27 +1,46 @@
package ws
const SuccessMsgType = 1
const ErrorMsgType = 0
const InfoMsgType = 2
// 消息类型
type MsgType uint8
const (
JsonMsg MsgType = 1
TextMsg MsgType = 2
BinaryMsg MsgType = 3
)
// 消息信息
type Msg struct {
ToUserId UserId
Data any
Type MsgType // 消息类型
}
// ************** 系统消息 **************
const SuccessSysMsgType = 1
const ErrorSysMsgType = 0
const InfoSysMsgType = 2
// websocket消息
type Msg struct {
Type int `json:"type"` // 消息类型
Title string `json:"title"` // 消息标题
Msg string `json:"msg"` // 消息内容
type SysMsg struct {
Type int `json:"type"` // 消息类型
Title string `json:"title"` // 消息标题
SysMsg string `json:"msg"` // 消息内容
}
// 普通消息
func NewMsg(title, msg string) *Msg {
return &Msg{Type: InfoMsgType, Title: title, Msg: msg}
func NewSysMsg(title, msg string) *SysMsg {
return &SysMsg{Type: InfoSysMsgType, Title: title, SysMsg: msg}
}
// 成功消息
func SuccessMsg(title, msg string) *Msg {
return &Msg{Type: SuccessMsgType, Title: title, Msg: msg}
func SuccessSysMsg(title, msg string) *SysMsg {
return &SysMsg{Type: SuccessSysMsgType, Title: title, SysMsg: msg}
}
// 错误消息
func ErrMsg(title, msg string) *Msg {
return &Msg{Type: ErrorMsgType, Title: title, Msg: msg}
func ErrSysMsg(title, msg string) *SysMsg {
return &SysMsg{Type: ErrorSysMsgType, Title: title, SysMsg: msg}
}

59
server/pkg/ws/ws.go
View File

@@ -1,10 +1,7 @@
package ws
import (
"encoding/json"
"mayfly-go/pkg/logx"
"net/http"
"time"
"github.com/gorilla/websocket"
)
@@ -17,58 +14,24 @@ var Upgrader = websocket.Upgrader{
},
}
var conns = make(map[uint64]*websocket.Conn, 100)
var Manager = NewClientManager() // 管理者
func init() {
checkConn()
go Manager.Start()
}
// 放置ws连接
func Put(userId uint64, conn *websocket.Conn) {
existConn := conns[userId]
if existConn != nil {
Delete(userId)
}
conn.SetCloseHandler(func(code int, text string) error {
Delete(userId)
return nil
})
conns[userId] = conn
// 添加ws客户端
func AddClient(userId uint64, conn *websocket.Conn) *Client {
cli := NewClient(UserId(userId), conn)
Manager.AddClient(cli)
return cli
}
func checkConn() {
heartbeat := time.Duration(60) * time.Second
tick := time.NewTicker(heartbeat)
go func() {
for range tick.C {
// 遍历所有连接，ping失败的则删除掉
for uid, conn := range conns {
err := conn.WriteControl(websocket.PingMessage, []byte("ping"), time.Now().Add(heartbeat/2))
if err != nil {
Delete(uid)
return
}
}
}
}()
}
// 删除ws连接
func Delete(userid uint64) {
logx.Debugf("移除websocket连接: uid = %d", userid)
conn := conns[userid]
if conn != nil {
conn.Close()
delete(conns, userid)
}
func CloseClient(userid uint64) {
Manager.CloseByUid(UserId(userid))
}
// 对指定用户发送消息
func SendMsg(userId uint64, msg *Msg) {
conn := conns[userId]
if conn != nil {
bytes, _ := json.Marshal(msg)
conn.WriteMessage(websocket.TextMessage, bytes)
}
func SendMsg(userId uint64, msg *SysMsg) {
Manager.SendJsonMsg(UserId(userId), msg)
}