跳到主要内容

WebSocket

问题​

Go 中如何实现 WebSocket？如何管理连接和广播消息？

答案​

使用 gorilla/websocket​

import "github.com/gorilla/websocket"

var upgrader = websocket.Upgrader{
    ReadBufferSize:  1024,
    WriteBufferSize: 1024,
    CheckOrigin: func(r *http.Request) bool {
        return true // 生产环境应检查 Origin
    },
}

func wsHandler(w http.ResponseWriter, r *http.Request) {
    conn, err := upgrader.Upgrade(w, r, nil)
    if err != nil {
        log.Println("upgrade error:", err)
        return
    }
    defer conn.Close()
    
    for {
        msgType, msg, err := conn.ReadMessage()
        if err != nil {
            break // 连接关闭
        }
        // Echo
        conn.WriteMessage(msgType, msg)
    }
}

连接管理（Hub 模式）​

type Hub struct {
    clients    map[*Client]bool
    broadcast  chan []byte
    register   chan *Client
    unregister chan *Client
    mu         sync.RWMutex
}

type Client struct {
    hub  *Hub
    conn *websocket.Conn
    send chan []byte
}

func (h *Hub) run() {
    for {
        select {
        case client := <-h.register:
            h.mu.Lock()
            h.clients[client] = true
            h.mu.Unlock()
        case client := <-h.unregister:
            h.mu.Lock()
            if _, ok := h.clients[client]; ok {
                delete(h.clients, client)
                close(client.send)
            }
            h.mu.Unlock()
        case message := <-h.broadcast:
            h.mu.RLock()
            for client := range h.clients {
                select {
                case client.send <- message:
                default:
                    // 发送缓冲满了，关闭慢客户端
                    close(client.send)
                    delete(h.clients, client)
                }
            }
            h.mu.RUnlock()
        }
    }
}

心跳保活​

const (
    pongWait   = 60 * time.Second
    pingPeriod = (pongWait * 9) / 10
)

// 读 goroutine
func (c *Client) readPump() {
    defer c.hub.unregister <- c
    c.conn.SetReadDeadline(time.Now().Add(pongWait))
    c.conn.SetPongHandler(func(string) error {
        c.conn.SetReadDeadline(time.Now().Add(pongWait))
        return nil
    })
    // 读循环...
}

// 写 goroutine
func (c *Client) writePump() {
    ticker := time.NewTicker(pingPeriod)
    defer ticker.Stop()
    
    for {
        select {
        case msg := <-c.send:
            c.conn.WriteMessage(websocket.TextMessage, msg)
        case <-ticker.C:
            c.conn.WriteMessage(websocket.PingMessage, nil)
        }
    }
}

常见面试问题​

Q1: 为什么每个 WebSocket 连接需要两个 goroutine？​

答案：WebSocket 的读和写可以并发执行，但 gorilla/websocket 的单个 Conn 不支持并发写。用两个 goroutine 分别负责读（readPump）和写（writePump），通过 Channel 通信，既保证了并发安全，又保证了消息及时发送。

Q2: 如何处理慢客户端？​

答案：使用带缓冲的 send Channel。如果 Channel 满了（客户端消费太慢），直接关闭连接，避免影响其他客户端的消息广播。

相关链接​