package ws import ( "fmt" "github.com/fox/fox/ipb" "github.com/fox/fox/log" "github.com/fox/fox/safeChan" "github.com/golang/protobuf/proto" "github.com/gorilla/websocket" "sync" "time" ) // 连接id var ( wsMsgType = websocket.BinaryMessage nextConnId uint32 ) // 客户端读写消息 type wsMessage struct { messageType int data []byte } // 客户端连接 type wsConnect struct { wsConn *websocket.Conn // 底层websocket inChan *safeChan.SafeChan[*wsMessage] // 读队列 outChan *safeChan.SafeChan[*wsMessage] // 写队列 id uint32 userId int64 onDisconnect func(IConn) once sync.Once } func newWsConnect(wsConn *websocket.Conn, onDisconnect func(IConn)) *wsConnect { c := &wsConnect{ wsConn: wsConn, inChan: safeChan.NewSafeChan[*wsMessage](1000), outChan: safeChan.NewSafeChan[*wsMessage](1000), id: nextConnId, userId: 0, onDisconnect: onDisconnect, } return c } // 把消息放进写队列 func (c *wsConnect) SendMsg(data []byte) error { return c.outChan.Write(&wsMessage{messageType: wsMsgType, data: data}) } // 关闭链接 func (c *wsConnect) Close() { c.once.Do(func() { //log.Debug(c.Log("关闭链接")) c.inChan.Close() c.outChan.Close() _ = c.wsConn.Close() if c.onDisconnect != nil { c.onDisconnect(c) } wsMgr.Remove(c) }) } // 循环从websocket中读取消息放入到读队列中 func (c *wsConnect) readWsLoop() { defer func() { //log.Debug(c.Log("readWsLoop协程退出")) c.Close() }() c.wsConn.SetReadLimit(maxMessageSize) _ = c.wsConn.SetReadDeadline(time.Now().Add(pongWait)) c.wsConn.SetPongHandler(func(string) error { //log.Debug(c.Log("received pong. from client")) _ = c.wsConn.SetReadDeadline(time.Now().Add(pongWait)) return nil }) for { // 读一个message msgType, data, err := c.wsConn.ReadMessage() if err != nil { if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure, websocket.CloseNormalClosure) { log.Error(c.Log("消息读取出现错误:%v", err)) } //log.Debug(c.Log("关闭连接:%v", err)) return } if msgType != wsMsgType { continue } msg := &wsMessage{messageType: msgType, data: data} if err = c.inChan.Write(msg); err != nil { log.Error(c.Log("读队列关闭:%v", err)) return } } } func (c *wsConnect) writeWsLoop() { ticker := time.NewTicker(pingPeriod) defer func() { //log.Debug(c.Log("writeWsLoop协程退出")) ticker.Stop() // 关闭连接 c.Close() }() for { select { // 取一个消息发送给客户端 case msg, ok := <-c.outChan.Reader(): if !ok { return } if err := c.wsConn.WriteMessage(msg.messageType, msg.data); err != nil { iMsg := &ipb.InternalMsg{} _ = proto.Unmarshal(msg.data, iMsg) log.Error(c.Log("发送消息错误:%v 消息长度:%v 消息内容:%v", err, len(msg.data), iMsg)) return } case <-ticker.C: if err := c.wsConn.WriteMessage(websocket.PingMessage, []byte("ping")); err != nil { log.Error(c.Log("发送心跳失败:%v", err)) return } else { //log.Debug(c.Log("发送心跳")) } } } } func (c *wsConnect) handle(process func(IConn, []byte)) { defer c.Close() for { select { case msg, ok := <-c.inChan.Reader(): if !ok { //log.Debug(c.Log("handle协程退出")) return } process(c, msg.data) } } } // 设置用户id func (c *wsConnect) setUserId(uid int64) { c.userId = uid } // 获取连接id func (c *wsConnect) Id() uint32 { return c.id } // 获取用户id func (c *wsConnect) UserId() int64 { return c.userId } func (c *wsConnect) Name() string { return fmt.Sprintf("用户:%v, 地址:%v", c.userId, c.wsConn.RemoteAddr()) } func (c *wsConnect) Addr() string { return c.wsConn.RemoteAddr().String() } func (c *wsConnect) Log(format string, v ...interface{}) string { s := fmt.Sprintf("连接:%v, id:%v uid:%v", c.wsConn.RemoteAddr().String(), c.id, c.userId) return s + fmt.Sprintf(format, v...) }