带你用golang手撸一个websocket中间件【1】：少年，你的项目要不要用到websocket呀？（你说通讯用来

开始

【1】：同学，你的项目要不要用到websocket呀？

你说通讯用来干嘛？在线聊天、消息推送、扫码登录、物联网设备管理、游戏、弹幕、在线支付、视频通话等等这些需要立即得到反馈但又不想给服务器增加太多负担的业务场景都需要用websocket！

　【案例区域】

手把手带你从零开始写一个视频通话功能

（其实吧，主要为了加薪，面向加薪编程 \(^o^)/~）

【2】：为什么要使用websocket？

我们先来模拟一个最简单的扫码登录网页的场景，网页端在生成的二维码后要怎么知道用户有没有用手机扫描这个二维码呢？在传统的项目中最长的方式就是不停的去请求后端接口问 “用户扫了没？用户扫了没？用户扫了没？(一直往复）”，直到用户扫完或者关闭了网页。

这种方式就是最常见的 长轮训（我最开始学写代码的时候也是用这种方式），这种方式是最简单的，但是也与之相对应的问题也很明显，占用太多后台资源 and 感官延迟，那有没有一种别的方案能不占资源又快的方式呢？

在类似这种需求下 websocket 诞生了，今天这里我们不谈那些枯燥的理论知识，只玩实操！

<!DOCTYPE html>
<html>
<head>
</head>
<body>
<script>

  var ws = new WebSocket("ws://127.0.0.1:8282");

  ws.onmessage = function(event) {
    
      console.log(event.data)
  };

  ws.onclose = function(event) {

      console.log('ws连接已关闭')
  }


</script>
</body>
</html>

上面是一个最简单的 websocket 连接代码

嗯~，到这里，你已经能正常的连接到服务器上并等待 服务器主动 给你发送消息了。

好的，那现在客户端准备好了，服务端呢，别急，现在我们从 0 开始一起手撸一个websocket服务端

服务端设计

开始搭建服务端前，我们必须先思考一下它架构方式。

【1】：为什么是用golang来开发服务端，用其他语言不行吗？

当然是可以的！在文章的标题中我有提到这个用golang来开发，是因为其他语言都有现成可以用的socket中间件，但是golang似乎还没有，那我们就来自己撸一个吧！

其次，在我使用了大半年golang对它的理解来说，我觉得golang更适合用来做这种服务中间件，性能不错呢！

【2】：服务端定位

开始动手写代码前应该提前思考这么一个问题🤔，这个websocket服务应该以一种怎样的方式存在于项目中呢？

（嵌套组件 OR 独立中间件）

我的想法是中间件，出于以下原因考虑：

我希望它以一个独立角色存在于我的项目中
我不需要它过多的侵入业务代码
只要帮我管理好客户端的链接和能正常、快速的推送消息即可

【3】：架构设计

上面我有提到，这个websocket服务端的两个主要功能是 管理连接和推送消息

好的，那首先围绕第一个问题，如何管理连接？当服务端出现N多个连接时要怎么知道谁是谁，消息应该推给谁？

(w(ﾟДﾟ)w，啊~ 头好疼，感觉要长脑子了）

所以在经过一顿头脑风暴和参考其他语言的websocket框架后，我决定用client、user、group来管理所有所有客户端连接，其实就是 分类管理，下面我分别解释。

client（连接） 任何一个websocket连接都会被视为一个client，所有的操作也是围绕着client在进行管理
user （用户） 每个用户可能会打开多个websocket或者说叫页面，那要把消息发给这个用户时就需要把消息通知到不同的client，所以这时就可以把几个client归类为一个user。
group （群组） 还有另外一种需求，需要把消息通知到某些用户或者某几个用户，这时就要考虑把client归类到一个group

在看下面的内容前，大家一定要先消化了解这三个概念！！！

还有另外一个问题，怎么让websocket服务不侵入业务代码

这里我大概画了一张草图，三者之间的关系可以这样理解

客户端除了连接websocket和发送心跳信息（ping/pong，这个是服务所必须的）外，只要被动接受Socket服务段推送过来的消息即可，前端页面在根据推送过来的消息做对应业务处理。
客户端需要发送消息时应当将请求发送到后端服务，再由后端服务经过业务逻辑处理后调取对应的Socket接口
Socket服务器除了管理连接和推送消息外不用再处理任何与业务相关的内容

掉头发时间

在博客我只展示一点点代码哈，其他的都已经完全开源到github了，各位看官需要的话自取哈

(记得别忘了给个stars呀)

我们要使用golang来实现socket服务，自然离不开 github.com/gorilla/websocket 这个核心库啦！

这里不得不说，golang的生态还是挺完善的。

gorilla/websocket帮我们解决了socket的连接和推送问题，剩下连接关系管理和服务接口就是我们要关注的重点了。

【1】：连接关系管理

先来给大家上两段代码

server.go

package websocket

import (
	"fmt"
	"sync"
	"time"

	"github.com/golang-module/carbon"
	"github.com/gorilla/websocket"
)

type WebSocketClientBase struct {
	ID            string
	Conn          *websocket.Conn
	LastHeartbeat int64
	BindUid       string
	JoinGroup     []string
}

type WebSocketUserBase struct {
	Uid      string
	ClientID []string
}

type WebSocketGroupBase struct {
	ClientID []string
}

var GatewayClients, GatewayUser, GatewayGroup sync.Map

/**
 * @description: 客户端心跳检测，超时即断开连接（主要是为了降低服务端承载压力）
 * @param {string} clientID
 * @return {*}
 */
func clientHeartbeatCheck(clientID string) {

	for {

		time.Sleep(5 * time.Second)

		clientInterface, exists := GatewayClients.Load(clientID)

		if !exists {

			break
		}

		client, _ := clientInterface.(*WebSocketClientBase)

		if (carbon.Now().Timestamp() - client.LastHeartbeat) > int64(HeartbeatTime) {

			fmt.Println("Client", clientID, "heartbeat timeout")

			client.Conn.Close()
			GatewayClients.Delete(clientID)
			break
		}
	}
}

/**
 * @description: 客户端断线时自动踢出Uid绑定列表
 * @param {string} clientID
 * @param {string} uid
 * @return {*}
 */
func clientUnBindUid(clientID string, uid string) {

	value, ok := GatewayUser.Load(uid)

	if ok {

		users := value.(*WebSocketUserBase)

		for k, v := range users.ClientID {

			if v == clientID {

				users.ClientID = append(users.ClientID[:k], users.ClientID[k+1:]...)
			}
		}

		if len(users.ClientID) == 0 {

			GatewayUser.Delete(uid)
		}

	}
}

/**
 * @description: 客户端断线时自动踢出已加入的群组
 * @param {string} clientID
 * @return {*}
 */
func clientLeaveGroup(clientID string) {
	// 使用 Load 方法获取值
	value, ok := GatewayClients.Load(clientID)
	if !ok {
		// 如果没有找到对应的值，处理相应的逻辑
		return
	}

	client := value.(*WebSocketClientBase)

	// 遍历 JoinGroup
	for _, v := range client.JoinGroup {
		// 使用 Load 方法获取值
		groupValue, groupOK := GatewayGroup.Load(v)
		if !groupOK {
			// 如果没有找到对应的值，处理相应的逻辑
			continue
		}

		group := groupValue.(*WebSocketGroupBase)

		// 在群组中找到对应的 clientID，并删除
		for j, id := range group.ClientID {
			if id == clientID {
				copy(group.ClientID[j:], group.ClientID[j+1:])
				group.ClientID = group.ClientID[:len(group.ClientID)-1]

				// 如果群组中没有成员了，删除群组
				if len(group.ClientID) == 0 {
					GatewayGroup.Delete(v)
				}

				break
			}
		}
	}
}

connect.go

package websocket

import (
	"fmt"
	"gateway-websocket/config"
	"net/http"
	"runtime/debug"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/golang-module/carbon"
	"github.com/google/uuid"
	"github.com/gorilla/websocket"
)

var (
	upGrader = websocket.Upgrader{
		// 设置消息接收缓冲区大小（byte），如果这个值设置得太小，可能会导致服务端在读取客户端发送的大型消息时遇到问题
		ReadBufferSize: config.GatewayConfig["ReadBufferSize"].(int),
		// 设置消息发送缓冲区大小（byte），如果这个值设置得太小，可能会导致服务端在发送大型消息时遇到问题
		WriteBufferSize: config.GatewayConfig["WriteBufferSize"].(int),
		// 消息包启用压缩
		EnableCompression: config.GatewayConfig["MessageCompression"].(bool),
		// ws握手超时时间
		HandshakeTimeout: time.Duration(config.GatewayConfig["WebsocketHandshakeTimeout"].(int)) * time.Second,
		// ws握手过程中允许跨域
		CheckOrigin: func(r *http.Request) bool {
			return true
		},
	}

	// 设置心跳检测间隔时长（秒）
	HeartbeatTime = config.GatewayConfig["HeartbeatTimeout"].(int)
)

/**
 * @description: 初始化客户端连接
 * @param {*websocket.Conn} conn
 * @return {*}
 */
func handleClientInit(conn *websocket.Conn) string {

	clientID := uuid.New().String()

	client := &WebSocketClientBase{
		ID:            clientID,
		Conn:          conn,
		LastHeartbeat: carbon.Now().Timestamp(),
	}

	// 使用 Store 方法存储值
	GatewayClients.Store(clientID, client)

	if err := conn.WriteMessage(config.GatewayConfig["MessageFormat"].(int), []byte(clientID)); err != nil {

		handleClientDisconnect(clientID)
		return ""
	}

	return clientID
}

/**
 * @description: 主动关闭客户端连接
 * @param {string} clientID
 * @return {*}
 */
func handleClientDisconnect(clientID string) {

	// 使用 Load 和 Delete 方法，不需要额外的锁定操作
	v, ok := GatewayClients.Load(clientID)
	if ok {

		client := v.(*WebSocketClientBase)

		if client.BindUid != "" {
			clientUnBindUid(clientID, client.BindUid)
		}

		if len(client.JoinGroup) > 0 {
			clientLeaveGroup(clientID)
		}

		GatewayClients.Delete(clientID)
	}
}

/**
 * @description: 向客户端回复心跳消息
 * @param {*websocket.Conn} conn
 * @param {string} clientID
 * @param {int} messageType
 * @param {[]byte} message
 * @return {*}
 */
func handleClientMessage(conn *websocket.Conn, clientID string, messageType int, message []byte) {

	// 使用 Load 方法获取值
	v, ok := GatewayClients.Load(clientID)
	if !ok {
		// 如果没有找到对应的值，处理相应的逻辑
		handleClientDisconnect(clientID)
		return
	}

	client := v.(*WebSocketClientBase)

	if messageType == config.GatewayConfig["MessageFormat"].(int) && string(message) == "ping" {

		if err := conn.WriteMessage(config.GatewayConfig["MessageFormat"].(int), []byte("pong")); err != nil {

			handleClientDisconnect(clientID)
			return
		}

		GatewayClients.Store(clientID, &WebSocketClientBase{
			ID:            clientID,
			Conn:          conn,
			LastHeartbeat: carbon.Now().Timestamp(),
			BindUid:       client.BindUid,
			JoinGroup:     client.JoinGroup,
		})
	}
}

func WsServer(c *gin.Context) {

	defer func() {
		if err := recover(); err != nil {
			fmt.Printf("WsServer panic: %v\n", err)
			debug.PrintStack()
		}
	}()

	// 将 HTTP 连接升级为 WebSocket 连接
	conn, err := upGrader.Upgrade(c.Writer, c.Request, nil)

	if err != nil {
		return
	}

	defer conn.Close()

	// 客户端唯一身份标识
	clientID := handleClientInit(conn)

	// 发送客户端唯一标识 ID
	if clientID == "" {
		return
	}

	go clientHeartbeatCheck(clientID)

	for {

		// 读取客户端发送过来的消息
		messageType, message, err := conn.ReadMessage()

		// 当收到err时则标识客户端连接出现异常，如断线
		if err != nil {

			handleClientDisconnect(clientID)

		} else {

			handleClientMessage(conn, clientID, messageType, message)
		}
	}

}

在上面的代码中，我创建了一个websocket的连接服务和使用了3个sync.Map来分别存放管理不同的客户端连接

【TIPS: 在做这种存在高并发场景的业务时不要使用Map而是用sync.Map，因为go的Map是非线程安全的，在并发时会造成资源竞争从而导致你的程序宕掉，这点一定要注意！！！】

Stop！！！，文章好像被拉的太长了(⊙o⊙)…，那就只展示一点点吧，剩下的代码我全开源到github啦，大家自行参阅。

操作实例

经过上面的过程后，服务端已经可以给客户端推送消息了。然后就是二个问题，要怎么管理这些客户端连接呢？websocket的服务器我们也准备好了，那后端要怎么去调用接口来操作client呢？

【这里我给大家准备好了不同后端语言的操作demo（其他语言会陆续完善）】

【TIPS: 如果是使用php操作连接的话需要安装grpc扩展哦~，否则调用时会抛出异常！！！】

上面我有提到管理连接都是围绕着 client、user、group 这三个概念在运行的，所以我也在框架中提供了一些gRPC接口方便大家的后端程序来管理连接。

接口名称	用途	归属分类
BroadcastMessage	向所有client广播消息	client
SendMessageToClient	向某个client发送消息	client
ClonseClient	关闭某个client的连接	client
GetAllOnlineClient	获取所有在线客户端	client
ClientIsOnline	判断client是否在线	client
CountOnlineClient	统计在线client数量	client
UnGroup	解散某个群组	group
SendMessageToGroup	向某个群组发送消息	group
CountOnlineGroup	统计群组内在线的client数量	group
CountGroup	统计群组数量	group
GetGroupOnlineClient	获取群组内所有在线的client	group
LeaveGroup	将client移出群组	group
JoinGroup	将client加入群组	group
GetUidByClient	获取Uid下所有client	user
CountOnlineUid	统计在线uid人数	user
GetAllOnlineUid	获取所有在线uid	user
SendMessageToUid	向某个uid下绑定的所有client发送消息	user
UidIsOnline	判断某个uid是否存在	user
UnBindUid	client解除绑定Uid	user
ClientBindUid	client绑定Uid	user
GetClientByUid	获取client绑定的Uid	user

这样思路是不是就清晰啦？

测试时间

代码写完，先把程序run起来

然后压测安排上