有了前两篇的铺垫，相信大家已经对 Golang 中 WebSocket 的使用有一定的了解了，
今天我们以一个更加真实的例子来学习如何在 Golang 中使用 WebSocket。
需求背景
在实际的项目中，往往有一些任务耗时比较长，然后我们会把这些任务做异步的处理，但是又要及时给客户端反馈任务的处理进度。
对于这种场景，我们可以使用 WebSocket 来实现。其他可以使用 WebSocket 进行通知的场景还有像管理后台一些通知（比如新订单通知）等。
在本篇文章中，就是要实现一个这样的消息推送系统，具体来说，它会有以下功能：
- 可以给特定的用户推送：建立连接的时候，就建立起
WebSocket连接与用户 ID 之间的关联 - 断开连接的时候，移除
WebSocket连接与用户的关联，并且关闭这个WebSocket连接 - 业务系统可以通过 HTTP 接口来给特定的用户推送
WebSocket消息：只要传递用户 ID 以及需要推送的消息即可
基础框架
下面是一个最简单版本的框架图：
它包含如下几个角色：
Client客户端，也就是实际中接收消息通知的浏览器Server服务端，在我们的例子中，服务端实际不处理业务逻辑，只处理跟客户端的消息交互：维持WebSocket连接，推送消息到特定的WebSocket连接- 业务逻辑：这个实际上不属于 demo 的一部分，但是
Server推送的数据是来自业务逻辑处理的结果
设计成这样的目的是为了将技术跟业务进行分离，业务逻辑上的变化不影响到底层技术，同样的，WebSocket 推送中心的技术上的变动也不会影响到实际的业务。
开始开发
一些结构体变动
Client结构体的变化
type Client struct {
hub *Hub
conn *websocket.Conn
send chan []byte
// 新增字段
uid int
}
因为我们需要建立起 WebSocket 连接与用户之间的关联，因此我们需要一个额外的字段来记录用户 ID，也就是上面的 uid 字段。
这个字段会在客户端建立连接后写入。
Hub结构体的变化
type Hub struct {
clients map[*Client]bool
register chan *Client
unregister chan *Client
// 记录 uid 跟 client 的对应关系
userClients map[int]*Client
// 读写锁，保护 userClients 以及 clients 的读写
sync.RWMutex
}
- 因为我们不再需要做广播，所以会移除
Hub中的broadcast字段。
取而代之的是，我们会直接在消息推送接口中写入到 uid 对应的 Client 的 send 通道。
当然我们也可以在 Hub 中另外加一个字段来记录要推送给不同 uid 的消息，但是我们的 Hub 的 run 方法是一个协程处理的，当需要推送的数据较多或者其中有
网络延迟的时候，会直接影响到推送给其他用户的消息。当然我们也可以改造一下 run 方法，启动多个协程来处理，不过这样比较复杂，本文会在 writePump 中处理。
（也就是建立 WebSocket 连接时的那个写操作协程）
- 同时为了更加快速地通过
uid来获取对应的WebSocket连接，新增了一个userClients字段。
这是一个 map 类型的字段，key 是 uid，值是对应的 Client 指针。
- 最后新增了一个
Mutex互斥锁
因为，在用户实际进行登录的时候需要写入 userClients 字段，而这是一个 map 类型字段，并不支持并发读写。
如果我们在接受并发连接的时候同时修改 userClients 的时候会导致 panic，因此我们使用了一个互斥锁来保证 userClients 的读写安全。
同时，clients 也是一个 map，但上一篇文章中没有使用 sync.Mutex 来保护它的读写，在并发操作的时候也是会有问题的，
所以 Mutex 同时也需要保护 clients 的读写。
func (h *Hub) run() {
for {
select {
case client := <-h.register:
h.Lock()
h.clients[client] = true
h.Unlock()
case client := <-h.unregister:
if _, ok := h.clients[client]; ok {
h.Lock()
delete(h.userClients, client.uid)
delete(h.clients, client)
h.Unlock()
close(client.send)
}
}
}
}
最后，我们会在 Hub 的 run 方法中写 userClients 或者 clients 字段的时候，先获取锁，写成功的时候释放锁。
建立连接
在本篇中，将会继续沿用上一篇的代码，只是其中一些细节会有所改动。建立连接这步操作，跟上一篇的一样：
// 将 HTTP 转换为 WebSocket 连接的 Upgrader
var upgrader = websocket.Upgrader{
ReadBufferSize: 1024,
WriteBufferSize: 1024,
}
// 处理 WebSocket 连接请求
func serveWs(hub *Hub, w http.ResponseWriter, r *http.Request) {
// 升级为 WebSocket 连接
conn, err := upgrader.Upgrade(w, r, nil)
if err != nil {
log.Println(err)
return
}
// 新建一个 Client
client := &Client{hub: hub, conn: conn, send: make(chan []byte, 256)}
// 注册到 Hub
client.hub.register <- client
// 推送消息的协程
go client.writePump()
// 结束消息的协程
go client.readPump()
}
接收消息
由于我们要做的只是一个推送消息的系统，所以我们只处理用户发来的登录请求，其他的消息会全部丢弃：
func (c *Client) readPump() {
defer func() {
c.hub.unregister <- c
_ = c.conn.Close()
}()
c.conn.SetReadLimit(maxMessageSize)
c.conn.SetReadDeadline(time.Time{}) // 永不超时
for {
// 从客户端接收消息
_, message, err := c.conn.ReadMessage()
if err != nil {
log.Println("readPump error: ", err)
break
}
// 只处理登录消息
var data = make(map[string]string)
err = json.Unmarshal(message, &data)
if err != nil {
break
}
// 写入 uid 以及 Hub 的 userClients
if uid, ok := data["uid"]; ok {
c.uid = uid
c.hub.Lock()
c.hub.userClients[uid] = c
c.hub.Unlock()
}
}
}
在本文中，假设客户端的登录消息格式为 {"uid": "123456"} 这种 json 格式。
在这里也操作了
userClients字段，同样需要使用互斥锁来保证操作的安全性。
发送消息
- 在我们的系统中，可以提供一个 HTTP 接口来跟业务系统进行交互：
// 发送消息的接口
// 参数：
// 1. uid：接收消息的用户 ID
// 2. message：需要发送给这个用户的消息
http.HandleFunc("/send", func(w http.ResponseWriter, r *http.Request) {
send(hub, w, r)
})
// 发送消息的方法
func send(hub *Hub, w http.ResponseWriter, r *http.Request) {
uid := r.FormValue("uid")
// 参数错误
if uid == "" {
w.WriteHeader(http.StatusBadRequest)
return
}
// 从 hub 中获取 client
hub.Lock()
client, ok := hub.userClients[uid]
hub.Unlock()
// 尚未建立连接
if !ok {
w.WriteHeader(http.StatusBadRequest)
return
}
// 发送消息
message := r.FormValue("message")
client.send <- []byte(message)
}
- 实际发送消息的操作
在 writePump 方法中，我们会将从 /send 接收到的数据发送给对应的用户：
// 发送消息的协程
func (c *Client) writePump() {
defer func() {
_ = c.conn.Close()
}()
for {
select {
case message, ok := <-c.send:
// 设置写超时时间
c.conn.SetWriteDeadline(time.Now().Add(writeWait))
// 连接已经被关闭了
if !ok {
c.conn.WriteMessage(websocket.CloseMessage, []byte{})
return
}
// 获取一个发送消息的 Writer
w, err := c.conn.NextWriter(websocket.TextMessage)
if err != nil {
return
}
// 写入消息到 Writer
w.Write(message)
// 关闭 Writer
if err := w.Close(); err != nil {
return
}
}
}
}
在这个方法中，我们会从 c.send 这个 chan 中获取需要发送给客户端的消息，然后进行发送操作。
测试
- 启动
main程序
go run main.go
- 打开一个浏览器的控制台，执行以下代码
ws = new WebSocket('ws://127.0.0.1:8181/ws')
ws.send('{"uid": "123"}')
这两行代码的作用是与 WebSocket 服务器建立连接，然后发送一个登录信息。
然后我们打开控制台的 Network -> WS -> Message 就可以看到浏览器发给服务端的消息：
- 使用 HTTP 客户端发送消息给 uid 为 123 的用户
假设我们的
WebSocket服务器绑定的端口为8181
打开终端，执行以下命令：
curl "http://localhost:8181/send?uid=123&message=Hello%20World"
然后我们可以在 Network -> WS -> Message 看到接收到了消息 Hello World。
结束了
到此为止，我们已经实现了一个初步可工作的 WebSocket 应用，当然还有很多可以优化的地方，
比如：
- 错误处理
Hub状态目前对外部来说是一个黑盒子，我们可以加个接口返回一下Hub的当前状态，比如当前连接数- 日志：出错的时候，日志可以帮助我们快速定位问题
这些功能会在后续继续完善，今天就到此为止了。
