fix(ws): add per-conn mutex to prevent concurrent WriteJSON panic

Root cause: Broadcast() and Push() launched multiple goroutines
writing to the same websocket.Conn simultaneously. Gorilla websocket
does not support concurrent WriteMessage/WriteJSON.

Fix: add connMu sync.Map (*websocket.Conn → *sync.Mutex) to WSManager.
Each WriteJSON now goes through writeWithLock() which acquires the
per-conn mutex. Unregister() also acquires the lock before Close()
to prevent write/close races.

No functional changes to API surface.
This commit is contained in:
xiaowei 2026-05-30 17:15:19 +08:00
parent c5835def8f
commit cd49a2801a
1 changed files with 25 additions and 5 deletions

View File

@ -15,8 +15,9 @@ var upgrader = websocket.Upgrader{
// WSManager 管理所有 WebSocket 连接,支持按 AgentID 定向推送
type WSManager struct {
mu sync.RWMutex
clients map[string]map[*websocket.Conn]bool // agentID → 连接集合
mu sync.RWMutex
clients map[string]map[*websocket.Conn]bool // agentID → 连接集合
connMu sync.Map // per-conn mutex: *websocket.Conn → *sync.Mutex
}
// 全局 WebSocket 管理器
@ -32,6 +33,7 @@ func (wm *WSManager) Register(agentID string, conn *websocket.Conn) {
wm.clients[agentID] = make(map[*websocket.Conn]bool)
}
wm.clients[agentID][conn] = true
wm.connMu.Store(conn, &sync.Mutex{})
log.Printf("[ws] agent %s connected (%d total)", agentID, len(wm.clients[agentID]))
}
@ -45,10 +47,28 @@ func (wm *WSManager) Unregister(agentID string, conn *websocket.Conn) {
delete(wm.clients, agentID)
}
}
conn.Close()
// 先获取该 conn 的锁,确保没有正在进行的 write然后关闭
if muVal, ok := wm.connMu.Load(conn); ok {
mu := muVal.(*sync.Mutex)
mu.Lock()
conn.Close()
mu.Unlock()
} else {
conn.Close()
}
wm.connMu.Delete(conn)
log.Printf("[ws] agent %s disconnected", agentID)
}
// writeWithLock 安全地写入 WebSocket JSON每个连接保证单 writer
func (wm *WSManager) writeWithLock(conn *websocket.Conn, msg interface{}) error {
muVal, _ := wm.connMu.Load(conn)
mu := muVal.(*sync.Mutex)
mu.Lock()
defer mu.Unlock()
return conn.WriteJSON(msg)
}
// Push 推送到指定 Agent 所有连接
func (wm *WSManager) Push(agentID string, msgType string, payload interface{}) {
wm.mu.RLock()
@ -66,7 +86,7 @@ func (wm *WSManager) Push(agentID string, msgType string, payload interface{}) {
for conn := range clients {
go func(c *websocket.Conn) {
if err := c.WriteJSON(msg); err != nil {
if err := wm.writeWithLock(c, msg); err != nil {
log.Printf("[ws] write error: %v", err)
wm.Unregister(agentID, c)
}
@ -87,7 +107,7 @@ func (wm *WSManager) Broadcast(msgType string, payload interface{}) {
for agentID, clients := range wm.clients {
for conn := range clients {
go func(c *websocket.Conn, aid string) {
if err := c.WriteJSON(msg); err != nil {
if err := wm.writeWithLock(c, msg); err != nil {
log.Printf("[ws] broadcast error: %v", err)
wm.Unregister(aid, c)
}