使用 Go 语言实现 WebSocket的核心逻辑

文章目录

- WebSocket 简介
- 时序图
- 核心逻辑
- - Client 结构与功能
  - 创建新客户端
  - 消息读取逻辑 (ReadPump)
  - 发送消息逻辑 (Send)
  - 客户端管理器 (ClientManager)
  - WebSocket 处理器
  - 处理心跳与长连接
- 总结

本文将基于 Go 语言，通过使用 gorilla/websocket 库来实现一个简单的聊天应用。该应用具备处理 WebSocket 连接、消息传输、以及用户连接管理等功能。我们将详细展示如何实现这些功能，并剖析背后的核心逻辑与原理。

WebSocket 简介

WebSocket 是一种全双工的通信协议，允许客户端和服务器之间在一个持久连接上进行双向数据传输。与 HTTP 的短连接不同，WebSocket 可以在建立连接后保持连接状态，从而实现实时通信。因此，WebSocket 非常适合用于聊天应用等需要实时数据传输的场景。

时序图

核心逻辑

在本示例中，我们主要实现了以下几个核心模块：

Client：表示单个 WebSocket 连接的客户端，负责处理消息的收发。
ClientManager：用于管理多个客户端的连接，处理客户端的增加、删除以及消息的路由。
WebSocket 处理逻辑：处理新连接的建立、消息的读取与发送。

Client 结构与功能

type Client struct {conn         *websocket.ConnmessageQueue chan []bytemu           sync.Mutexuser         string
}

Client 结构体用于表示一个 WebSocket 客户端连接。每个客户端包含：

conn：当前的 WebSocket 连接。
messageQueue：用于存储待发送的消息队列。
mu：用于保证并发安全的互斥锁。
user：表示客户端的用户标识。

创建新客户端

func NewClient(user string, conn *websocket.Conn) *Client {return &Client{conn:         conn,user:         user,messageQueue: make(chan []byte, 100),}
}

NewClient 函数用于创建新的客户端实例。每个客户端都有一个独立的消息队列，用于存储要发送给客户端的消息。

消息读取逻辑 (ReadPump)

func (c *Client) ReadPump() {defer func() {c.conn.Close()}()for {mt, message, err := c.conn.ReadMessage()if err != nil {log.Println("read:", err)manager.mu.Lock()delete(manager.clients, c.user)_ = c.conn.Close()manager.mu.Unlock()break}if mt == websocket.TextMessage || mt == websocket.PingMessage {c.mu.Lock()c.messageQueue <- messagec.mu.Unlock()}}
}

ReadPump 方法用于持续从 WebSocket 连接中读取消息，并将接收到的消息存储到 messageQueue 队列中。该方法通过一个无限循环，不断读取 WebSocket 的消息。当出现错误时，例如客户端断开连接，便会关闭当前连接并将该客户端从客户端管理器中移除。

其中，ReadMessage() 方法用于从 WebSocket 连接中读取消息，返回的 mt 表示消息类型。常见的类型包括文本消息（TextMessage）和 ping 消息（PingMessage）。对于这些消息类型，消息会被推送到 messageQueue 以便后续处理。

发送消息逻辑 (Send)

func Send(user string, returnMessage []byte, logger logx.Logger) {manager.mu.RLock()client, exists := manager.clients[user]manager.mu.RUnlock()if !exists {logger.Infof("client not found for user:%s message:%s", user, string(returnMessage))return}client.mu.Lock()err := client.conn.WriteMessage(websocket.TextMessage, returnMessage)client.mu.Unlock()if err != nil {logger.Errorf("client.conn.WriteMessage error %s", err.Error())manager.mu.Lock()delete(manager.clients, user)manager.mu.Unlock()_ = client.conn.Close()}
}

Send 函数负责向指定的用户发送消息。首先，它会检查用户是否存在于 ClientManager 中，如果不存在则记录日志并返回。如果用户存在，则通过 WriteMessage() 方法将消息发送给客户端。若发送消息时发生错误，会将该用户从连接管理器中移除，并关闭该 WebSocket 连接。

客户端管理器 (ClientManager)

type ClientManager struct {clients map[string]*Clientmu      sync.RWMutex
}var manager = ClientManager{clients: make(map[string]*Client),
}

ClientManager 用于管理多个客户端的连接，clients 字段是一个存储所有客户端连接的映射，键是用户标识，值是客户端对象。通过读写锁 (sync.RWMutex)，确保在并发访问时的线程安全。

WebSocket 处理器

ChatWebsocketHandler 是处理 WebSocket 连接的 HTTP 处理函数。

func ChatWebsocketHandler(svcCtx *svc.ServiceContext) http.HandlerFunc {return func(w http.ResponseWriter, r *http.Request) {conn, err := upgrader.Upgrade(w, r, nil)logger := logx.WithContext(r.Context())if err != nil {logger.Errorf("upgrade:%+v", err)return}user := r.URL.Query().Get("user")if user == "" {logger.Errorf("user is empty:")_ = conn.Close()return}client := NewClient(user, conn)manager.mu.Lock()oldClient, exists := manager.clients[user]if exists {_ = oldClient.conn.Close()}manager.clients[user] = clientmanager.mu.Unlock()go client.ReadPump()// 省略其他消息处理逻辑...}
}

连接升级：首先使用 upgrader.Upgrade() 将 HTTP 请求升级为 WebSocket 连接。
用户认证：通过 URL 查询参数获取用户 ID，并创建对应的 Client。
旧连接处理：如果该用户已经有一个旧的 WebSocket 连接，则会关闭旧连接。
启动消息读取：通过启动 ReadPump() 协程，持续读取该用户的 WebSocket 消息。

处理心跳与长连接

在 WebSocket 通信中，维持长连接的一个常用做法是使用心跳机制。

if req.Heartbeat {// 处理心跳消息err = client.conn.WriteMessage(websocket.PongMessage, []byte(""))if err != nil {logger.Errorf("write pong message failed:", err)manager.mu.Lock()delete(manager.clients, user)manager.mu.Unlock()_ = client.conn.Close()return}
}