sonet/internal/gateway_ws/session/net_client.go

package session

import (
	"fmt"
	"github.com/gorilla/websocket"
	"sonet/pkg/utils/logger"
	"sync"
	"sync/atomic"
	"time"
)

// NetClient 长连接客户端
type NetClient struct {
	Conn                        *websocket.Conn
	writeDeadline, readDeadline time.Duration
	writeLock                   *sync.Mutex
	Subject                     *NetSubject
	closed                      *atomic.Bool
}

func NewNetClient(conn *websocket.Conn, readDeadline, writeDeadline time.Duration) *NetClient {
	return &NetClient{
		Conn:          conn,
		readDeadline:  readDeadline,
		writeDeadline: writeDeadline,
		writeLock:     new(sync.Mutex),
		closed:        &atomic.Bool{},
	}
}

func (c *NetClient) Close() {
	c.writeLock.Lock()
	defer c.writeLock.Unlock()
	if c.closed.Load() {
		return
	}
	c.closed.Store(true)
	err := c.Conn.Close()
	if err != nil {
		logger.Error("NetClient Close error: ", err)
	}
}

func (c *NetClient) IsClosed() bool {
	return c.closed.Load()
}

func (c *NetClient) Write(bytes []byte) (err error) {
	c.writeLock.Lock()
	defer c.writeLock.Unlock()
	err = c.Conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
	if err != nil {
		err = fmt.Errorf("NetClient SetWriteDeadline error: %s", err.Error())
		return
	}
	err = c.Conn.WriteMessage(websocket.BinaryMessage, bytes)
	if err != nil {
		err = fmt.Errorf("NetClient write message error: %s", err.Error())
	}
	return
}

func (c *NetClient) MustWrite(bytes []byte) {
	err := c.Write(bytes)
	if err != nil {
		logger.Error(err)
		return
	}
}

// ReadMessage 阻塞读取消息
func (c *NetClient) ReadMessage() (ignore bool, bytes []byte, err error) {
	// err = c.Conn.SetReadDeadline(time.Now().Add(c.readDeadline))
	err = c.Conn.SetReadDeadline(time.Time{})
	if err != nil {
		return
	}

	var messageType int
	messageType, bytes, err = c.Conn.ReadMessage()
	if messageType != websocket.BinaryMessage {
		ignore = true
		logger.Infof("read websocket message type: %d", messageType)
		return
	}
	return
}