dfb6d2bfc1
hub.go 重写用send chan+单写协程模式避免并发写竞争 handler.go 加DEV_MODE dev-token支持+broadcast端点 config.go 加DevMode/RedisURL字段
100 lines
2.0 KiB
Go
100 lines
2.0 KiB
Go
package hub
|
|
|
|
import (
|
|
"sync"
|
|
|
|
"github.com/gorilla/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[*Connection]bool
|
|
}
|
|
|
|
// NewHub 创建 Hub 实例
|
|
func NewHub() *Hub {
|
|
return &Hub{
|
|
clients: make(map[string]map[*Connection]bool),
|
|
}
|
|
}
|
|
|
|
// 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[*Connection]bool)
|
|
}
|
|
h.clients[userID][c] = true
|
|
return c
|
|
}
|
|
|
|
// Unregister 注销一个用户连接并关闭其发送通道
|
|
func (h *Hub) Unregister(c *Connection) {
|
|
h.mu.Lock()
|
|
defer h.mu.Unlock()
|
|
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 向指定用户的所有在线连接推送消息
|
|
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 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)
|
|
}
|
|
}
|
|
}
|