Skip to content

第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协议。

来看握手请求的关键头部:

http
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
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库。

先看最基础的握手实现:

go
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 连接建立后的初始化

握手成功只是第一步,连接建立后还有一系列初始化工作要做。我把这个过程叫做"连接生命周期初始化":

go
// 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
            }
        }
    }
}

这里有一个关键设计:读写分离双协程readPumpwritePump各跑一个goroutine,互不干扰。为什么这么设计?因为WebSocket是全双工的,读和写可以同时进行。如果用单协程处理读写,一个慢消息会阻塞整个连接的处理。

Send通道的缓冲大小设为256,这是一个经验值。如果你的业务场景消息突发量大(比如群聊消息广播),可以调大。但不要设太大,否则消息积压会吃光内存。更好的做法是配合背压机制,当队列满到80%时丢弃低优先级消息。

通道不是越大越好,缓冲大小是你对消息积压的容忍度。积压超过容忍度,就该丢而不是存。


二、消息协议设计(JSON/Protobuf)

2.1 协议设计的重要性

协议是WebSocket网关的灵魂。很多团队在协议设计上栽的跟头比在代码逻辑上栽的还多。我见过一个团队,消息格式今天用JSON,明天加个字段,后天改个字段名,半年后客户端版本碎片化严重,新老消息格式混在一起,整个系统乱成一锅粥。

协议设计的核心原则:向前兼容、高效编解码、可扩展

2.2 JSON协议设计

JSON是最简单的协议格式,调试方便,可读性好。对于内部系统或中小规模应用,JSON完全够用。

我推荐的JSON消息格式:

go
// 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
}

为什么Bodyjson.RawMessage而不是interface{}?因为interface{}会立即反序列化整个JSON,而RawMessage是延迟解析的。你可以先解析Header,根据Type字段决定要不要解析Body,以及用什么结构体解析。这样能显著减少不必要的反序列化开销。

协议设计的第一原则是向前兼容。今天加的字段,明天不能删;今天选的格式,后天不能换。除非你想写三版客户端兼容代码。

来看一个完整的JSON消息编解码示例:

go
// 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的优势在于:

  1. 编解码速度比JSON快3-5倍
  2. 二进制体积比JSON小30%-50%
  3. 强类型,编译时就能发现字段错误
  4. 天然支持向前兼容(新字段加optional即可)

先定义proto文件:

protobuf
// 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:

go
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。

go
// 根据消息类型选择编码方式
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 协议版本管理

协议一定会演进,版本管理必须从一开始就设计好。我踩过最深的坑就是没有版本号,上线后老客户端收到新格式消息直接崩溃。

go
// 版本兼容性检查
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网关的核心组件。一个设计良好的连接池要解决以下问题:

  1. 快速查找:通过UserID快速找到连接
  2. 并发安全:多个goroutine同时读写连接池
  3. 连接迁移:用户重连时,旧连接要优雅关闭
  4. 状态监控:实时知道当前在线连接数
go
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 注册与注销逻辑

go
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

go
type Client struct {
    // ... 其他字段
    closeOnce sync.Once
}

func (c *Client) Close() {
    c.closeOnce.Do(func() {
        close(c.Send)
        c.Conn.Close()
    })
}

3.3 多设备同时在线

很多IM系统都需要支持一个用户多设备同时在线(比如手机和电脑都登录)。这时候连接池的结构要改:

go
// 多设备连接管理
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秒(两个心跳周期+冗余)
go
// 心跳参数
const (
    PingInterval     = 30 * time.Second
    PongWait         = 60 * time.Second
    ReadTimeout      = 90 * time.Second
    WriteWait        = 10 * time.Second
    MaxMissedPongs   = 2 // 允许丢失的最大Pong次数
)

步骤二:服务端心跳实现

go
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超时)连接已经断了,但客户端不知道。所以客户端也应该定时发心跳消息:

go
// 客户端心跳消息(应用层心跳,区别于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 过期连接清理

即使有心跳机制,还是会有一些僵尸连接残留(比如心跳定时器还没触发进程就挂了)。需要一个兜底的定期清理机制:

go
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 连接池监控指标

生产环境一定要有监控,否则你根本不知道连接池的健康状态。我推荐的监控指标清单:

go
// 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 消息路由架构

消息路由是网关的大脑。客户端发来一条消息,网关需要决定:这条消息要发到哪里去?是发给另一个用户?还是发给一个群?还是转发给后端业务服务?

我设计的路由架构分三层:

客户端消息 → 消息分发器(第一层) → 路由策略(第二层) → 目标投递(第三层)

第一层负责解析消息类型,第二层根据消息类型选择路由策略,第三层执行实际的消息投递。

go
// 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不在线,消息需要离线存储。

