package ws import ( "encoding/json" "sync" "sync/atomic" "time" "go.uber.org/zap" ) const ( defaultUserBufferSize = 128 maxReplayEvents = 200 ) // Event 表示一个WebSocket事件 type Event struct { ID uint64 `json:"event_id"` Type string `json:"type"` Event string `json:"event"` // 事件名称(兼容多种协议) TS int64 `json:"ts"` Payload interface{} `json:"payload"` } // Client 表示一个WebSocket客户端连接 type Client struct { ID uint64 UserID string Send chan []byte Quit chan struct{} } // Message 表示WebSocket消息 type Message struct { Type string `json:"type"` Payload json.RawMessage `json:"payload"` } // ResponseMessage 表示发送给客户端的消息 type ResponseMessage struct { EventID uint64 `json:"event_id"` Type string `json:"type"` TS int64 `json:"ts"` Payload interface{} `json:"payload,omitempty"` } // ErrorMessage 表示错误消息 type ErrorMessage struct { Type string `json:"type"` Payload struct { Code string `json:"code"` Message string `json:"message"` } `json:"payload"` } // subscriber 表示一个事件订阅者 type subscriber struct { id uint64 ch chan Event quit chan struct{} } // Hub WebSocket连接管理器 type Hub struct { seq uint64 mu sync.RWMutex clients map[string]map[uint64]*Client // userID -> clientID -> Client history map[string][]Event // userID -> []Event (用于断线重连历史回放) subscribers map[string]map[uint64]*subscriber // userID -> subscriberID -> subscriber } // NewHub 创建WebSocket Hub func NewHub() *Hub { return &Hub{ clients: make(map[string]map[uint64]*Client), history: make(map[string][]Event), subscribers: make(map[string]map[uint64]*subscriber), } } // NextID 获取下一个事件ID func (h *Hub) NextID() uint64 { return atomic.AddUint64(&h.seq, 1) } // Register 注册客户端连接 func (h *Hub) Register(client *Client) []Event { h.mu.Lock() defer h.mu.Unlock() if _, ok := h.clients[client.UserID]; !ok { h.clients[client.UserID] = make(map[uint64]*Client) } h.clients[client.UserID][client.ID] = client // 获取历史事件用于回放 replay := make([]Event, 0) for _, e := range h.history[client.UserID] { replay = append(replay, e) } zap.L().Debug("WebSocket client registered", zap.String("user_id", client.UserID), zap.Uint64("client_id", client.ID), zap.Int("total_clients", len(h.clients[client.UserID])), ) return replay } // Unregister 注销客户端连接 func (h *Hub) Unregister(client *Client) { h.mu.Lock() defer h.mu.Unlock() if userClients, ok := h.clients[client.UserID]; ok { if _, exists := userClients[client.ID]; exists { delete(userClients, client.ID) close(client.Quit) if len(userClients) == 0 { delete(h.clients, client.UserID) } zap.L().Debug("WebSocket client unregistered", zap.String("user_id", client.UserID), zap.Uint64("client_id", client.ID), ) } } } // HasClients 检查用户是否有活跃连接 func (h *Hub) HasClients(userID string) bool { h.mu.RLock() defer h.mu.RUnlock() return len(h.clients[userID]) > 0 } // GetClientCount 获取用户的连接数 func (h *Hub) GetClientCount(userID string) int { h.mu.RLock() defer h.mu.RUnlock() return len(h.clients[userID]) } // PublishToUser 向指定用户发送事件 func (h *Hub) PublishToUser(userID string, eventType string, payload interface{}) Event { ev := Event{ ID: h.NextID(), Type: eventType, TS: time.Now().UnixMilli(), Payload: payload, } h.publish(userID, ev) return ev } // PublishToUsers 向多个用户发送事件 func (h *Hub) PublishToUsers(userIDs []string, eventType string, payload interface{}) { for _, uid := range userIDs { h.PublishToUser(uid, eventType, payload) } } // publish 内部发布方法 func (h *Hub) publish(userID string, ev Event) { h.mu.Lock() // 存储历史事件(最多保留200条) history := append(h.history[userID], ev) if len(history) > maxReplayEvents { history = history[len(history)-maxReplayEvents:] } h.history[userID] = history // 复制客户端列表避免锁竞争 targets := make([]*Client, 0, len(h.clients[userID])) for _, c := range h.clients[userID] { targets = append(targets, c) } h.mu.Unlock() // 构建消息 msg := ResponseMessage{ EventID: ev.ID, Type: ev.Type, TS: ev.TS, Payload: ev.Payload, } data, err := json.Marshal(msg) if err != nil { zap.L().Error("Failed to marshal WebSocket message", zap.String("user_id", userID), zap.String("event_type", ev.Type), zap.Error(err), ) return } // 非阻塞发送,慢消费者丢弃消息 for _, c := range targets { select { case <-c.Quit: // 客户端已断开 case c.Send <- data: // 消息已发送 default: // 发送缓冲区满,丢弃消息 zap.L().Warn("WebSocket client buffer full, dropping message", zap.String("user_id", userID), zap.Uint64("client_id", c.ID), ) } } } // SendError 向客户端发送错误消息 func (h *Hub) SendError(client *Client, code string, message string) { errMsg := ErrorMessage{ Type: "error", } errMsg.Payload.Code = code errMsg.Payload.Message = message data, err := json.Marshal(errMsg) if err != nil { return } select { case <-client.Quit: case client.Send <- data: default: } } // Broadcast 广播消息给所有连接 func (h *Hub) Broadcast(eventType string, payload interface{}) { h.mu.RLock() userIDs := make([]string, 0, len(h.clients)) for uid := range h.clients { userIDs = append(userIDs, uid) } h.mu.RUnlock() h.PublishToUsers(userIDs, eventType, payload) } // GetOnlineUsers 获取在线用户列表 func (h *Hub) GetOnlineUsers() []string { h.mu.RLock() defer h.mu.RUnlock() users := make([]string, 0, len(h.clients)) for uid := range h.clients { users = append(users, uid) } return users } // GetOnlineCount 获取在线用户数 func (h *Hub) GetOnlineCount() int { h.mu.RLock() defer h.mu.RUnlock() return len(h.clients) } // Subscribe 订阅事件 func (h *Hub) Subscribe(userID string, afterID uint64) (chan Event, func(), []Event) { subID := h.NextID() sub := &subscriber{ id: subID, ch: make(chan Event, defaultUserBufferSize), quit: make(chan struct{}), } h.mu.Lock() if _, ok := h.subscribers[userID]; !ok { h.subscribers[userID] = make(map[uint64]*subscriber) } h.subscribers[userID][subID] = sub replay := make([]Event, 0) for _, e := range h.history[userID] { if e.ID > afterID { replay = append(replay, e) } } h.mu.Unlock() cancel := func() { h.mu.Lock() defer h.mu.Unlock() if userSubs, ok := h.subscribers[userID]; ok { if s, exists := userSubs[subID]; exists { close(s.quit) delete(userSubs, subID) close(s.ch) } if len(userSubs) == 0 { delete(h.subscribers, userID) } } } return sub.ch, cancel, replay } // EncodeData 编码事件数据 func EncodeData(ev Event) (string, error) { body, err := json.Marshal(ev.Payload) if err != nil { return "", err } return string(body), nil }