Go Web 编程快速入门 11 - WebSocket实时通信：实时消息推送和双向通信

WebSocket 提供了浏览器与服务端的全双工通信能力，适合实时消息、在线协作、游戏等场景。本章延续第 04.1 章的风格，从最小可运行示例到工程化封装与房间管理，帮助你快速构建稳定的实时系统。

1 WebSocket基础与握手

1.1 升级器最小封装

go 复制代码

package ws

import (
    "net/http"
    "time"
    "github.com/gorilla/websocket"
)

type Config struct {
    ReadTimeout    time.Duration
    WriteTimeout   time.Duration
    PingInterval   time.Duration
    MaxMessageSize int64
    CheckOrigin    bool
}

func DefaultConfig() Config {
    return Config{
        ReadTimeout:    60 * time.Second,
        WriteTimeout:   10 * time.Second,
        PingInterval:   30 * time.Second,
        MaxMessageSize: 1 << 20, // 1MB
        CheckOrigin:    false,
    }
}

type Upgrader struct {
    u      websocket.Upgrader
    config Config
}

func NewUpgrader(cfg Config) *Upgrader {
    if cfg.ReadTimeout == 0 { cfg = DefaultConfig() }
    return &Upgrader{
        config: cfg,
        u: websocket.Upgrader{
            ReadBufferSize:  4096,
            WriteBufferSize: 4096,
            CheckOrigin: func(r *http.Request) bool { return !cfg.CheckOrigin || r.Header.Get("Origin") != "" },
        },
    }
}

func (u *Upgrader) Upgrade(w http.ResponseWriter, r *http.Request) (*websocket.Conn, error) {
    conn, err := u.u.Upgrade(w, r, nil)
    if err != nil { return nil, err }
    conn.SetReadLimit(u.config.MaxMessageSize)
    conn.SetReadDeadline(time.Now().Add(u.config.ReadTimeout))
    conn.SetPongHandler(func(string) error {
        conn.SetReadDeadline(time.Now().Add(u.config.ReadTimeout))
        return nil
    })
    return conn, nil
}

1.2 连接读写泵（最小实现）

go 复制代码

type Conn struct {
    c       *websocket.Conn
    cfg     Config
    send    chan []byte
    closing chan struct{}
}

func NewConn(c *websocket.Conn, cfg Config) *Conn {
    conn := &Conn{c: c, cfg: cfg, send: make(chan []byte, 256), closing: make(chan struct{})}
    go conn.readPump(); go conn.writePump()
    return conn
}

func (c *Conn) readPump() {
    defer c.Close()
    for {
        select {
        case <-c.closing:
            return
        default:
            mt, msg, err := c.c.ReadMessage()
            if err != nil { return }
            if mt == websocket.TextMessage {
                // 这里可以接入路由或房间广播
            }
        }
    }
}

func (c *Conn) writePump() {
    ticker := time.NewTicker(c.cfg.PingInterval)
    defer ticker.Stop()
    for {
        select {
        case msg := <-c.send:
            c.c.SetWriteDeadline(time.Now().Add(c.cfg.WriteTimeout))
            if err := c.c.WriteMessage(websocket.TextMessage, msg); err != nil { return }
        case <-ticker.C:
            c.c.SetWriteDeadline(time.Now().Add(c.cfg.WriteTimeout))
            if err := c.c.WriteMessage(websocket.PingMessage, nil); err != nil { return }
        case <-c.closing:
            return
        }
    }
}

func (c *Conn) Close() { close(c.closing); _ = c.c.Close() }
func (c *Conn) Send(data []byte) { select { case c.send <- data: default: } }

2 连接管理与房间

2.1 连接池与房间模型

go 复制代码

type Hub struct {
    conns  map[string]*Conn           // sessionID -> Conn
    rooms  map[string]map[string]bool // roomID -> sessionID set
}

func NewHub() *Hub {
    return &Hub{conns: make(map[string]*Conn), rooms: make(map[string]map[string]bool)}
}

func (h *Hub) Add(sessionID string, c *Conn) { h.conns[sessionID] = c }
func (h *Hub) Remove(sessionID string)        { delete(h.conns, sessionID) }

func (h *Hub) Join(roomID, sessionID string) {
    if h.rooms[roomID] == nil { h.rooms[roomID] = make(map[string]bool) }
    h.rooms[roomID][sessionID] = true
}

func (h *Hub) Leave(roomID, sessionID string) { if h.rooms[roomID] != nil { delete(h.rooms[roomID], sessionID) } }

func (h *Hub) Broadcast(roomID string, data []byte) {
    for sid := range h.rooms[roomID] {
        if c := h.conns[sid]; c != nil { c.Send(data) }
    }
}

2.2 路由与消息分发

你可以在 readPump 中解析消息类型（文本/JSON），根据房间或目标用户进行分发。为保证健壮性，建议：

控制消息大小与速率（令牌桶）。
统一格式化错误与系统消息。
对外暴露"加入/离开房间、私聊、群发"三类基础命令。

3 心跳与断线重连

服务端定期发送 Ping，客户端回复 Pong。
超过 PongTimeout 未响应则关闭连接并清理资源。
客户端应在断线后自动重连，并回放必要的状态（如重新加入房间）。

4 完整示例：聊天室服务

go 复制代码

package main

import (
    "encoding/json"
    "fmt"
    "log"
    "net/http"
    "time"
    "github.com/gorilla/websocket"
)

var upgrader = ws.NewUpgrader(ws.DefaultConfig())
var hub = NewHub()

func main() {
    http.HandleFunc("/ws", handleWS)
    http.HandleFunc("/rooms/broadcast", httpBroadcast)
    log.Println("WebSocket 服务: http://localhost:8080/ws")
    log.Fatal(http.ListenAndServe(":8080", nil))
}

func handleWS(w http.ResponseWriter, r *http.Request) {
    conn, err := upgrader.Upgrade(w, r)
    if err != nil { http.Error(w, err.Error(), http.StatusBadRequest); return }
    c := NewConn(conn, ws.DefaultConfig())
    sessionID := fmt.Sprintf("s_%d", time.Now().UnixNano())
    hub.Add(sessionID, c)
    hub.Join("lobby", sessionID) // 默认加入大厅
    c.Send([]byte("欢迎加入 lobby"))
}

func httpBroadcast(w http.ResponseWriter, r *http.Request) {
    type req struct{ Room string; Content string }
    var body req
    if err := json.NewDecoder(r.Body).Decode(&body); err != nil { http.Error(w, "请求体错误", 400); return }
    hub.Broadcast(body.Room, []byte(body.Content))
    _ = json.NewEncoder(w).Encode(map[string]string{"status": "ok"})
}

通过上述最小实现，你已具备搭建聊天室的核心能力。后续可以加入鉴权、消息持久化、在线状态、服务端水平扩展（如使用 Redis/PubSub 做跨实例广播）等增强特性。