Golang WebSocket开发 实时通信服务搭建

github.com/gorilla/websocket

package main

import (
    "log"
    "net/http"
    "time"

    "github.com/gorilla/websocket"
)

// Hub 是我们服务的核心，负责管理所有活跃的客户端，并向它们广播消息。
type Hub struct {
    clients    map[*Client]bool // 注册的客户端
    broadcast  chan []byte      // 从客户端接收的入站消息，待广播
    register   chan *Client     // 客户端注册请求
    unregister chan *Client     // 客户端取消注册请求
}

func newHub() *Hub {
    return &Hub{
        broadcast:  make(chan []byte),
        register:   make(chan *Client),
        unregister: make(chan *Client),
        clients:    make(map[*Client]bool),
    }
}

// run 方法在独立的goroutine中运行，处理Hub的内部事件循环。
func (h *Hub) run() {
    for {
        select {
        case client := <-h.register:
            h.clients[client] = true
            log.Printf("新客户端连接: %s", client.conn.RemoteAddr())
        case client := <-h.unregister:
            if _, ok := h.clients[client]; ok {
                delete(h.clients, client)
                close(client.send) // 关闭发送通道，通知客户端的writePump停止
                log.Printf("客户端断开: %s", client.conn.RemoteAddr())
            }
        case message := <-h.broadcast:
            for client := range h.clients {
                select {
                case client.send <- message: // 尝试将消息发送给客户端
                default: // 如果客户端的发送通道阻塞，说明客户端可能处理不过来或者已断开
                    close(client.send)
                    delete(h.clients, client)
                    log.Printf("客户端发送缓冲区已满或已断开，取消注册: %s", client.conn.RemoteAddr())
                }
            }
        }
    }
}

// Client 代表一个独立的WebSocket连接。
type Client struct {
    hub  *Hub
    conn *websocket.Conn // 实际的WebSocket连接
    send chan []byte     // 用于发送消息的缓冲通道
}

const (
    writeWait      = 10 * time.Second // 写入消息到对端的时间限制
    pongWait       = 60 * time.Second // 等待对端pong消息的时间
    pingPeriod     = (pongWait * 9) / 10 // 发送ping消息的周期，应小于pongWait
    maxMessageSize = 512 // 允许从对端接收的最大消息大小
)

var upgrader = websocket.Upgrader{
    ReadBufferSize:  1024,
    WriteBufferSize: 1024,
    CheckOrigin: func(r *http.Request) bool {
        // 生产环境中，这里需要严格检查Origin头，避免CORS问题和恶意连接
        // 比如：return r.Header.Get("Origin") == "http://your-allowed-domain.com"
        return true // 示例中为方便起见，允许所有来源
    },
}

// readPump 从WebSocket连接读取消息，并将其发送到Hub。
func (c *Client) readPump() {
    defer func() {
        c.hub.unregister <- c // 客户端断开时，通知Hub取消注册
        c.conn.Close()
    }()
    c.conn.SetReadLimit(maxMessageSize)
    c.conn.SetReadDeadline(time.Now().Add(pongWait)) // 设置读取超时，用于心跳检测
    c.conn.SetPongHandler(func(string) error { // 收到pong消息时，重置读取超时
        c.conn.SetReadDeadline(time.Now().Add(pongWait))
        return nil
    })
    for {
        _, message, err := c.conn.ReadMessage()
        if err != nil {
            // 判断是否是正常的关闭错误，否则打印日志
            if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
                log.Printf("读取消息错误: %v", err)
            }
            break // 发生错误，退出循环，清理连接
        }
        log.Printf("收到消息 (%s): %s", c.conn.RemoteAddr(), message)
        // 实际应用中，这里会对消息进行解析和处理，然后可能广播或定向发送
        c.hub.broadcast <- message // 示例中直接广播收到的消息
    }
}

// writePump 从Hub的发送通道读取消息，并将其写入WebSocket连接。
func (c *Client) writePump() {
    ticker := time.NewTicker(pingPeriod) // 定时发送ping消息
    defer func() {
        ticker.Stop()
        c.conn.Close()
    }()
    for {
        select {
        case message, ok := <-c.send:
            c.conn.SetWriteDeadline(time.Now().Add(writeWait)) // 设置写入超时
            if !ok {
                // Hub关闭了通道，表示Hub要关闭此连接
                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(<-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 {
                return // 发送失败，退出循环
            }
        }
    }
}

// serveWs 处理WebSocket连接升级请求。
func serveWs(hub *Hub, w http.ResponseWriter, r *http.Request) {
    conn, err := upgrader.Upgrade(w, r, nil)
    if err != nil {
        log.Println("WebSocket升级失败:", err)
        return
    }
    // 创建新客户端，并将其注册到Hub
    client := &Client{hub: hub, conn: conn, send: make(chan []byte, 256)} // 256的缓冲通道，提供一定的背压能力
    client.hub.register <- client

    // 在独立的Goroutine中运行读写泵，确保不阻塞主线程
    go client.writePump()
    go client.readPump()
}

func main() {
    hub := newHub()
    go hub.run() // 启动Hub的事件循环

    http.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) {
        serveWs(hub, w, r)
    })

    log.Println("WebSocket服务器正在 :8080 端口启动...")
    err := http.ListenAndServe(":8080", nil)
    if err != nil {
        log.Fatal("服务器启动失败: ", err)
    }
}

Hub

Client

readPump

writePump

send

readPump

writePump

defer

context.Context

Client.send

select

default

pingPeriod

pongWait