fusenapi/server/websocket/internal/handler/datatransferhandler.go

261 lines
6.4 KiB
Go
Raw Normal View History

2023-07-24 08:07:05 +00:00
package handler
import (
2023-07-24 11:44:28 +00:00
"encoding/json"
2023-07-25 04:06:19 +00:00
"fmt"
2023-07-24 11:44:28 +00:00
"fusenapi/constants"
2023-07-25 09:10:50 +00:00
"fusenapi/server/websocket/internal/logic"
2023-07-24 08:07:05 +00:00
"fusenapi/server/websocket/internal/svc"
2023-07-24 11:44:28 +00:00
"fusenapi/server/websocket/internal/types"
"github.com/google/uuid"
2023-07-24 08:07:05 +00:00
"github.com/gorilla/websocket"
2023-07-24 11:44:28 +00:00
"github.com/zeromicro/go-zero/core/logx"
2023-07-24 08:07:05 +00:00
"net/http"
"sync"
2023-07-24 11:44:28 +00:00
"time"
2023-07-24 08:07:05 +00:00
)
var (
//升级
upgrade = websocket.Upgrader{
//允许跨域
CheckOrigin: func(r *http.Request) bool {
return true
},
}
//连接map池
2023-07-25 09:10:50 +00:00
mapConnPool = sync.Map{}
2023-07-24 08:07:05 +00:00
)
2023-07-24 11:44:28 +00:00
// 每个连接的连接属性
2023-07-24 08:07:05 +00:00
type wsConnectItem struct {
2023-07-25 09:10:50 +00:00
conn *websocket.Conn //websocket的连接
closeChan chan struct{} //关闭chan
isClose bool //是否已经关闭
flag string
inChan chan []byte //接受消息缓冲通道
outChan chan []byte //发送回客户端的消息
mutex sync.Mutex
renderImage map[string]struct{} //需要渲染的图片
2023-07-24 08:07:05 +00:00
}
func DataTransferHandler(svcCtx *svc.ServiceContext) http.HandlerFunc {
return func(w http.ResponseWriter, r *http.Request) {
2023-07-25 09:10:50 +00:00
// 创建一个业务逻辑层实例
var req types.DataTransferReq
l := logic.NewDataTransferLogic(r.Context(), svcCtx)
l.DataTransfer(&req, nil)
2023-07-25 04:06:19 +00:00
//升级websocket
conn, err := upgrade.Upgrade(w, r, nil)
if err != nil {
logx.Error("http upgrade websocket err:", err)
return
}
defer conn.Close()
rsp := types.DataTransferRsp{}
2023-07-24 11:44:28 +00:00
// 解析JWT token,并对空用户进行判断
2023-07-25 09:10:50 +00:00
/*claims, err := svcCtx.ParseJwtToken(r)
2023-07-24 08:07:05 +00:00
// 如果解析JWT token出错,则返回未授权的JSON响应并记录错误消息
if err != nil {
2023-07-25 04:06:19 +00:00
rsp.MsgType = constants.WEBSOCKET_UNAUTH
rsp.Message = "unAuth"
b, _ := json.Marshal(rsp)
_ = conn.WriteMessage(websocket.TextMessage, b)
2023-07-24 08:07:05 +00:00
return
}
if claims != nil {
// 从token中获取对应的用户信息
2023-07-24 11:44:28 +00:00
_, err = auth.GetUserInfoFormMapClaims(claims)
2023-07-24 08:07:05 +00:00
// 如果获取用户信息出错,则返回未授权的JSON响应并记录错误消息
if err != nil {
2023-07-25 04:06:19 +00:00
rsp.MsgType = constants.WEBSOCKET_UNAUTH
rsp.Message = "unAuth!!"
b, _ := json.Marshal(rsp)
_ = conn.WriteMessage(websocket.TextMessage, b)
2023-07-24 08:07:05 +00:00
return
}
} else {
// 如果claims为nil,则认为用户身份为白板用户
2023-07-25 04:06:19 +00:00
rsp.MsgType = constants.WEBSOCKET_UNAUTH
rsp.Message = "unAuth!!!"
b, _ := json.Marshal(rsp)
_ = conn.WriteMessage(websocket.TextMessage, b)
2023-07-24 08:07:05 +00:00
return
2023-07-25 09:10:50 +00:00
}*/
2023-07-24 11:44:28 +00:00
//生成连接唯一标识
2023-07-25 04:06:19 +00:00
flag := uuid.New().String() + time.Now().Format("20060102150405")
ws := wsConnectItem{
2023-07-25 09:10:50 +00:00
conn: conn,
flag: flag,
closeChan: make(chan struct{}, 1),
inChan: make(chan []byte, 100),
outChan: make(chan []byte, 100),
renderImage: make(map[string]struct{}),
isClose: false,
2023-07-24 11:44:28 +00:00
}
//保存连接
2023-07-25 09:10:50 +00:00
mapConnPool.Store(flag, ws)
2023-07-25 04:06:19 +00:00
defer ws.close()
//把连接成功消息发回去
2023-07-25 10:04:01 +00:00
rsp.T = constants.WEBSOCKET_CONNECT_SUCCESS
rsp.D = flag
2023-07-24 11:44:28 +00:00
b, _ := json.Marshal(rsp)
conn.WriteMessage(websocket.TextMessage, b)
2023-07-25 04:06:19 +00:00
//循环读客户端信息
go ws.readLoop()
2023-07-25 09:10:50 +00:00
//循环把数据发送给客户端
go ws.writeLoop()
2023-07-25 04:06:19 +00:00
//推消息到云渲染
go ws.sendLoop()
//心跳
2023-07-25 09:10:50 +00:00
ws.heartbeat()
}
}
// 心跳
func (w *wsConnectItem) heartbeat() {
rsp := types.DataTransferRsp{
2023-07-25 10:04:01 +00:00
T: constants.WEBSOCKET_HEARTBEAT,
D: "",
2023-07-25 09:10:50 +00:00
}
2023-07-25 11:37:56 +00:00
b, _ := json.Marshal(rsp)
2023-07-25 09:10:50 +00:00
for {
time.Sleep(time.Second * 10)
select {
case <-w.closeChan:
return
default:
2023-07-25 10:34:56 +00:00
}
//发送心跳信息
if err := w.conn.WriteMessage(websocket.TextMessage, b); err != nil {
logx.Error("发送心跳信息异常,关闭连接:", w.flag, err)
w.close()
return
2023-07-25 04:06:19 +00:00
}
}
}
// 关闭连接
func (w *wsConnectItem) close() {
2023-07-25 09:10:50 +00:00
w.mutex.Lock()
defer w.mutex.Unlock()
2023-07-25 04:06:19 +00:00
logx.Info("websocket:", w.flag, " is closing...")
w.conn.Close()
2023-07-25 09:10:50 +00:00
mapConnPool.Delete(w.flag)
2023-07-25 04:06:19 +00:00
if !w.isClose {
w.isClose = true
close(w.closeChan)
2023-07-25 09:10:50 +00:00
close(w.outChan)
2023-07-25 04:06:19 +00:00
close(w.inChan)
2023-07-24 08:07:05 +00:00
}
2023-07-25 04:06:19 +00:00
logx.Info("websocket:", w.flag, " is closed")
2023-07-24 08:07:05 +00:00
}
2023-07-24 11:44:28 +00:00
2023-07-25 09:10:50 +00:00
func (w *wsConnectItem) writeLoop() {
for {
select {
case <-w.closeChan: //如果关闭了
return
case data := <-w.outChan:
w.conn.WriteMessage(websocket.TextMessage, data)
}
}
}
2023-07-25 04:06:19 +00:00
// 接受客户端发来的消息
func (w *wsConnectItem) readLoop() {
for {
select {
case <-w.closeChan: //如果关闭了
return
default:
2023-07-25 10:34:56 +00:00
}
_, data, err := w.conn.ReadMessage()
if err != nil {
logx.Error("接受信息错误:", err)
//关闭连接
w.close()
return
2023-07-25 04:06:19 +00:00
}
2023-07-25 10:34:56 +00:00
//消息传入缓冲通道
w.inChan <- data
2023-07-25 04:06:19 +00:00
}
}
2023-07-24 11:44:28 +00:00
2023-07-25 04:06:19 +00:00
// 把收到的消息发往不同的地方处理
func (w *wsConnectItem) sendLoop() {
for {
select {
case <-w.closeChan:
return
case data := <-w.inChan:
2023-07-25 09:10:50 +00:00
w.dealwithReciveData(data)
}
}
}
// 处理接受到的数据
func (w *wsConnectItem) dealwithReciveData(data []byte) {
var parseInfo types.DataTransferReq
if err := json.Unmarshal(data, &parseInfo); err != nil {
logx.Error("invalid format of websocket message")
return
}
2023-07-25 10:04:01 +00:00
switch parseInfo.T {
2023-07-25 09:41:38 +00:00
//图片渲染
case constants.WEBSOCKET_RENDER_IMAGE:
2023-07-25 09:10:50 +00:00
var renderImageData []types.RenderImageReqMsg
2023-07-25 10:04:01 +00:00
if err := json.Unmarshal([]byte(parseInfo.D), &renderImageData); err != nil {
2023-07-25 09:10:50 +00:00
logx.Error("invalid format of websocket render image message", err)
return
2023-07-25 04:06:19 +00:00
}
2023-07-25 09:10:50 +00:00
logx.Info("收到请求云渲染图片数据:", renderImageData)
w.mutex.Lock()
defer w.mutex.Unlock()
//把需要渲染的图片加进去
for _, v := range renderImageData {
2023-07-25 09:12:37 +00:00
key := w.getRenderImageMapKey(v.ProductId, v.SizeId, v.TemplateId)
2023-07-25 09:10:50 +00:00
w.renderImage[key] = struct{}{}
}
default:
2023-07-25 04:06:19 +00:00
}
2023-07-24 11:44:28 +00:00
}
2023-07-25 09:10:50 +00:00
2023-07-25 11:33:59 +00:00
// 把渲染好的数据放入outchan
2023-07-25 11:35:52 +00:00
func (w *wsConnectItem) setOutRenderImage(req types.RenderNotifyReq) {
w.mutex.Lock()
defer w.mutex.Unlock()
2023-07-25 11:33:59 +00:00
for _, notifyItem := range req.NotifyList {
2023-07-25 11:35:52 +00:00
renderKey := w.getRenderImageMapKey(notifyItem.ProductId, notifyItem.SizeId, notifyItem.TemplateId)
2023-07-25 11:33:59 +00:00
//查询
2023-07-25 11:35:52 +00:00
_, ok := w.renderImage[renderKey]
2023-07-25 11:33:59 +00:00
if !ok {
continue
}
responseData := types.RenderImageRspMsg{
ProductId: notifyItem.ProductId,
SizeId: notifyItem.SizeId,
TemplateId: notifyItem.TemplateId,
Source: "我是渲染资源",
}
b, _ := json.Marshal(responseData)
select {
2023-07-25 11:35:52 +00:00
case <-w.closeChan:
2023-07-25 11:33:59 +00:00
return
2023-07-25 11:35:52 +00:00
case w.outChan <- b:
2023-07-25 11:33:59 +00:00
logx.Info("notify send render result to out chan")
}
//删掉已经处理的渲染任务
2023-07-25 11:35:52 +00:00
delete(w.renderImage, renderKey)
2023-07-25 11:33:59 +00:00
}
}
2023-07-25 09:10:50 +00:00
// 获取需要渲染图片的map key
2023-07-25 09:12:37 +00:00
func (w *wsConnectItem) getRenderImageMapKey(productId, sizeId, templateId int64) string {
2023-07-25 09:10:50 +00:00
return fmt.Sprintf("%d-%d-%d", productId, sizeId, templateId)
}