feat(push-gateway): 修复WebSocket并发写竞争并添加DEV_MODE鉴权
hub.go 重写用send chan+单写协程模式避免并发写竞争 handler.go 加DEV_MODE dev-token支持+broadcast端点 config.go 加DevMode/RedisURL字段
This commit is contained in:
@@ -6,54 +6,94 @@ import (
|
||||
"github.com/gorilla/websocket"
|
||||
)
|
||||
|
||||
// Hub 管理 WebSocket 客户端连接
|
||||
// Connection 包装单个 WebSocket 连接及其异步发送通道
|
||||
type Connection struct {
|
||||
UserID string
|
||||
conn *websocket.Conn
|
||||
send chan []byte
|
||||
}
|
||||
|
||||
// Send 将消息投递到该连接的发送通道(非阻塞,通道满则丢弃)
|
||||
func (c *Connection) Send(message []byte) {
|
||||
select {
|
||||
case c.send <- message:
|
||||
default:
|
||||
// 通道满,丢弃消息避免阻塞 hub
|
||||
}
|
||||
}
|
||||
|
||||
// Outgoing 返回该连接的发送通道,供写协程消费
|
||||
func (c *Connection) Outgoing() <-chan []byte {
|
||||
return c.send
|
||||
}
|
||||
|
||||
// Hub 管理所有在线 WebSocket 连接,按 userID 索引
|
||||
type Hub struct {
|
||||
mu sync.RWMutex
|
||||
clients map[string]map[*websocket.Conn]bool // userID -> connections
|
||||
clients map[string]map[*Connection]bool
|
||||
}
|
||||
|
||||
// NewHub 创建 Hub 实例
|
||||
func NewHub() *Hub {
|
||||
return &Hub{
|
||||
clients: make(map[string]map[*websocket.Conn]bool),
|
||||
clients: make(map[string]map[*Connection]bool),
|
||||
}
|
||||
}
|
||||
|
||||
func (h *Hub) Register(userID string, conn *websocket.Conn) {
|
||||
// Register 注册一个用户连接,返回 Connection 供调用方持有
|
||||
func (h *Hub) Register(userID string, conn *websocket.Conn) *Connection {
|
||||
c := &Connection{
|
||||
UserID: userID,
|
||||
conn: conn,
|
||||
send: make(chan []byte, 64),
|
||||
}
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
|
||||
if h.clients[userID] == nil {
|
||||
h.clients[userID] = make(map[*websocket.Conn]bool)
|
||||
h.clients[userID] = make(map[*Connection]bool)
|
||||
}
|
||||
h.clients[userID][conn] = true
|
||||
h.clients[userID][c] = true
|
||||
return c
|
||||
}
|
||||
|
||||
func (h *Hub) Unregister(userID string, conn *websocket.Conn) {
|
||||
// Unregister 注销一个用户连接并关闭其发送通道
|
||||
func (h *Hub) Unregister(c *Connection) {
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
|
||||
if conns, ok := h.clients[userID]; ok {
|
||||
delete(conns, conn)
|
||||
if len(conns) == 0 {
|
||||
delete(h.clients, userID)
|
||||
}
|
||||
conns, ok := h.clients[c.UserID]
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
if _, exists := conns[c]; exists {
|
||||
delete(conns, c)
|
||||
close(c.send)
|
||||
}
|
||||
if len(conns) == 0 {
|
||||
delete(h.clients, c.UserID)
|
||||
}
|
||||
}
|
||||
|
||||
// SendToUser 向指定用户的所有连接推送消息
|
||||
// SendToUser 向指定用户的所有在线连接推送消息
|
||||
func (h *Hub) SendToUser(userID string, message []byte) error {
|
||||
h.mu.RLock()
|
||||
defer h.mu.RUnlock()
|
||||
|
||||
conns, ok := h.clients[userID]
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
|
||||
for conn := range conns {
|
||||
if err := conn.WriteMessage(websocket.TextMessage, message); err != nil {
|
||||
return err
|
||||
}
|
||||
for c := range conns {
|
||||
c.Send(message)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Broadcast 向所有在线连接广播消息
|
||||
func (h *Hub) Broadcast(message []byte) {
|
||||
h.mu.RLock()
|
||||
defer h.mu.RUnlock()
|
||||
for _, conns := range h.clients {
|
||||
for c := range conns {
|
||||
c.Send(message)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user