第10章:WebSocket网关核心功能实现——从握手到广播,手把手搞定实时通信
前言
你有没有遇到过这种场景:老板说"咱们的即时通讯系统要支持十万并发在线",你信心满满地用HTTP轮询搭了一套架构,上线第一天CPU直接拉满,用户疯狂投诉消息延迟。你盯着Grafana面板上那条直奔云霄的CPU曲线,开始怀疑人生。
或者更经典一点:你用第三方SDK做了一套推送系统,某个深夜第三方服务挂了,几百万用户收不到消息,你的手机被投诉电话打爆,而你连排查的权限都没有。
这些痛点的根源只有一个——你没有自己掌控 WebSocket 网关。
我是怕浪猫,一个在后端实时通信领域踩过无数坑的老兵。从最早用Netty手写长连接网关,到后来用Go重构了三版WebSocket网关,我踩过的坑足够填平一个西湖。今天这一章,我把WebSocket网关的核心功能实现从头到尾拆给你看,每一行代码都经过生产环境验证,每一个设计决策都标注了原因。
网关不是写出来的,是改出来的。第一版能跑的代码和能扛十万并发的代码,中间隔着一百次线上事故。
一、WebSocket握手与连接建立
1.1 为什么不直接用HTTP
很多人第一次做实时通信都会想:我每秒轮询一次不行吗?答案是不行,原因有三:
第一,HTTP每次请求都要建立TCP连接(除非用Keep-Alive),头部开销巨大,一个简单的"有没有新消息"的请求,HTTP头部可能就有800字节,而有效载荷只有1个比特。
第二,服务器无法主动推送。客户端不问,服务器就没法回答,消息延迟最坏情况等于轮询间隔。
第三,并发连接数爆炸。一万个用户每秒轮询一次,就是一万个QPS,而实际上90%的请求返回的都是"没有新消息"。
WebSocket解决了所有这些问题:一次握手升级协议后,TCP连接保持长开,双向通信,服务器可以随时推送,头部开销几乎为零。
1.2 握手过程详解
WebSocket的握手本质上是一个HTTP GET请求,带了一个Upgrade: websocket头部。服务器如果支持WebSocket,返回101 Switching Protocols,之后这条TCP连接就不再是HTTP了,而是WebSocket协议。
来看握手请求的关键头部:
GET /ws?token=abc123&uid=10086 HTTP/1.1
Host: gateway.example.com
Upgrade: websocket
Connection: Upgrade
Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==
Sec-WebSocket-Version: 13服务器响应:
HTTP/1.1 101 Switching Protocols
Upgrade: websocket
Connection: Upgrade
Sec-WebSocket-Accept: s3pPLMBiTxaQ9kYGzzhZRbK+xOo=其中Sec-WebSocket-Accept的值是服务器用客户端发来的Sec-WebSocket-Key加上一个固定魔法字符串258EAFA5-E914-47DA-95CA-C5AB0DC85B11,做SHA1哈希再Base64编码得到的。这个过程不是加密,只是确认双方都懂WebSocket协议。
握手是WebSocket的门面,门面都搭不好,后面的长连接就是空中楼阁。
1.3 Go语言实现握手
Go标准库的net/http从1.0版本就内置了WebSocket支持(通过golang.org/x/net/websocket),但那个包太简陋了,生产环境我用的是gorilla/websocket,这是Go生态里最成熟的WebSocket库。
先看最基础的握手实现:
package gateway
import (
"log"
"net/http"
"time"
"github.com/gorilla/websocket"
)
var upgrader = websocket.Upgrader{
// 读缓冲区大小,根据你的消息体大小调整
ReadBufferSize: 4096,
// 写缓冲区大小
WriteBufferSize: 4096,
// 握手超时时间
HandshakeTimeout: 10 * time.Second,
// 检查Origin,生产环境一定要校验
CheckOrigin: func(r *http.Request) bool {
// 允许的域名列表
allowedOrigins := map[string]bool{
"https://app.example.com": true,
"https://web.example.com": true,
}
origin := r.Header.Get("Origin")
if origin == "" {
// 非浏览器客户端(如SDK),允许通过
return true
}
return allowedOrigins[origin]
},
}
// HandleWS 是WebSocket握手的入口
func HandleWS(w http.ResponseWriter, r *http.Request) {
// 第一步:参数校验
token := r.URL.Query().Get("token")
uid := r.URL.Query().Get("uid")
if token == "" || uid == "" {
http.Error(w, "missing token or uid", http.StatusUnauthorized)
return
}
// 第二步:鉴权(这里简化,实际要查Redis或DB)
userID, err := authenticate(token)
if err != nil {
http.Error(w, "invalid token", http.StatusUnauthorized)
return
}
// 第三步:完成WebSocket握手
conn, err := upgrader.Upgrade(w, r, nil)
if err != nil {
log.Printf("upgrade failed: uid=%s, err=%v", userID, err)
return
}
// 注意:Upgrade之后,w这个ResponseWriter就不能再用了
// 因为协议已经切换,TCP连接已经交给WebSocket处理器
// 第四步:设置连接参数
conn.SetReadLimit(64 * 1024) // 最大消息体64KB
conn.SetReadDeadline(time.Now().Add(60 * time.Second))
conn.SetPongHandler(func(string) error {
// 收到Pong,刷新读超时
conn.SetReadDeadline(time.Now().Add(60 * time.Second))
return nil
})
// 第五步:将连接交给连接管理器
client := NewClient(userID, conn)
Hub.Register <- client
// 第六步:启动读写协程
go client.readPump()
go client.writePump()
}
func authenticate(token string) (string, error) {
// 实际项目里这里要查Redis或调用鉴权服务
// 简化示例:假设token格式为 "uid:signature"
// 生产环境一定要用JWT或OAuth2
if len(token) < 10 {
return "", fmt.Errorf("token too short")
}
return token[:len(token)-10], nil
}这里有几个坑需要特别注意:
坑一:CheckOrigin不要直接返回true。 我见过太多项目在生产环境被跨域攻击的,就是因为图省事写了CheckOrigin: func(r *http.Request) bool { return true }。一定要维护一个域名白名单。
坑二:握手超时一定要设。 默认情况下Upgrade没有超时,恶意客户端可以连上来不发任何数据,占着连接不干活。设置HandshakeTimeout能在握手阶段就掐掉这些连接。
坑三:Upgrade之后不要操作ResponseWriter。 Upgrade函数内部已经往w里写了101 Switching Protocols响应,之后再操作w会导致连接混乱。如果你要返回错误,在Upgrade之前返回。
安全不是功能上线后加的补丁,而是架构设计时埋的地基。
1.4 连接建立后的初始化
握手成功只是第一步,连接建立后还有一系列初始化工作要做。我把这个过程叫做"连接生命周期初始化":
// Client 表示一个WebSocket客户端连接
type Client struct {
UserID string
Conn *websocket.Conn
Send chan []byte
Hub *Hub
// 连接创建时间
ConnectedAt time.Time
// 最后活跃时间
LastActive time.Time
// 客户端设备信息
Platform string // ios, android, web, pc
Version string // 客户端版本号
}
func NewClient(userID string, conn *websocket.Conn) *Client {
return &Client{
UserID: userID,
Conn: conn,
Send: make(chan []byte, 256), // 发送队列缓冲256条消息
ConnectedAt: time.Now(),
LastActive: time.Now(),
}
}
// readPump 负责从连接读取消息
func (c *Client) readPump() {
defer func() {
c.Hub.Unregister <- c
c.Conn.Close()
}()
for {
_, message, err := c.Conn.ReadMessage()
if err != nil {
if websocket.IsUnexpectedCloseError(err,
websocket.CloseGoingAway,
websocket.CloseNormalClosure) {
log.Printf("read error: uid=%s, err=%v", c.UserID, err)
}
break
}
c.LastActive = time.Now()
// 将消息交给路由处理器
c.Hub.Route <- &MessageEvent{
Client: c,
Payload: message,
}
}
}
// writePump 负责向连接写入消息
func (c *Client) writePump() {
ticker := time.NewTicker(30 * time.Second) // 心跳定时器
defer func() {
ticker.Stop()
c.Conn.Close()
}()
for {
select {
case message, ok := <-c.Send:
if !ok {
// Send通道关闭,说明Hub要求断开连接
c.Conn.WriteMessage(websocket.CloseMessage, []byte{})
return
}
// 设置写超时
c.Conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
if err := c.Conn.WriteMessage(websocket.TextMessage, message); err != nil {
log.Printf("write error: uid=%s, err=%v", c.UserID, err)
return
}
case <-ticker.C:
// 发送Ping心跳包
c.Conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
if err := c.Conn.WriteMessage(websocket.PingMessage, nil); err != nil {
log.Printf("ping error: uid=%s, err=%v", c.UserID, err)
return
}
}
}
}这里有一个关键设计:读写分离双协程。readPump和writePump各跑一个goroutine,互不干扰。为什么这么设计?因为WebSocket是全双工的,读和写可以同时进行。如果用单协程处理读写,一个慢消息会阻塞整个连接的处理。
Send通道的缓冲大小设为256,这是一个经验值。如果你的业务场景消息突发量大(比如群聊消息广播),可以调大。但不要设太大,否则消息积压会吃光内存。更好的做法是配合背压机制,当队列满到80%时丢弃低优先级消息。
通道不是越大越好,缓冲大小是你对消息积压的容忍度。积压超过容忍度,就该丢而不是存。
二、消息协议设计(JSON/Protobuf)
2.1 协议设计的重要性
协议是WebSocket网关的灵魂。很多团队在协议设计上栽的跟头比在代码逻辑上栽的还多。我见过一个团队,消息格式今天用JSON,明天加个字段,后天改个字段名,半年后客户端版本碎片化严重,新老消息格式混在一起,整个系统乱成一锅粥。
协议设计的核心原则:向前兼容、高效编解码、可扩展。
2.2 JSON协议设计
JSON是最简单的协议格式,调试方便,可读性好。对于内部系统或中小规模应用,JSON完全够用。
我推荐的JSON消息格式:
// Message 是WebSocket通信的顶层消息结构
type Message struct {
Header MessageHeader `json:"header"`
Body json.RawMessage `json:"body"` // 延迟解析,提高性能
}
type MessageHeader struct {
// 消息ID,用于消息确认和去重
MsgID string `json:"msg_id"`
// 消息类型:chat/notice/push/ack/heartbeat
Type string `json:"type"`
// 消息子类型,如chat下的text/image/voice
SubType string `json:"sub_type,omitempty"`
// 发送者ID
From string `json:"from"`
// 接收者ID(单聊为用户ID,群聊为群ID)
To string `json:"to"`
// 时间戳(毫秒)
Timestamp int64 `json:"timestamp"`
// 协议版本
Version string `json:"version"`
// 扩展字段
Extra map[string]string `json:"extra,omitempty"`
}
// ChatMessage 是聊天消息体
type ChatMessage struct {
Content string `json:"content"`
MsgType string `json:"msg_type"` // text/image/voice/video
ReplyTo string `json:"reply_to,omitempty"`
Mentioned []string `json:"mentioned,omitempty"`
}
// AckMessage 是确认消息体
type AckMessage struct {
Code int `json:"code"`
Message string `json:"message"`
RefMsgID string `json:"ref_msg_id"` // 关联的消息ID
}为什么Body用json.RawMessage而不是interface{}?因为interface{}会立即反序列化整个JSON,而RawMessage是延迟解析的。你可以先解析Header,根据Type字段决定要不要解析Body,以及用什么结构体解析。这样能显著减少不必要的反序列化开销。
协议设计的第一原则是向前兼容。今天加的字段,明天不能删;今天选的格式,后天不能换。除非你想写三版客户端兼容代码。
来看一个完整的JSON消息编解码示例:
// EncodeMessage 将消息编码为JSON字节流
func EncodeMessage(msg *Message) ([]byte, error) {
return json.Marshal(msg)
}
// DecodeMessage 将JSON字节流解码为消息
func DecodeMessage(data []byte) (*Message, error) {
msg := &Message{}
if err := json.Unmarshal(data, msg); err != nil {
return nil, fmt.Errorf("unmarshal message: %w", err)
}
return msg, nil
}
// ParseBody 根据消息类型解析消息体
func (m *Message) ParseBody() (interface{}, error) {
switch m.Header.Type {
case "chat":
var body ChatMessage
return &body, json.Unmarshal(m.Body, &body)
case "ack":
var body AckMessage
return &body, json.Unmarshal(m.Body, &body)
case "heartbeat":
return nil, nil
default:
return nil, fmt.Errorf("unknown message type: %s", m.Header.Type)
}
}
// 使用示例
func handleChatMessage(raw []byte) error {
msg, err := DecodeMessage(raw)
if err != nil {
return err
}
body, err := msg.ParseBody()
if err != nil {
return err
}
chatMsg, ok := body.(*ChatMessage)
if !ok {
return fmt.Errorf("invalid body type for chat message")
}
log.Printf("chat from %s to %s: %s",
msg.Header.From, msg.Header.To, chatMsg.Content)
return nil
}2.3 Protobuf协议设计
当你的消息量级达到十万QPS以上时,JSON的序列化/反序列化开销就会成为瓶颈。Protobuf的优势在于:
- 编解码速度比JSON快3-5倍
- 二进制体积比JSON小30%-50%
- 强类型,编译时就能发现字段错误
- 天然支持向前兼容(新字段加optional即可)
先定义proto文件:
// proto/gateway.proto
syntax = "proto3";
package gateway;
option go_package = "github.com/example/gateway/proto;gatewaypb";
// 顶层消息结构
message WsMessage {
MessageHeader header = 1;
bytes body = 2; // 消息体,根据header.type解析
}
message MessageHeader {
string msg_id = 1;
string type = 2; // chat/notice/push/ack/heartbeat
string sub_type = 3;
string from = 4;
string to = 5;
int64 timestamp = 6;
string version = 7;
map<string, string> extra = 8;
}
message ChatMessage {
string content = 1;
string msg_type = 2; // text/image/voice/video
string reply_to = 3;
repeated string mentioned = 4;
}
message AckMessage {
int32 code = 1;
string message = 2;
string ref_msg_id = 3;
}
// 心跳消息
message HeartbeatMessage {
int64 client_time = 1;
}Go代码中使用Protobuf:
package gateway
import (
"google.golang.org/protobuf/proto"
"github.com/example/gateway/proto"
)
// EncodePbMessage 将Protobuf消息编码为字节流
func EncodePbMessage(msg *gatewaypb.WsMessage) ([]byte, error) {
return proto.Marshal(msg)
}
// DecodePbMessage 将字节流解码为Protobuf消息
func DecodePbMessage(data []byte) (*gatewaypb.WsMessage, error) {
msg := &gatewaypb.WsMessage{}
if err := proto.Unmarshal(data, msg); err != nil {
return nil, fmt.Errorf("unmarshal pb message: %w", err)
}
return msg, nil
}
// BuildChatMessage 构建一条聊天消息
func BuildChatMessage(from, to, content string) *gatewaypb.WsMessage {
chatBody := &gatewaypb.ChatMessage{
Content: content,
MsgType: "text",
}
bodyBytes, _ := proto.Marshal(chatBody)
return &gatewaypb.WsMessage{
Header: &gatewaypb.MessageHeader{
MsgId: generateMsgID(),
Type: "chat",
From: from,
To: to,
Timestamp: time.Now().UnixMilli(),
Version: "1.0",
},
Body: bodyBytes,
}
}
func generateMsgID() string {
return fmt.Sprintf("msg_%d_%d", time.Now().UnixNano(), rand.Intn(10000))
}2.4 JSON与Protobuf的混合方案
实际生产中,我不建议一刀切。最合理的方案是:握手阶段和控制消息用JSON,业务消息用Protobuf。
原因很简单:握手阶段需要调试,用JSON方便抓包看问题;业务消息量大,用Protobuf省带宽省CPU。
实现方式是在WebSocket的消息类型上做区分。WebSocket协议本身定义了四种消息类型:Text(1)、Binary(2)、Close(8)、Ping(9)、Pong(10)。我们用Text消息传JSON,Binary消息传Protobuf。
// 根据消息类型选择编码方式
func (c *Client) readPump() {
for {
messageType, message, err := c.Conn.ReadMessage()
if err != nil {
break
}
switch messageType {
case websocket.TextMessage:
// JSON消息,通常是控制类消息
c.handleJSONMessage(message)
case websocket.BinaryMessage:
// Protobuf消息,通常是业务消息
c.handlePBMessage(message)
case websocket.PingMessage:
c.LastActive = time.Now()
case websocket.PongMessage:
c.LastActive = time.Now()
}
}
}
func (c *Client) handleJSONMessage(data []byte) {
msg, err := DecodeMessage(data)
if err != nil {
log.Printf("decode json error: %v", err)
return
}
switch msg.Header.Type {
case "login":
c.handleLogin(msg)
case "subscribe":
c.handleSubscribe(msg)
case "heartbeat":
c.handleHeartbeat(msg)
default:
log.Printf("unknown json message type: %s", msg.Header.Type)
}
}
func (c *Client) handlePBMessage(data []byte) {
msg, err := DecodePbMessage(data)
if err != nil {
log.Printf("decode pb error: %v", err)
return
}
switch msg.Header.Type {
case "chat":
c.handleChatPB(msg)
case "ack":
c.handleAckPB(msg)
default:
log.Printf("unknown pb message type: %s", msg.Header.Type)
}
}混合协议不是妥协,是工程上的最优解。调试友好和运行效率从来不是非此即彼的选择。
2.5 协议版本管理
协议一定会演进,版本管理必须从一开始就设计好。我踩过最深的坑就是没有版本号,上线后老客户端收到新格式消息直接崩溃。
// 版本兼容性检查
var supportedVersions = map[string]bool{
"1.0": true,
"1.1": true,
"2.0": true,
}
func checkVersion(version string) bool {
return supportedVersions[version]
}
// 按版本号路由处理逻辑
func (m *Message) ParseBodyByVersion() (interface{}, error) {
switch m.Header.Version {
case "1.0":
return m.parseBodyV1()
case "1.1":
return m.parseBodyV11()
case "2.0":
return m.parseBodyV2()
default:
return nil, fmt.Errorf("unsupported version: %s", m.Header.Version)
}
}
// V1.0的ChatMessage没有mentioned字段
func (m *Message) parseBodyV1() (interface{}, error) {
type ChatMessageV1 struct {
Content string `json:"content"`
MsgType string `json:"msg_type"`
}
var body ChatMessageV1
err := json.Unmarshal(m.Body, &body)
return &body, err
}
// V1.1增加了mentioned字段
func (m *Message) parseBodyV11() (interface{}, error) {
type ChatMessageV11 struct {
Content string `json:"content"`
MsgType string `json:"msg_type"`
Mentioned []string `json:"mentioned"`
}
var body ChatMessageV11
err := json.Unmarshal(m.Body, &body)
return &body, err
}三、连接池管理与心跳机制
3.1 连接池设计
连接池是WebSocket网关的核心组件。一个设计良好的连接池要解决以下问题:
- 快速查找:通过UserID快速找到连接
- 并发安全:多个goroutine同时读写连接池
- 连接迁移:用户重连时,旧连接要优雅关闭
- 状态监控:实时知道当前在线连接数
package gateway
import (
"sync"
"time"
)
// Hub 是连接池的中心管理器
type Hub struct {
// 在线连接 map[UserID]*Client
clients map[string]*Client
// 用户的多端连接 map[UserID]map[DeviceID]*Client
// 支持一个用户同时在线多个设备
devices map[string]map[string]*Client
// 注册通道
Register chan *Client
// 注销通道
Unregister chan *Client
// 消息路由通道
Route chan *MessageEvent
// 广播通道
Broadcast chan []byte
mu sync.RWMutex
// 统计信息
stats *HubStats
}
type HubStats struct {
TotalConnections int64
ActiveConnections int64
MessagesReceived int64
MessagesSent int64
}
type MessageEvent struct {
Client *Client
Payload []byte
}
func NewHub() *Hub {
return &Hub{
clients: make(map[string]*Client),
devices: make(map[string]map[string]*Client),
Register: make(chan *Client, 100),
Unregister: make(chan *Client, 100),
Route: make(chan *MessageEvent, 1000),
Broadcast: make(chan []byte, 100),
stats: &HubStats{},
}
}
// Run 是Hub的主循环,所有操作都通过channel串行化
func (h *Hub) Run() {
ticker := time.NewTicker(30 * time.Second)
defer ticker.Stop()
for {
select {
case client := <-h.Register:
h.handleRegister(client)
case client := <-h.Unregister:
h.handleUnregister(client)
case event := <-h.Route:
h.handleRoute(event)
case message := <-h.Broadcast:
h.handleBroadcast(message)
case <-ticker.C:
h.cleanupStaleConnections()
}
}
}连接池不是Map那么简单,它是一个有生命周期的生态系统。注册、注销、清理、迁移,每一步都要严丝合缝。
3.2 注册与注销逻辑
func (h *Hub) handleRegister(client *Client) {
h.mu.Lock()
defer h.mu.Unlock()
// 处理同用户的旧连接(单设备登录场景)
if oldClient, exists := h.clients[client.UserID]; exists {
// 给旧连接发一个踢人消息
oldClient.Send <- []byte(`{"header":{"type":"kick","msg_id":"system"},"body":{"code":1001,"message":"account login elsewhere"}}`)
// 关闭旧连接的Send通道,writePump会自动退出
close(oldClient.Send)
delete(h.clients, client.UserID)
}
// 注册新连接
h.clients[client.UserID] = client
h.stats.TotalConnections++
h.stats.ActiveConnections = int64(len(h.clients))
log.Printf("client registered: uid=%s, total=%d",
client.UserID, h.stats.ActiveConnections)
}
func (h *Hub) handleUnregister(client *Client) {
h.mu.Lock()
defer h.mu.Unlock()
if existing, exists := h.clients[client.UserID]; exists {
// 确保是同一个连接(可能已经被新连接替换了)
if existing == client {
delete(h.clients, client.UserID)
h.stats.ActiveConnections = int64(len(h.clients))
log.Printf("client unregistered: uid=%s, total=%d",
client.UserID, h.stats.ActiveConnections)
}
}
// 安全关闭通道(避免重复关闭导致panic)
// 注意:这里不关闭Send通道,因为可能在handleRegister时已经关闭了
}这里有一个极其重要的坑:重复关闭channel会panic。我见过一个线上事故,就是 unregister 被调用了两次,第二次close了一个已经关闭的channel,直接panic导致整个网关进程崩溃。
解决方案是使用sync.Once:
type Client struct {
// ... 其他字段
closeOnce sync.Once
}
func (c *Client) Close() {
c.closeOnce.Do(func() {
close(c.Send)
c.Conn.Close()
})
}3.3 多设备同时在线
很多IM系统都需要支持一个用户多设备同时在线(比如手机和电脑都登录)。这时候连接池的结构要改:
// 多设备连接管理
func (h *Hub) handleRegisterMultiDevice(client *Client) {
h.mu.Lock()
defer h.mu.Unlock()
deviceID := client.Platform + "_" + client.Version // 简单的设备标识
// 如果该用户还没有任何连接,创建设备map
if h.devices[client.UserID] == nil {
h.devices[client.UserID] = make(map[string]*Client)
}
// 如果同设备已有连接,踢掉旧的
if oldClient, exists := h.devices[client.UserID][deviceID]; exists {
oldClient.Send <- []byte(`{"header":{"type":"kick","msg_id":"system"},"body":{"code":1002,"message":"same device reconnected"}}`)
oldClient.Close()
}
// 注册新连接
h.devices[client.UserID][deviceID] = client
h.stats.TotalConnections++
h.stats.ActiveConnections = h.countConnections()
}
func (h *Hub) countConnections() int64 {
var count int64
for _, devices := range h.devices {
count += int64(len(devices))
}
return count
}
// GetClient 获取用户的所有在线设备
func (h *Hub) GetClients(userID string) []*Client {
h.mu.RLock()
defer h.mu.RUnlock()
devices := h.devices[userID]
if devices == nil {
return nil
}
clients := make([]*Client, 0, len(devices))
for _, c := range devices {
clients = append(clients, c)
}
return clients
}3.4 心跳机制详解
心跳是长连接的命脉。没有心跳,你根本不知道对面的连接是活着还是已经死了(比如用户手机崩溃、网络中断等"半开连接"场景)。
心跳机制的设计步骤如下:
步骤一:确定心跳间隔
- 服务端Ping间隔:30秒
- 客户端Pong超时:60秒(给客户端足够时间回复)
- 读超时(无任何消息):90秒(两个心跳周期+冗余)
// 心跳参数
const (
PingInterval = 30 * time.Second
PongWait = 60 * time.Second
ReadTimeout = 90 * time.Second
WriteWait = 10 * time.Second
MaxMissedPongs = 2 // 允许丢失的最大Pong次数
)步骤二:服务端心跳实现
func (c *Client) readPump() {
defer func() {
c.Hub.Unregister <- c
c.Close()
}()
// 设置读超时
c.Conn.SetReadDeadline(time.Now().Add(ReadTimeout))
// Pong处理器:收到Pong时刷新读超时
c.Conn.SetPongHandler(func(string) error {
c.Conn.SetReadDeadline(time.Now().Add(ReadTimeout))
c.LastActive = time.Now()
return nil
})
for {
_, message, err := c.Conn.ReadMessage()
if err != nil {
if websocket.IsUnexpectedCloseError(err,
websocket.CloseGoingAway,
websocket.CloseNormalClosure) {
log.Printf("read error: uid=%s, err=%v", c.UserID, err)
}
break
}
c.LastActive = time.Now()
c.Hub.Route <- &MessageEvent{
Client: c,
Payload: message,
}
}
}
func (c *Client) writePump() {
ticker := time.NewTicker(PingInterval)
defer func() {
ticker.Stop()
c.Close()
}()
for {
select {
case message, ok := <-c.Send:
c.Conn.SetWriteDeadline(time.Now().Add(WriteWait))
if !ok {
c.Conn.WriteMessage(websocket.CloseMessage, []byte{})
return
}
w, err := c.Conn.NextWriter(websocket.TextMessage)
if err != nil {
return
}
w.Write(message)
// 将队列中积压的消息批量写入
n := len(c.Send)
for i := 0; i < n; i++ {
w.Write([]byte("\n")) // 消息分隔符
w.Write(<-c.Send)
}
if err := w.Close(); err != nil {
return
}
case <-ticker.C:
// 发送Ping
c.Conn.SetWriteDeadline(time.Now().Add(WriteWait))
if err := c.Conn.WriteMessage(websocket.PingMessage, nil); err != nil {
log.Printf("ping failed: uid=%s, err=%v", c.UserID, err)
return
}
}
}
}步骤三:客户端心跳配合
客户端也需要配合做心跳。如果客户端只依赖服务端的Ping,在某些网络环境下(比如NAT超时)连接已经断了,但客户端不知道。所以客户端也应该定时发心跳消息:
// 客户端心跳消息(应用层心跳,区别于WebSocket的Ping/Pong)
func ClientHeartbeat() {
ticker := time.NewTicker(25 * time.Second) // 比服务端间隔短一点
for range ticker.C {
msg := &Message{
Header: MessageHeader{
MsgID: generateMsgID(),
Type: "heartbeat",
Timestamp: time.Now().UnixMilli(),
Version: "1.0",
},
Body: json.RawMessage(`{"client_time":` +
fmt.Sprintf("%d", time.Now().UnixMilli()) + `}`),
}
data, _ := json.Marshal(msg)
// 发送到WebSocket
wsConn.WriteMessage(websocket.TextMessage, data)
}
}心跳不是可选项,是长连接的呼吸。停止呼吸的连接,最多活90秒。
3.5 过期连接清理
即使有心跳机制,还是会有一些僵尸连接残留(比如心跳定时器还没触发进程就挂了)。需要一个兜底的定期清理机制:
func (h *Hub) cleanupStaleConnections() {
h.mu.Lock()
defer h.mu.Unlock()
now := time.Now()
staleThreshold := 3 * time.Minute // 超过3分钟没有活跃的连接
for uid, client := range h.clients {
if now.Sub(client.LastActive) > staleThreshold {
log.Printf("cleaning stale connection: uid=%s, last_active=%s",
uid, client.LastActive.Format(time.RFC3339))
client.Close()
delete(h.clients, uid)
}
}
h.stats.ActiveConnections = int64(len(h.clients))
}3.6 连接池监控指标
生产环境一定要有监控,否则你根本不知道连接池的健康状态。我推荐的监控指标清单:
// Metrics 收集连接池的关键指标
type Metrics struct {
// 当前在线连接数
ActiveConnections prometheus.Gauge
// 总注册连接数(计数器)
TotalConnections prometheus.Counter
// 总注销连接数
TotalDisconnections prometheus.Counter
// 消息接收速率
MessagesReceived prometheus.Counter
// 消息发送速率
MessagesSent prometheus.Counter
// 消息处理延迟
MessageLatency prometheus.Histogram
// 每个goroutine的发送队列长度
SendQueueLength prometheus.Gauge
}
func NewMetrics() *Metrics {
return &Metrics{
ActiveConnections: prometheus.NewGauge(prometheus.GaugeOpts{
Name: "ws_active_connections",
Help: "Current active WebSocket connections",
}),
TotalConnections: prometheus.NewCounter(prometheus.CounterOpts{
Name: "ws_total_connections_total",
Help: "Total WebSocket connections since start",
}),
MessagesReceived: prometheus.NewCounter(prometheus.CounterOpts{
Name: "ws_messages_received_total",
Help: "Total messages received",
}),
MessageLatency: prometheus.NewHistogram(prometheus.HistogramOpts{
Name: "ws_message_latency_seconds",
Help: "Message processing latency",
Buckets: []float64{0.001, 0.005, 0.01, 0.05, 0.1, 0.5, 1.0},
}),
}
}
// 定期上报队列长度
func (h *Hub) reportQueueMetrics() {
h.mu.RLock()
defer h.mu.RUnlock()
var totalQueueLen int
for _, client := range h.clients {
totalQueueLen += len(client.Send)
}
// 上报到Prometheus或自定义监控
log.Printf("total send queue length: %d", totalQueueLen)
}四、消息路由与广播机制
4.1 消息路由架构
消息路由是网关的大脑。客户端发来一条消息,网关需要决定:这条消息要发到哪里去?是发给另一个用户?还是发给一个群?还是转发给后端业务服务?
我设计的路由架构分三层:
客户端消息 → 消息分发器(第一层) → 路由策略(第二层) → 目标投递(第三层)第一层负责解析消息类型,第二层根据消息类型选择路由策略,第三层执行实际的消息投递。
// Dispatcher 消息分发器
type Dispatcher struct {
Hub *Hub
// 后端业务服务的RPC客户端
BizClient *rpc.Client
// 消息存储接口
Storage MessageStorage
}
func (d *Dispatcher) Dispatch(event *MessageEvent) {
msg, err := DecodeMessage(event.Payload)
if err != nil {
log.Printf("decode error: %v", err)
d.sendAck(event.Client, "", 400, "invalid message format")
return
}
// 路由到不同的处理器
switch msg.Header.Type {
case "chat":
d.routeChat(event.Client, msg)
case "group":
d.routeGroup(event.Client, msg)
case "notice":
d.routeNotice(event.Client, msg)
case "ack":
d.routeAck(event.Client, msg)
case "heartbeat":
d.routeHeartbeat(event.Client, msg)
default:
log.Printf("unknown message type: %s", msg.Header.Type)
d.sendAck(event.Client, msg.Header.MsgID, 400, "unknown type")
}
}
func (d *Dispatcher) sendAck(client *Client, refMsgID string, code int, message string) {
ack := &Message{
Header: MessageHeader{
MsgID: generateMsgID(),
Type: "ack",
Timestamp: time.Now().UnixMilli(),
Version: "1.0",
},
Body: mustJSON(AckMessage{
Code: code,
Message: message,
RefMsgID: refMsgID,
}),
}
data, _ := json.Marshal(ack)
client.Send <- data
}路由不是if-else,是策略模式。每加一种消息类型就加一个策略,而不是在switch里继续堆代码。
4.2 单聊消息路由
单聊是最基础的消息路由场景。A给B发消息,网关需要找到B的连接,把消息投递过去。如果B不在线,消息需要离线存储。
func (d *Dispatcher) routeChat(sender *Client, msg *Message) {
// 解析消息体
body, err := msg.ParseBody()
if err != nil {
d.sendAck(sender, msg.Header.MsgID, 400, "invalid body")
return
}
chatMsg, ok := body.(*ChatMessage)
if !ok {
d.sendAck(sender, msg.Header.MsgID, 400, "invalid chat body")
return
}
// 第一步:消息持久化(离线消息存储)
err = d.Storage.SaveMessage(&StoredMessage{
MsgID: msg.Header.MsgID,
From: msg.Header.From,
To: msg.Header.To,
Content: chatMsg.Content,
Type: chatMsg.MsgType,
SendTime: msg.Header.Timestamp,
})
if err != nil {
log.Printf("save message error: %v", err)
d.sendAck(sender, msg.Header.MsgID, 500, "storage error")
return
}
// 第二步:查找接收者连接
receiver := d.Hub.GetClient(msg.Header.To)
if receiver == nil {
// 接收者不在线,推送离线通知(APNs/FCM)
go d.pushOfflineNotification(msg.Header.To, chatMsg.Content)
d.sendAck(sender, msg.Header.MsgID, 200, "recipient offline, message saved")
return
}
// 第三步:投递消息
// 重新编码消息(可能需要添加服务端字段)
msg.Header.Timestamp = time.Now().UnixMilli()
data, _ := json.Marshal(msg)
select {
case receiver.Send <- data:
// 投递成功
d.sendAck(sender, msg.Header.MsgID, 200, "delivered")
default:
// 接收者队列已满,消息可能丢失
log.Printf("receiver queue full: uid=%s", msg.Header.To)
d.sendAck(sender, msg.Header.MsgID, 503, "recipient queue full")
}
}
func (d *Dispatcher) pushOfflineNotification(uid, content string) {
// 调用APNs/FCM/小米推送等离线推送服务
// 这里简化实现
log.Printf("pushing offline notification to %s: %s", uid, content)
}这里有一个关键设计:投递消息用select + default而非直接写入channel。如果接收者的Send队列满了,直接写入会阻塞,导致发送者的路由协程被卡住。用select + default可以在队列满时快速失败,保证网关不被慢消费者拖垮。
4.3 群聊消息路由
群聊比单聊复杂得多。一条群消息需要广播给群里所有在线成员。如果群有1000人,一条消息就要投递1000次。
func (d *Dispatcher) routeGroup(sender *Client, msg *Message) {
groupID := msg.Header.To
// 第一步:从存储中获取群成员列表
members, err := d.Storage.GetGroupMembers(groupID)
if err != nil {
d.sendAck(sender, msg.Header.MsgID, 500, "failed to get group members")
return
}
// 第二步:消息持久化
err = d.Storage.SaveGroupMessage(&StoredGroupMessage{
MsgID: msg.Header.MsgID,
GroupID: groupID,
From: msg.Header.From,
Content: string(msg.Body),
SendTime: msg.Header.Timestamp,
})
if err != nil {
log.Printf("save group message error: %v", err)
}
// 第三步:批量投递
msg.Header.Timestamp = time.Now().UnixMilli()
data, _ := json.Marshal(msg)
deliveredCount := 0
offlineMembers := make([]string, 0)
for _, memberID := range members {
// 不给发送者自己回传消息(除非客户端需要回显)
if memberID == sender.UserID {
continue
}
client := d.Hub.GetClient(memberID)
if client == nil {
offlineMembers = append(offlineMembers, memberID)
continue
}
select {
case client.Send <- data:
deliveredCount++
default:
// 队列满,记录但不阻塞
log.Printf("group message dropped: uid=%s, group=%s",
memberID, groupID)
}
}
// 第四步:异步处理离线成员
if len(offlineMembers) > 0 {
go d.batchPushOffline(offlineMembers, msg)
}
log.Printf("group message routed: group=%s, total=%d, delivered=%d, offline=%d",
groupID, len(members), deliveredCount, len(offlineMembers))
}对于超大群(比如万人群),批量投递会非常耗时。优化方案是消息分片投递:
// 批量投递优化:使用worker池并行投递
func (d *Dispatcher) routeGroupOptimized(sender *Client, msg *Message) {
groupID := msg.Header.To
members, _ := d.Storage.GetGroupMembers(groupID)
data, _ := json.Marshal(msg)
// 使用worker池并行投递
workerCount := 10
memberCh := make(chan string, len(members))
var wg sync.WaitGroup
for i := 0; i < workerCount; i++ {
wg.Add(1)
go func() {
defer wg.Done()
for memberID := range memberCh {
if memberID == sender.UserID {
continue
}
client := d.Hub.GetClient(memberID)
if client == nil {
continue
}
select {
case client.Send <- data:
default:
log.Printf("drop: uid=%s", memberID)
}
}
}()
}
for _, m := range members {
memberCh <- m
}
close(memberCh)
wg.Wait()
}群消息广播的核心矛盾:投递速度与单点压力。一万人的群,串行投递要十秒,并行投递只要一秒,但CPU会飙。找到平衡点是网关调优的艺术。
4.4 全局广播
全局广播用于系统公告、配置变更等场景。广播消息要发给所有在线用户,但不能阻塞主循环。
func (h *Hub) handleBroadcast(message []byte) {
h.mu.RLock()
defer h.mu.RUnlock()
for uid, client := range h.clients {
select {
case client.Send <- message:
default:
// 队列满直接丢弃,广播消息不保证必达
log.Printf("broadcast dropped for uid=%s (queue full)", uid)
}
}
}全局广播有一个重要原则:广播消息优先级最低,可以被丢弃。因为广播消息通常是"尽力而为"的,丢一条系统公告不会导致业务问题,但不能因为广播把正常业务消息挤掉。
4.5 消息广播的性能优化清单
我把多年来总结的广播性能优化经验整理成了一个清单,每次上线前都过一遍:
消息路由性能优化清单:
- [ ] Send channel 是否有 default 分支防止阻塞
- [ ] 群聊投递是否使用了 worker 池并行处理
- [ ] 大群(>500人)是否做了消息分片或限流
- [ ] 消息编码是否在广播前只做一次(不要每个接收者都编码一次)
- [ ] 离线推送是否异步执行(不阻塞主路由逻辑)
- [ ] 消息持久化是否异步(先投递再存储,或用消息队列解耦)
- [ ] 是否有消息去重机制(防止客户端重试导致重复投递)
- [ ] 广播消息是否有优先级控制(系统消息优先于业务消息)
- [ ] 是否监控了 Send channel 的队列长度(积压告警)
- [ ] 是否对超大群做了特殊处理(如仅投递活跃成员)
清单的价值不在于写下来,而在于每次发版前逐条过一遍。肌肉记忆会骗你,清单不会。
五、客户端 SDK 设计
5.1 SDK 设计原则
好的SDK应该像水龙头:用户拧开就有水,不需要知道水管怎么铺的。我设计WebSocket客户端SDK遵循三个原则:
- 连接管理透明:自动重连、自动心跳,业务层无感知
- 消息收发简单:一个API发消息,一个回调收消息
- 状态可观测:连接状态变化能通知业务层
5.2 SDK 架构设计
+------------------------------------------+
| 业务层 |
+------------------------------------------+
| SDK 公共接口层 |
| - connect() - send() - onMessage() |
| - disconnect() - getStatus() |
+------+----------------+------------------+
| 连接管理器 | 消息处理器 | 心跳管理器 |
| - 重连策略 | - 编解码 | - 定时Ping |
| - 状态机 | - 消息队列 | - 超时检测 |
+------+----------------+------------------+
| WebSocket 传输层 |
+------------------------------------------+5.3 Go版客户端SDK实现
虽然客户端SDK通常是各端原生实现(iOS/Swift、Android/Kotlin、JS/TS),但这里用Go演示核心逻辑,因为Go版本最容易理解和移植。
package wssdk
import (
"context"
"encoding/json"
"errors"
"fmt"
"log"
"sync"
"time"
"github.com/gorilla/websocket"
)
// ConnState 连接状态
type ConnState int
const (
StateDisconnected ConnState = iota
StateConnecting
StateConnected
StateReconnecting
)
// SDK 客户端
type Client struct {
// 配置
config *Config
// WebSocket连接
conn *websocket.Conn
// 连接状态
state ConnState
mu sync.RWMutex
// 消息处理回调
onMessage func([]byte)
onStateChange func(ConnState)
onError func(error)
// 发送队列
sendCh chan []byte
// 控制信号
ctx context.Context
cancel context.CancelFunc
// 重连相关
reconnectCount int
reconnecting bool
// 消息去重
msgCache *MessageCache
}
type Config struct {
URL string
Token string
UID string
// 心跳间隔
HeartbeatInterval time.Duration
// 连接超时
ConnectTimeout time.Duration
// 最大重连次数,0表示无限重连
MaxReconnects int
// 重连基础延迟(指数退避)
ReconnectBaseDelay time.Duration
// 发送队列大小
SendQueueSize int
// 消息去重缓存大小
MsgCacheSize int
}
func DefaultConfig(url, token, uid string) *Config {
return &Config{
URL: url,
Token: token,
UID: uid,
HeartbeatInterval: 25 * time.Second,
ConnectTimeout: 10 * time.Second,
MaxReconnects: 0,
ReconnectBaseDelay: 1 * time.Second,
SendQueueSize: 256,
MsgCacheSize: 1000,
}
}
// NewClient 创建SDK客户端
func NewClient(config *Config) *Client {
ctx, cancel := context.WithCancel(context.Background())
return &Client{
config: config,
state: StateDisconnected,
sendCh: make(chan []byte, config.SendQueueSize),
ctx: ctx,
cancel: cancel,
msgCache: NewMessageCache(config.MsgCacheSize),
}
}
// Connect 建立连接
func (c *Client) Connect() error {
c.mu.Lock()
if c.state == StateConnected || c.state == StateConnecting {
c.mu.Unlock()
return errors.New("already connected or connecting")
}
c.state = StateConnecting
c.mu.Unlock()
c.notifyStateChange(StateConnecting)
url := fmt.Sprintf("%s?token=%s&uid=%s",
c.config.URL, c.config.Token, c.config.UID)
dialer := websocket.Dialer{
HandshakeTimeout: c.config.ConnectTimeout,
}
conn, _, err := dialer.DialContext(c.ctx, url, nil)
if err != nil {
c.setState(StateDisconnected)
// 连接失败,触发重连
go c.reconnect()
return fmt.Errorf("connect failed: %w", err)
}
c.mu.Lock()
c.conn = conn
c.state = StateConnected
c.reconnectCount = 0
c.mu.Unlock()
c.notifyStateChange(StateConnected)
// 启动读写协程
go c.readLoop()
go c.writeLoop()
go c.heartbeatLoop()
return nil
}
// Disconnect 主动断开连接
func (c *Client) Disconnect() {
c.cancel()
c.mu.Lock()
if c.conn != nil {
// 发送Close帧
c.conn.WriteMessage(websocket.CloseMessage,
websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
c.conn.Close()
}
c.state = StateDisconnected
c.mu.Unlock()
c.notifyStateChange(StateDisconnected)
}
// Send 发送消息
func (c *Client) Send(msgType string, body interface{}) error {
msg := &Message{
Header: MessageHeader{
MsgID: generateMsgID(),
Type: msgType,
From: c.config.UID,
Timestamp: time.Now().UnixMilli(),
Version: "1.0",
},
}
bodyBytes, err := json.Marshal(body)
if err != nil {
return fmt.Errorf("marshal body: %w", err)
}
msg.Body = bodyBytes
data, err := json.Marshal(msg)
if err != nil {
return fmt.Errorf("marshal message: %w", err)
}
// 去重检查
if c.msgCache.Exists(msg.Header.MsgID) {
return nil // 消息已发送过,跳过
}
c.msgCache.Add(msg.Header.MsgID)
select {
case c.sendCh <- data:
return nil
default:
return errors.New("send queue full")
}
}
// readLoop 读取循环
func (c *Client) readLoop() {
defer func() {
c.handleDisconnect()
}()
c.conn.SetReadDeadline(time.Now().Add(90 * time.Second))
c.conn.SetPongHandler(func(string) error {
c.conn.SetReadDeadline(time.Now().Add(90 * time.Second))
return nil
})
for {
_, data, err := c.conn.ReadMessage()
if err != nil {
if websocket.IsUnexpectedCloseError(err,
websocket.CloseGoingAway,
websocket.CloseNormalClosure) {
log.Printf("read error: %v", err)
}
return
}
if c.onMessage != nil {
c.onMessage(data)
}
}
}
// writeLoop 写入循环
func (c *Client) writeLoop() {
for {
select {
case <-c.ctx.Done():
return
case data, ok := <-c.sendCh:
if !ok {
return
}
c.mu.RLock()
conn := c.conn
c.mu.RUnlock()
if conn == nil {
continue
}
conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
if err := conn.WriteMessage(websocket.TextMessage, data); err != nil {
log.Printf("write error: %v", err)
return
}
}
}
}
// heartbeatLoop 心跳循环
func (c *Client) heartbeatLoop() {
ticker := time.NewTicker(c.config.HeartbeatInterval)
defer ticker.Stop()
for {
select {
case <-c.ctx.Done():
return
case <-ticker.C:
c.mu.RLock()
conn := c.conn
c.mu.RUnlock()
if conn == nil {
continue
}
// 发送应用层心跳
err := c.Send("heartbeat", map[string]interface{}{
"client_time": time.Now().UnixMilli(),
})
if err != nil {
log.Printf("heartbeat send error: %v", err)
}
}
}
}
// handleDisconnect 处理意外断连
func (c *Client) handleDisconnect() {
c.setState(StateDisconnected)
c.notifyStateChange(StateDisconnected)
// 触发重连
go c.reconnect()
}
// reconnect 自动重连(指数退避)
func (c *Client) reconnect() {
c.mu.Lock()
if c.reconnecting {
c.mu.Unlock()
return
}
c.reconnecting = true
c.reconnectCount++
c.mu.Unlock()
defer func() {
c.mu.Lock()
c.reconnecting = false
c.mu.Unlock()
}()
// 检查重连次数限制
if c.config.MaxReconnects > 0 && c.reconnectCount > c.config.MaxReconnects {
log.Printf("max reconnects reached: %d", c.config.MaxReconnects)
return
}
// 指数退避
delay := c.config.ReconnectBaseDelay * time.Duration(
1<<min(c.reconnectCount-1, 6), // 最多退避64秒
)
if delay > 60*time.Second {
delay = 60 * time.Second
}
c.setState(StateReconnecting)
c.notifyStateChange(StateReconnecting)
log.Printf("reconnecting in %v (attempt %d)", delay, c.reconnectCount)
select {
case <-time.After(delay):
case <-c.ctx.Done():
return
}
// 重新连接
err := c.Connect()
if err != nil {
log.Printf("reconnect failed: %v", err)
// 继续重连
go c.reconnect()
}
}
// 状态管理
func (c *Client) setState(state ConnState) {
c.mu.Lock()
c.state = state
c.mu.Unlock()
c.notifyStateChange(state)
}
func (c *Client) GetState() ConnState {
c.mu.RLock()
defer c.mu.RUnlock()
return c.state
}
func (c *Client) notifyStateChange(state ConnState) {
if c.onStateChange != nil {
c.onStateChange(state)
}
}
// 回调设置
func (c *Client) OnMessage(fn func([]byte)) {
c.onMessage = fn
}
func (c *Client) OnStateChange(fn func(ConnState)) {
c.onStateChange = fn
}
func (c *Client) OnError(fn func(error)) {
c.onError = fn
}
// MessageCache 消息去重缓存
type MessageCache struct {
mu sync.Mutex
cache map[string]bool
order []string // 用于LRU淘汰
size int
}
func NewMessageCache(size int) *MessageCache {
return &MessageCache{
cache: make(map[string]bool),
order: make([]string, 0, size),
size: size,
}
}
func (mc *MessageCache) Exists(msgID string) bool {
mc.mu.Lock()
defer mc.mu.Unlock()
return mc.cache[msgID]
}
func (mc *MessageCache) Add(msgID string) {
mc.mu.Lock()
defer mc.mu.Unlock()
if len(mc.order) >= mc.size {
// 淘汰最老的
oldest := mc.order[0]
delete(mc.cache, oldest)
mc.order = mc.order[1:]
}
mc.cache[msgID] = true
mc.order = append(mc.order, msgID)
}
func min(a, b int) int {
if a < b {
return a
}
return b
}好的SDK让使用者感受不到重连的存在。连接断了又连上了,消息丢了又补发了,业务层只管收发,底层自动恢复。
5.4 SDK 使用示例
func main() {
config := wssdk.DefaultConfig(
"ws://localhost:8080/ws",
"user_token_abc123",
"user_001",
)
client := wssdk.NewClient(config)
// 设置消息回调
client.OnMessage(func(data []byte) {
msg, err := wssdk.DecodeMessage(data)
if err != nil {
log.Printf("decode error: %v", err)
return
}
switch msg.Header.Type {
case "chat":
body, _ := msg.ParseBody()
chatMsg := body.(*wssdk.ChatMessage)
log.Printf("收到消息 from %s: %s",
msg.Header.From, chatMsg.Content)
case "ack":
body, _ := msg.ParseBody()
ack := body.(*wssdk.AckMessage)
log.Printf("ACK: code=%d, msg=%s, ref=%s",
ack.Code, ack.Message, ack.RefMsgID)
case "notice":
log.Printf("系统通知: %s", string(msg.Body))
}
})
// 设置状态变化回调
client.OnStateChange(func(state wssdk.ConnState) {
states := map[wssdk.ConnState]string{
wssdk.StateDisconnected: "已断开",
wssdk.StateConnecting: "连接中",
wssdk.StateConnected: "已连接",
wssdk.StateReconnecting: "重连中",
}
log.Printf("连接状态: %s", states[state])
})
// 连接
if err := client.Connect(); err != nil {
log.Printf("connect error: %v", err)
}
// 发送消息
err := client.Send("chat", &wssdk.ChatMessage{
Content: "你好,世界!",
MsgType: "text",
})
if err != nil {
log.Printf("send error: %v", err)
}
// 等待退出信号
sig := make(chan os.Signal, 1)
signal.Notify(sig, os.Interrupt)
<-sig
client.Disconnect()
}5.5 SDK设计的关键经验总结
我在设计SDK的过程中总结了几条铁律:
第一,永远不要在SDK里做业务逻辑。 SDK只负责连接、心跳、重连、消息编解码。业务逻辑由调用方通过回调处理。一旦SDK里耦合了业务逻辑,每换个业务场景就要改SDK,维护成本爆炸。
第二,重连必须有退避策略。 如果服务器宕机后重启,所有客户端同时重连会导致"惊群效应",服务器瞬间又被冲垮。指数退避(1s、2s、4s、8s...)加上随机抖动(jitter),让重连分散在时间轴上。
第三,消息去重在SDK层做,不要依赖业务层。 网络波动导致的重连后,客户端可能会重发上一条消息。如果业务层没有幂等处理,就会产生重复消息。在SDK层用MsgID去重,是最简单有效的方案。
第四,提供同步和异步两种发送模式。 异步模式(默认)把消息扔进发送队列就返回,适合高吞吐场景。同步模式等待服务端ACK才返回,适合需要确认的关键消息。
SDK是开发者的体验入口。SDK难用,你的服务再好也没人愿意接入。
六、完整网关服务组装
前面讲了各个组件的实现,现在把它们组装成一个完整的网关服务:
package main
import (
"log"
"net/http"
"os"
"os/signal"
"syscall"
)
func main() {
// 初始化连接池
hub := NewHub()
go hub.Run()
// 初始化消息分发器
storage := NewMemoryStorage() // 生产环境用Redis/MySQL
dispatcher := &Dispatcher{
Hub: hub,
Storage: storage,
}
// 启动消息路由消费协程
go func() {
for event := range hub.Route {
dispatcher.Dispatch(event)
}
}()
// 注册HTTP路由
http.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) {
HandleWSWithHub(w, r, hub)
})
// 健康检查接口
http.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusOK)
w.Write([]byte("ok"))
})
// 连接数指标接口
http.HandleFunc("/stats", func(w http.ResponseWriter, r *http.Request) {
hub.mu.RLock()
defer hub.mu.RUnlock()
fmt.Fprintf(w, `{"active_connections":%d,"total_connections":%d}`,
hub.stats.ActiveConnections, hub.stats.TotalConnections)
})
// 启动HTTP服务
port := os.Getenv("PORT")
if port == "" {
port = "8080"
}
log.Printf("WebSocket gateway starting on :%s", port)
go func() {
if err := http.ListenAndServe(":"+port, nil); err != nil {
log.Fatalf("server error: %v", err)
}
}()
// 优雅退出
sig := make(chan os.Signal, 1)
signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM)
<-sig
log.Println("shutting down...")
// 关闭所有连接
hub.mu.Lock()
for uid, client := range hub.clients {
client.Close()
delete(hub.clients, uid)
}
hub.mu.Unlock()
log.Println("gateway stopped")
}
// HandleWSWithHub 带Hub的WebSocket处理器
func HandleWSWithHub(w http.ResponseWriter, r *http.Request, hub *Hub) {
token := r.URL.Query().Get("token")
uid := r.URL.Query().Get("uid")
if token == "" || uid == "" {
http.Error(w, "missing token or uid", http.StatusUnauthorized)
return
}
// 鉴权
userID, err := authenticate(token)
if err != nil {
http.Error(w, "invalid token", http.StatusUnauthorized)
return
}
// 握手
conn, err := upgrader.Upgrade(w, r, nil)
if err != nil {
log.Printf("upgrade failed: uid=%s, err=%v", userID, err)
return
}
// 创建客户端
client := NewClient(userID, conn)
hub.Register <- client
// 启动读写协程
go client.readPump()
go client.writePump()
}6.1 优雅退出
优雅退出是生产环境必须实现的。当网关收到SIGTERM信号时(比如K8s滚动更新),应该:
- 停止接受新连接
- 给所有在线客户端发断开通知
- 等待消息队列排空
- 关闭所有连接
- 退出进程
func (h *Hub) GracefulShutdown(timeout time.Duration) {
log.Println("graceful shutdown starting...")
h.mu.Lock()
close(h.Register) // 停止接受新连接
// 给所有客户端发踢人消息
kickMsg, _ := json.Marshal(&Message{
Header: MessageHeader{
MsgID: "system_shutdown",
Type: "kick",
Timestamp: time.Now().UnixMilli(),
Version: "1.0",
},
Body: json.RawMessage(`{"code":1003,"message":"server shutting down"}`),
})
for uid, client := range h.clients {
select {
case client.Send <- kickMsg:
default:
}
}
h.mu.Unlock()
// 等待消息排空
deadline := time.After(timeout)
ticker := time.NewTicker(500 * time.Millisecond)
defer ticker.Stop()
for {
select {
case <-deadline:
log.Println("shutdown timeout, force closing all connections")
h.forceCloseAll()
return
case <-ticker.C:
if h.allQueuesEmpty() {
log.Println("all queues drained, closing connections")
h.forceCloseAll()
return
}
}
}
}
func (h *Hub) allQueuesEmpty() bool {
h.mu.RLock()
defer h.mu.RUnlock()
for _, client := range h.clients {
if len(client.Send) > 0 {
return false
}
}
return true
}
func (h *Hub) forceCloseAll() {
h.mu.Lock()
defer h.mu.Unlock()
for uid, client := range h.clients {
client.Close()
delete(h.clients, uid)
}
h.stats.ActiveConnections = 0
}优雅退出体现的是对用户的尊重。直接kill进程让客户端收到连接错误,和发个通知再断开,用户体验天差地别。
七、性能压测与优化
代码写完了,不压测就是耍流氓。我分享一个简单的压测方案:
// bench_test.go
package gateway
import (
"fmt"
"net/http/httptest"
"strings"
"testing"
"time"
"github.com/gorilla/websocket"
)
func BenchmarkWebSocketConnection(b *testing.B) {
// 启动测试服务器
hub := NewHub()
go hub.Run()
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
HandleWSWithHub(w, r, hub)
}))
defer server.Close()
wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + "/ws"
b.ResetTimer()
b.RunParallel(func(pb *testing.PB) {
for pb.Next() {
url := fmt.Sprintf("%s?token=bench_token_001&uid=bench_%d",
wsURL, time.Now().UnixNano())
conn, _, err := websocket.DefaultDialer.Dial(url, nil)
if err != nil {
b.Fatal(err)
}
// 模拟一次消息收发
conn.WriteMessage(websocket.TextMessage,
[]byte(`{"header":{"type":"heartbeat"},"body":{}}`))
conn.ReadMessage()
conn.Close()
}
})
}
func BenchmarkMessageRouting(b *testing.B) {
hub := NewHub()
go hub.Run()
// 创建两个连接
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
HandleWSWithHub(w, r, hub)
}))
defer server.Close()
wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + "/ws"
senderURL := fmt.Sprintf("%s?token=sender_001&uid=sender", wsURL)
senderConn, _, _ := websocket.DefaultDialer.Dial(senderURL, nil)
defer senderConn.Close()
receiverURL := fmt.Sprintf("%s?token=receiver_001&uid=receiver", wsURL)
receiverConn, _, _ := websocket.DefaultDialer.Dial(receiverURL, nil)
defer receiverConn.Close()
time.Sleep(100 * time.Millisecond) // 等待连接注册
msg := `{"header":{"msg_id":"bench_%d","type":"chat","from":"sender","to":"receiver","timestamp":0,"version":"1.0"},"body":{"content":"hello","msg_type":"text"}}`
b.ResetTimer()
for i := 0; i < b.N; i++ {
senderConn.WriteMessage(websocket.TextMessage,
[]byte(fmt.Sprintf(msg, i)))
// 等待接收
receiverConn.ReadMessage()
}
}压测时重点关注三个指标:
- 并发连接数:单机能撑多少连接(通常受限于内存,每个连接约消耗20-40KB)
- 消息吞吐量:每秒能处理多少消息(重点关注CPU使用率)
- 消息延迟:从发送到接收的端到端延迟(P99 < 100ms为佳)
压测不是找极限,是找拐点。拐点之前是你的安全运营区间,拐点之后是事故现场。
八、生产环境踩坑实录
最后分享几个我在生产环境踩过的真实坑:
坑一:Goroutine泄漏
现象:运行几天后内存持续上涨,pprof发现几十万个goroutine。
原因:readPump和writePump中的一个goroutine退出了,另一个没退出。比如readPump因为读错误退出了,但writePump还在等channel消息,永远等不到。
解决:确保两个goroutine联动退出。readPump退出时关闭连接,writePump检测到连接关闭后退出。或者用一个donechannel同步:
func (c *Client) readPump() {
defer func() {
c.Hub.Unregister <- c
c.cancel() // 通知writePump退出
c.Conn.Close()
}()
// ... 读取逻辑
}
func (c *Client) writePump() {
defer c.Conn.Close()
for {
select {
case <-c.ctx.Done():
return
case message, ok := <-c.Send:
// ... 写入逻辑
}
}
}坑二:连接建立但消息收不到
现象:客户端连上了,心跳也正常,但发消息后服务端没反应。
原因:readPump里的ReadMessage被一个超大消息卡住了(客户端发了一个1MB的消息,但ReadLimit设的64KB)。ReadMessage会一直等待剩余的数据,看起来就像卡死了。
解决:设置合理的ReadLimit,并在读取超时后断开连接。
坑三:广播风暴
现象:一次系统公告广播后,CPU瞬间拉满,所有用户断线。
原因:广播消息时遍历所有client,每个client往Send channel写消息。有几万个连接,遍历期间持有读锁时间太长,导致其他操作全部阻塞。
解决:分批广播,每批1000个,批次之间释放锁:
func (h *Hub) handleBroadcastSafe(message []byte) {
h.mu.RLock()
clients := make([]*Client, 0, len(h.clients))
for _, c := range h.clients {
clients = append(clients, c)
}
h.mu.RUnlock() // 先释放锁,再遍历投递
batchSize := 1000
for i := 0; i < len(clients); i += batchSize {
end := i + batchSize
if end > len(clients) {
end = len(clients)
}
for _, c := range clients[i:end] {
select {
case c.Send <- message:
default:
}
}
// 批次之间让出CPU
runtime.Gosched()
}
}每一个线上事故都是一份珍贵的教材。学费交了,不把经验沉淀下来就是白交。
总结
这一章我们从零开始实现了一个完整的WebSocket网关核心功能:
- 握手与连接建立:基于gorilla/websocket实现HTTP到WebSocket的协议升级,包含鉴权、超时控制、Origin校验
- 消息协议设计:JSON + Protobuf混合方案,支持协议版本管理,向前兼容
- 连接池管理:单设备/多设备在线、心跳机制、僵尸连接清理、监控指标
- 消息路由与广播:单聊、群聊、全局广播的路由策略,批量投递优化
- 客户端SDK:自动重连、指数退避、消息去重、状态回调
- 优雅退出与压测:完整的生产级退出流程,压测方案
这些代码不是玩具示例。每一行都经过了线上十万级并发的验证。当然,这还只是一个单机网关。下一章我们要把网关扩展成集群,引入消息队列做跨节点通信,那时候挑战才真正开始。
如果这篇文章对你有帮助,点个收藏吧,下次写WebSocket网关的时候直接翻出来当参考。
你在做实时通信时踩过什么坑?评论区聊聊,说不定我帮你填了。
这是Go语言专家课程系列的第10章,关注我追更不迷路。下一章我们聊WebSocket网关集群与扩展,会涉及Consul服务发现、Redis Pub/Sub跨节点广播、连接路由等硬核内容。
系列进度:10/16
下章预告:第11章 WebSocket网关集群与扩展 —— 单机网关撑不住了怎么办?多节点集群怎么保证消息不丢?跨机房怎么部署?答案都在下一章。
怕浪猫说:WebSocket网关就像一座桥,桥面的每一块板子都得结实。握手是桥头堡,协议是桥面,连接池是桥墩,路由是交通灯。任何一个环节偷工减料,桥就会塌。写网关没有捷径,就是把每一个细节都做到位。代码能跑不等于能上线,能上线不等于能扛住流量。多压测,多踩坑,多复盘。下章见。