本次开发主要使用了Gorilla Websocket软件包

客户端结构体

维护socket连接，保存客户端信息

代码

package contorller

import (
	"encoding/json"
	"fmt"
	"git.myscrm.cn/vr/server-yk-app-builder/model"
	"git.myscrm.cn/vr/server-yk-app-builder/pkg/enum"
	"git.myscrm.cn/vr/server-yk-app-builder/pkg/logger"
	"git.myscrm.cn/vr/server-yk-app-builder/pkg/response"
	"github.com/gorilla/websocket"
	"strconv"
	"time"
)

const (
	// pingWait is the maximum time in seconds to wait for a ping from
	pingWait = 20 * time.Second
	// Maximum message size allowed from peer.
	maxMessageSize = 512
)

type wsClient struct {
	manager  *wsManager
	id       string
	desc     string
	conn     *websocket.Conn
	status   enum.ClientStatus
	send     chan []byte
	isClosed chan bool
	lastPing time.Time
}

func newWsClient(manager *wsManager, conn *websocket.Conn, id, name, ip string) *wsClient {
	return &wsClient{
		manager:  manager,
		id:       id,
		desc:     fmt.Sprintf("客户端(%s, %s)", name, ip),
		conn:     conn,
		status:   enum.ClientStatusOffline,
		send:     make(chan []byte),
		isClosed: make(chan bool),
		lastPing: time.Now(),
	}
}

// Read 读取客户端发送过来的消息
func (c *wsClient) Read() {
	defer func() {
		c.unRegister()
		logger.Info(c.desc, "read协程退出")
	}()

	c.conn.SetReadLimit(maxMessageSize)
	c.conn.SetPingHandler(func(text string) error {
		// 只需要知道客户端还活着就行，不需要回复
		c.lastPing = time.Now()
		// 更新客户端状态
		clientStatus, _ := strconv.ParseInt(text, 10, 32)
		c.status = enum.ClientStatus(int32(clientStatus))
		return nil
	})

	for {
		msgType, data, err := c.conn.ReadMessage()
		if err != nil {
			logger.Error(c.desc, "c.conn.ReadMessage", err.Error())
			break
		}

		switch msgType {
		case websocket.TextMessage:
			var msg *model.WsMessage
			err = json.Unmarshal(data, &msg)
			if err != nil {
				logger.Error(c.desc, "json.Unmarshal", err.Error())
				break
			}

			switch msg.Type {
			default:
				logger.Info(c.desc, "未知消息类型", string(data))
				c.send <- data
			}
		}
	}
}

// Write 把对应消息写回客户端
func (c *wsClient) Write() {
	defer func() {
		logger.Info(c.desc, "write协程退出")
		c.unRegister()
	}()
	for {
		select {
		case <-c.isClosed:
			return
		case msg := <-c.send:
			err := c.conn.WriteMessage(websocket.TextMessage, msg)
			if err != nil {
				logger.Error(c.desc, "c.conn.WriteMessage", err.Error())
				return
			}
		}
	}
}

// Check 检测客户端是否超时
func (c *wsClient) Check() {
	defer func() {
		logger.Info(c.desc, "check协程退出")
	}()
	ticker := time.NewTicker(pingWait / 6)
	for {
		select {
		case <-c.isClosed:
			return
		case <-ticker.C:
			// 主动关闭连接
			if time.Now().Sub(c.lastPing) > pingWait {
				response.WsReturnErr(c.conn, enum.WsDataErr, "客户端超时，主动关闭连接")
				logger.Info(c.desc, "客户端超时，主动关闭连接")
				return
			}
		}
	}
}

func (c *wsClient) unRegister() {
	if c.manager.clients[c.id] != nil {
		c.manager.unRegister <- c
	}
}

使用

应当在控制器/入口处

func (c *cSocket) Ws(ctx *gin.Context) {
	ws, err := c.upgrader.Upgrade(ctx.Writer, ctx.Request, nil)
	if err != nil {
		return nil, "", err
	}

	ip := ctx.ClientIP()
	name := ctx.GetString("device_name")

	client := newWsClient(WsManager, ws, clientId, name, ip)
	client.manager.register <- client

	go client.Read()
	go client.Write()
	go client.Check()
}

客户端管理器

管理客户端，支持登录、注销、广播

代码

package contorller

import (
	"encoding/json"
	"git.myscrm.cn/vr/server-yk-app-builder/model"
	"git.myscrm.cn/vr/server-yk-app-builder/pkg/enum"
	"git.myscrm.cn/vr/server-yk-app-builder/pkg/logger"
)

var WsManager = NewWsManager()

type wsManager struct {
	clients    map[string]*wsClient // 记录在线用户
	broadcast  chan []byte          // 触发消息广播
	register   chan *wsClient       // 触发设备或用户登陆
	unRegister chan *wsClient       // 触发设备或用户退出
}

func NewWsManager() *wsManager {
	return &wsManager{
		clients:    make(map[string]*wsClient),
		broadcast:  make(chan []byte),
		register:   make(chan *wsClient),
		unRegister: make(chan *wsClient),
	}
}

func (m *wsManager) Start() {
	for {
		select {
		case conn := <-m.register:
			m.clients[conn.id] = conn
			conn.status = enum.ClientStatusOnline
			logger.Info(conn.desc, "已连接", "当前在线客户端", len(m.clients))
			data, _ := json.Marshal(&model.WsMessage{Type: enum.WsDataNotify, Data: "socket 初始化成功"})
			conn.send <- data
		}
	}
}

func (m *wsManager) BoardCast() {
	for {
		select {
		case message := <-m.broadcast:
			for _, conn := range m.clients {
				conn.send <- message
			}
		}
	}
}

func (m *wsManager) Quit() {
	for {
		select {
		case conn := <-m.unRegister:
			if _, ok := m.clients[conn.id]; ok {
				delete(m.clients, conn.id)
				conn.conn.Close()
				conn.status = enum.ClientStatusOffline
				close(conn.isClosed)
			}
		}
	}
}

func (m *wsManager) GetClient(id string) *wsClient {
	if client, ok := m.clients[id]; ok {
		return client
	}
	return nil
}

使用

一般为main.go或启动处

func RegisterHttpRoute(serverMux *gin.Engine) error {
	go contorller.WsManager.Start()
	go contorller.WsManager.Quit()

	// 设置恐慌恢复，错误处理
	serverMux.Use(middleware.PanicHandler(), middleware.ErrorHandler())
	serverMux.NoMethod(middleware.NotFoundHandler())
	serverMux.NoRoute(middleware.NotFoundHandler())

	// 注册 websocket 接口
	serverMux.GET("/socket", middleware.SocketAuth(), contorller.Socket.Ws)

	...
}