Jacky's Blog Jacky's Blog
  • 首页
  • 关于
  • 项目
  • 大事记
  • 留言板
  • 友情链接
  • 分类
    • 干货
    • 随笔
    • 项目
    • 公告
    • 纪念
    • 尝鲜
    • 算法
    • 深度学习
首页 › 干货 › Gin WebSocket 消息推送

Gin WebSocket 消息推送

Jacky
10月 1, 2021干货阅读 1,575
目录
  1. 前言
  2. 代码
  3. 分析
  4. 测试效果

前言

这次打算在自己参与的大项目上用 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

赞(3)
本文系作者 @Jacky 原创发布在 Jacky's Blog。未经许可,禁止转载。
Nginx 安装配置指南
上一篇
air 实时热更新 Go 应用
下一篇
再想想
暂无评论
近期评论
  • Jacky发表在《Nginx UI》
  • daiwenzh5发表在《Nginx UI》
  • Jacky发表在《Nginx UI》
  • daiwenzh5发表在《Nginx UI》
  • Jacky发表在《Nginx UI》
3
  • 3
  • 0
Copyright © 2016-2023 Jacky's Blog. Designed by nicetheme.
粤ICP备16016168号-1
  • 首页
  • 关于
  • 项目
  • 大事记
  • 留言板
  • 友情链接
  • 分类
    • 干货
    • 随笔
    • 项目
    • 公告
    • 纪念
    • 尝鲜
    • 算法
    • 深度学习
# Mac # # Apple # # OS X # # iOS # # macOS #
Jacky
PHP C C++ Python | 舞象之年 | 物联网工程
174
文章
169
评论
267
喜欢