fox/ws/wsConn.go

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
	connMgr      *connManager
}

func newWsConnect(wsConn *websocket.Conn, onDisconnect func(IConn), connMgr *connManager) *wsConnect {
	c := &wsConnect{
		wsConn:       wsConn,
		inChan:       safeChan.NewSafeChan[*wsMessage](1000),
		outChan:      safeChan.NewSafeChan[*wsMessage](1000),
		id:           nextConnId,
		userId:       0,
		onDisconnect: onDisconnect,
		connMgr:      connMgr,
	}
	return c
}

// 把消息放进写队列
func (c *wsConnect) SendMsg(data []byte) error {
	if c == nil {
		return fmt.Errorf("wsConnect is nil")
	}
	return c.outChan.Write(&wsMessage{messageType: wsMsgType, data: data})
}

// 关闭链接
func (c *wsConnect) Close() {
	if c == nil {
		return
	}
	c.once.Do(func() {
		//log.Debug(c.Log("关闭链接"))
		c.inChan.Close()
		c.outChan.Close()
		_ = c.wsConn.Close()
		if c.onDisconnect != nil {
			c.onDisconnect(c)
		}
		c.connMgr.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...)
}