go
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次。

go
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))
}

对于超大群(比如万人群),批量投递会非常耗时。优化方案是消息分片投递

go
// 批量投递优化:使用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 全局广播

全局广播用于系统公告、配置变更等场景。广播消息要发给所有在线用户,但不能阻塞主循环。

go
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遵循三个原则:

  1. 连接管理透明:自动重连、自动心跳,业务层无感知
  2. 消息收发简单:一个API发消息,一个回调收消息
  3. 状态可观测:连接状态变化能通知业务层

5.2 SDK 架构设计

+------------------------------------------+
|              业务层                       |
+------------------------------------------+
|         SDK 公共接口层                    |
|  - connect()  - send()  - onMessage()    |
|  - disconnect()  - getStatus()           |
+------+----------------+------------------+
| 连接管理器 | 消息处理器   | 心跳管理器     |
| - 重连策略 | - 编解码     | - 定时Ping     |
| - 状态机   | - 消息队列   | - 超时检测     |
+------+----------------+------------------+
|           WebSocket 传输层               |
+------------------------------------------+

5.3 Go版客户端SDK实现

虽然客户端SDK通常是各端原生实现(iOS/Swift、Android/Kotlin、JS/TS),但这里用Go演示核心逻辑,因为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 使用示例

go
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难用,你的服务再好也没人愿意接入。


六、完整网关服务组装

前面讲了各个组件的实现,现在把它们组装成一个完整的网关服务:

go
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滚动更新),应该:

  1. 停止接受新连接
  2. 给所有在线客户端发断开通知
  3. 等待消息队列排空
  4. 关闭所有连接
  5. 退出进程
go
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进程让客户端收到连接错误,和发个通知再断开,用户体验天差地别。


七、性能压测与优化

代码写完了,不压测就是耍流氓。我分享一个简单的压测方案:

go
// 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()
    }
}

压测时重点关注三个指标:

  1. 并发连接数:单机能撑多少连接(通常受限于内存,每个连接约消耗20-40KB)
  2. 消息吞吐量:每秒能处理多少消息(重点关注CPU使用率)
  3. 消息延迟:从发送到接收的端到端延迟(P99 < 100ms为佳)

压测不是找极限,是找拐点。拐点之前是你的安全运营区间,拐点之后是事故现场。


八、生产环境踩坑实录

最后分享几个我在生产环境踩过的真实坑:

坑一:Goroutine泄漏

现象:运行几天后内存持续上涨,pprof发现几十万个goroutine。

原因:readPumpwritePump中的一个goroutine退出了,另一个没退出。比如readPump因为读错误退出了,但writePump还在等channel消息,永远等不到。

解决:确保两个goroutine联动退出。readPump退出时关闭连接,writePump检测到连接关闭后退出。或者用一个donechannel同步:

go
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个,批次之间释放锁:

go
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网关核心功能:

  1. 握手与连接建立:基于gorilla/websocket实现HTTP到WebSocket的协议升级,包含鉴权、超时控制、Origin校验
  2. 消息协议设计:JSON + Protobuf混合方案,支持协议版本管理,向前兼容
  3. 连接池管理:单设备/多设备在线、心跳机制、僵尸连接清理、监控指标
  4. 消息路由与广播:单聊、群聊、全局广播的路由策略,批量投递优化
  5. 客户端SDK:自动重连、指数退避、消息去重、状态回调
  6. 优雅退出与压测:完整的生产级退出流程,压测方案

这些代码不是玩具示例。每一行都经过了线上十万级并发的验证。当然,这还只是一个单机网关。下一章我们要把网关扩展成集群,引入消息队列做跨节点通信,那时候挑战才真正开始。


如果这篇文章对你有帮助,点个收藏吧,下次写WebSocket网关的时候直接翻出来当参考。

你在做实时通信时踩过什么坑?评论区聊聊,说不定我帮你填了。

这是Go语言专家课程系列的第10章,关注我追更不迷路。下一章我们聊WebSocket网关集群与扩展,会涉及Consul服务发现、Redis Pub/Sub跨节点广播、连接路由等硬核内容。


系列进度:10/16

下章预告:第11章 WebSocket网关集群与扩展 —— 单机网关撑不住了怎么办?多节点集群怎么保证消息不丢?跨机房怎么部署?答案都在下一章。


怕浪猫说:WebSocket网关就像一座桥,桥面的每一块板子都得结实。握手是桥头堡,协议是桥面,连接池是桥墩,路由是交通灯。任何一个环节偷工减料,桥就会塌。写网关没有捷径,就是把每一个细节都做到位。代码能跑不等于能上线,能上线不等于能扛住流量。多压测,多踩坑,多复盘。下章见。

热爱生活,喜好美食,追求未来!