package handler import ( "encoding/json" "fmt" "fusenapi/constants" "fusenapi/server/websocket/internal/logic" "fusenapi/server/websocket/internal/svc" "fusenapi/server/websocket/internal/types" "github.com/google/uuid" "github.com/gorilla/websocket" "github.com/zeromicro/go-zero/core/logx" "net/http" "sync" "time" ) var ( //升级 upgrade = websocket.Upgrader{ //允许跨域 CheckOrigin: func(r *http.Request) bool { return true }, } //连接map池 mapConnPool = sync.Map{} ) // 每个连接的连接属性 type wsConnectItem struct { 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{} //需要渲染的图片 } func DataTransferHandler(svcCtx *svc.ServiceContext) http.HandlerFunc { return func(w http.ResponseWriter, r *http.Request) { // 创建一个业务逻辑层实例 var req types.DataTransferReq l := logic.NewDataTransferLogic(r.Context(), svcCtx) l.DataTransfer(&req, nil) //升级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{} // 解析JWT token,并对空用户进行判断 /*claims, err := svcCtx.ParseJwtToken(r) // 如果解析JWT token出错,则返回未授权的JSON响应并记录错误消息 if err != nil { rsp.MsgType = constants.WEBSOCKET_UNAUTH rsp.Message = "unAuth" b, _ := json.Marshal(rsp) _ = conn.WriteMessage(websocket.TextMessage, b) return } if claims != nil { // 从token中获取对应的用户信息 _, err = auth.GetUserInfoFormMapClaims(claims) // 如果获取用户信息出错,则返回未授权的JSON响应并记录错误消息 if err != nil { rsp.MsgType = constants.WEBSOCKET_UNAUTH rsp.Message = "unAuth!!" b, _ := json.Marshal(rsp) _ = conn.WriteMessage(websocket.TextMessage, b) return } } else { // 如果claims为nil,则认为用户身份为白板用户 rsp.MsgType = constants.WEBSOCKET_UNAUTH rsp.Message = "unAuth!!!" b, _ := json.Marshal(rsp) _ = conn.WriteMessage(websocket.TextMessage, b) return }*/ //生成连接唯一标识 flag := uuid.New().String() + time.Now().Format("20060102150405") ws := wsConnectItem{ 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, } //保存连接 mapConnPool.Store(flag, ws) defer ws.close() //把连接成功消息发回去 rsp.T = constants.WEBSOCKET_CONNECT_SUCCESS rsp.D = flag b, _ := json.Marshal(rsp) conn.WriteMessage(websocket.TextMessage, b) //循环读客户端信息 go ws.readLoop() //循环把数据发送给客户端 go ws.writeLoop() //推消息到云渲染 go ws.sendLoop() //心跳 ws.heartbeat() } } // 心跳 func (w *wsConnectItem) heartbeat() { rsp := types.DataTransferRsp{ T: constants.WEBSOCKET_HEARTBEAT, D: "", } for { time.Sleep(time.Second * 10) select { case <-w.closeChan: return default: //发送心跳信息 b, _ := json.Marshal(rsp) if err := w.conn.WriteMessage(websocket.TextMessage, b); err != nil { logx.Error("发送心跳信息异常,关闭连接:", w.flag, err) w.close() return } } } } // 关闭连接 func (w *wsConnectItem) close() { w.mutex.Lock() defer w.mutex.Unlock() logx.Info("websocket:", w.flag, " is closing...") w.conn.Close() mapConnPool.Delete(w.flag) if !w.isClose { w.isClose = true close(w.closeChan) close(w.outChan) close(w.inChan) } logx.Info("websocket:", w.flag, " is closed") } func (w *wsConnectItem) writeLoop() { for { select { case <-w.closeChan: //如果关闭了 return case data := <-w.outChan: w.conn.WriteMessage(websocket.TextMessage, data) } } } // 接受客户端发来的消息 func (w *wsConnectItem) readLoop() { for { select { case <-w.closeChan: //如果关闭了 return default: _, data, err := w.conn.ReadMessage() if err != nil { logx.Error("接受信息错误:", err) //关闭连接 w.close() return } //消息传入缓冲通道 w.inChan <- data } } } // 把收到的消息发往不同的地方处理 func (w *wsConnectItem) sendLoop() { for { select { case <-w.closeChan: return case data := <-w.inChan: 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 } switch parseInfo.T { //图片渲染 case constants.WEBSOCKET_RENDER_IMAGE: var renderImageData []types.RenderImageReqMsg if err := json.Unmarshal([]byte(parseInfo.D), &renderImageData); err != nil { logx.Error("invalid format of websocket render image message", err) return } logx.Info("收到请求云渲染图片数据:", renderImageData) w.mutex.Lock() defer w.mutex.Unlock() //把需要渲染的图片加进去 for _, v := range renderImageData { key := w.getRenderImageMapKey(v.ProductId, v.SizeId, v.TemplateId) w.renderImage[key] = struct{}{} } default: } } // 获取需要渲染图片的map key func (w *wsConnectItem) getRenderImageMapKey(productId, sizeId, templateId int64) string { return fmt.Sprintf("%d-%d-%d", productId, sizeId, templateId) }