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 }