如何使用Gorilla WebSocket向指定用户推送通知

核心问题分析 #

当前代码使用全局单一通道Chann，导致所有连接的WebSocket客户端都会收到通知信号，无法实现仅推送给目标用户的需求。需要维护用户ID与对应WebSocket连接的映射关系，实现精准推送。

改造步骤 #

1. 维护用户-连接映射 #

定义线程安全的映射，存储用户ID到其WebSocket连接的关联：

import (
    "sync"
    "github.com/gorilla/websocket"
)

// 全局维护用户ID到WebSocket连接的映射，加锁保证线程安全
var (
    clientConns = make(map[string]*websocket.Conn)
    connMutex   sync.RWMutex
)

2. 修改连接处理逻辑 #

用户建立WebSocket连接时将连接存入映射，连接关闭时从映射中移除：

func (h *notificationHandler) Connect(w http.ResponseWriter, r *http.Request) {
    id := r.URL.Query().Get("id")
    if id == "" {
        http.Error(w, "missing user id", http.StatusBadRequest)
        return
    }

    c, err := upgrader.Upgrade(w, r, nil)
    if err != nil {
        log.Print("upgrade:", err)
        return
    }
    defer c.Close()

    // 将连接加入映射
    connMutex.Lock()
    clientConns[id] = c
    connMutex.Unlock()

    // 连接关闭时移除映射
    defer func() {
        connMutex.Lock()
        delete(clientConns, id)
        connMutex.Unlock()
    }()

    // 处理客户端消息（按需保留，若无需接收客户端消息可简化）
    for {
        _, _, err := c.ReadMessage()
        if err != nil {
            log.Println("read error:", err)
            break
        }
    }
}

3. 修改通知推送逻辑 #

创建新通知时，根据user_id找到对应WebSocket连接，直接推送信号：

func (h *notificationHandler) CreateNotification(w http.ResponseWriter, r *http.Request) error {
    // 解析请求参数
    var req struct {
        UserID string `json:"user_id"`
        Text   string `json:"text"`
    }
    if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
        return err
    }

    // 插入通知到数据库
    _, err := h.db.Exec("INSERT INTO notifications (user_id, text) VALUES (?, ?)", req.UserID, req.Text)
    if err != nil {
        return err
    }

    // 定向推送通知信号
    connMutex.RLock()
    c, exists := clientConns[req.UserID]
    connMutex.RUnlock()

    if exists {
        go func() {
            signal := map[string]string{"event": "signal"}
            signalBytes, err := json.Marshal(signal)
            if err != nil {
                log.Println("marshal error:", err)
                return
            }
            if err := c.WriteMessage(websocket.TextMessage, signalBytes); err != nil {
                log.Println("write error:", err)
                // 推送失败时清理无效连接
                connMutex.Lock()
                delete(clientConns, req.UserID)
                connMutex.Unlock()
            }
        }()
    }

    // 返回成功响应
    w.WriteHeader(http.StatusCreated)
    return json.NewEncoder(w).Encode(map[string]string{"status": "ok"})
}

4. 清理冗余代码 #

原代码中的全局Chann和signal函数已不再需要，可直接删除。

关键注意事项 #

• 线程安全：对clientConns的读写必须加锁，避免并发读写引发panic。
• 无效连接清理：推送失败时及时移除无效连接，避免后续资源浪费。
• 客户端重连：需在客户端实现连接断开后的自动重连逻辑，保证通知可达。
• 多连接场景：若同一用户存在多设备连接，可根据需求选择覆盖旧连接或向所有连接推送。

内容的提问来源于stack exchange，提问作者Almaz