前言

这次打算在自己参与的大项目上用 WebSocket 来为在线用户做消息推送了，具体的需求是要能为同一个用户多个客户端（多个浏览器页面，不同浏览器，不同终端）推送一份相同的消息。

在参考了多个项目之后我使用二维 map 存用户 user_id 和 WebSocket 的连接 *websocket.Conn，在不使用 go channel 的情况下简单的实现了 Gin WebSocket 的消息推送功能。

代码

注册路由

r.GET("live", live.WsHandler)

具体实现

package live

import (
        "github.com/0xJacky/SZTUEIS-api/model"
	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
	"log"
	"net/http"
	"sync"
	"time"
)

var (
	clients = make(map[uint]map[*websocket.Conn]bool)
	mux     sync.Mutex
)

var upGrader = websocket.Upgrader{
	ReadBufferSize:   1024,
	WriteBufferSize:  1024,
	HandshakeTimeout: 20 * time.Second,
	// 取消 ws 跨域校验
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

// WsHandler 处理ws请求
func WsHandler(c *gin.Context) {
	var conn *websocket.Conn
	var err error

	conn, err = upGrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		log.Printf("Failed to set websocket upgrade: %+v", err)
		return
	}

	var json struct{
		Jwt string `json:"jwt"`
	}

	err = conn.ReadJSON(&json)

	if err != nil {
		log.Println("read json err:", err)
		_ = conn.Close()
		return
	}

	// 验证逻辑
	var userId uint
	userId, err = model.ValidateJWT(json.Jwt)

	if err != nil {
		log.Println("validate jwt err:", err)
		_ = conn.Close()
		return
	}

	// 将 conn 保存到字典
	addClient(userId, conn)
}


func addClient(id uint, conn *websocket.Conn) {
	mux.Lock()
	if clients[id] == nil {
		clients[id] = make(map[*websocket.Conn]bool)
	}
	clients[id][conn] = true
	mux.Unlock()
}

func getClients(id uint) (conns []*websocket.Conn) {
	mux.Lock()
	_conns, ok := clients[id]
	if ok {
		for k := range _conns {
			conns = append(conns, k)
		}
	}
	mux.Unlock()
	return
}

func deleteClient(id uint, conn *websocket.Conn) {
	mux.Lock()
	_ = conn.Close()
	delete(clients[id], conn)
	mux.Unlock()
}

func SetMessage(userId uint, content interface{}) {
	conns := getClients(userId)
	for i := range conns {
		i := i
		err := conns[i].WriteJSON(content)
		if err != nil {
			log.Println("write json err:", err)
			deleteClient(userId, conns[i])
		}
	}
}

分析

在上述代码中，我使用 map[uint][*websocket.Conn]bool 二维字典将存储用户的多个连接。

在 WsHandler 函数中实现了将 HTTP 连接升级到 HTTP/1.1 实现长链接，然后根据前端发送的 jwt 做验证，得到当前用户的 user_id，将 user_id 和长链接通过 addClient 函数注册，保存到字典中。

在需要发送消息的时候，通过访问字典并做迭代，即可向指定用户的全部客户端发送相同的消息。

需要注意一个字典不能同时被多个 goroutine 访问并操作，作为经典的读者写者问题，要用到 mux (sync.Mutex) 做信号量，通过给 mux 加锁来实现安全的访问和操作，并在使用完毕后解锁，否则就死锁了。

测试效果

Gin WebSocket 消息推送-Jacky's Blog