Go-WebSocket实时通信从聊天室到消息推送系统实战

原创于 2026-08-20 22:30:16 发布 · 174 阅读

2 ·

本内容遵循CC 4.0 BY-SA版权协议

GEO检测

收录于

Go-WebSocket实时通信从聊天室到消息推送系统实战

文章导语

WebSocket是实现实时双向通信的利器。在Go中，goroutine-per-connection模型与WebSocket天然契合。本文基于gorilla/websocket库，实现从基础聊天室到生产级消息推送系统。

一、WebSocket基础服务

import "github.com/gorilla/websocket"

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

func handleWebSocket(w http.ResponseWriter, r *http.Request) {
    conn, err := upgrader.Upgrade(w, r, nil)
    if err != nil {
        log.Print("升级失败:", err)
        return
    }
    defer conn.Close()
    
    for {
        messageType, msg, err := conn.ReadMessage()
        if err != nil {
            break
        }
        // Echo回复
        conn.WriteMessage(messageType, msg)
    }
}

二、聊天室Hub模式

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

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

三、心跳保活

func (c *Client) ReadPump(hub *Hub) {
    defer func() {
        hub.unregister <- c
        c.conn.Close()
    }()
    
    c.conn.SetReadDeadline(time.Now().Add(pongWait))
    c.conn.SetPongHandler(func(string) error {
        c.conn.SetReadDeadline(time.Now().Add(pongWait))
        return nil
    })
    
    for {
        _, msg, err := c.conn.ReadMessage()
        if err != nil {
            break
        }
        hub.broadcast <- msg
    }
}

func (c *Client) WritePump() {
    ticker := time.NewTicker(pingPeriod)
    defer func() {
        ticker.Stop()
        c.conn.Close()
    }()
    
    for {
        select {
        case msg, ok := <-c.send:
            if !ok { return }
            c.conn.WriteMessage(websocket.TextMessage, msg)
        case <-ticker.C:
            c.conn.WriteMessage(websocket.PingMessage, nil)
        }
    }
}