observe-server/common/ws/client.go

package ws

import (
	"github.com/go-admin-team/go-admin-core/logger"

	"github.com/gorilla/websocket"
)

// Client 单个 websocket 信息
type Client struct {
	ID, Group string
	Socket    *websocket.Conn
	Message   chan []byte
}

// MessageData 单个发送数据信息
type MessageData struct {
	ID, Group string
	Message   []byte
}

// MessageDataString 单个发送数据信息
type MessageDataString struct {
	ID, Group string
	Message   string
}

// GroupMessageData 组广播数据信息
type GroupMessageData struct {
	Group   string
	Message []byte
}

// BroadCastMessageData 广播发送数据信息
type BroadCastMessageData struct {
	Message []byte
}

// Read
func (c *Client) Read() {
	defer func() {
		WebsocketManager.UnRegister <- c
		logger.Infof("client [%s] disconnect", c.ID)
		if err := c.Socket.Close(); err != nil {
			logger.Errorf("client [%s] disconnect err: %s", c.ID, err)
		}
	}()

	for {
		logger.Debug("in for Read handler")
		messageType, message, err := c.Socket.ReadMessage()
		if err != nil || messageType == websocket.CloseMessage {
			break
		}
		logger.Debugf("client [%s] receive message: %s", c.ID, string(message))
		c.Message <- message
	}
}

// Write 写信息，从 channel 变量 Send 中读取数据写入 websocket 连接
func (c *Client) Write() {
	defer func() {
		WebsocketManager.UnRegister <- c //网络错误或客户端链接关闭，退出处理逻辑
		logger.Infof("client [%s] disconnect", c.ID)
		if err := c.Socket.Close(); err != nil {
			logger.Errorf("client [%s] disconnect err: %s", c.ID, err)
		}
	}()

	for {
		select {
		case message, ok := <-c.Message:
			if !ok {
				_ = c.Socket.WriteMessage(websocket.CloseMessage, []byte{})
				return
			}
			// logger.Debug("write message: ", message)
			err := c.Socket.WriteMessage(websocket.TextMessage, message)
			if err != nil {
				logger.Errorf("client [%s] writemessage err: %s", c.ID, err)
			}
		}
	}
}