package ws import ( "context" "encoding/json" "sync" "time" "github.com/google/uuid" "github.com/redis/go-redis/v9" "go.uber.org/zap" ) // busMsg Redis Pub/Sub 消息格式 type busMsg struct { InstanceID string `json:"instance_id"` MsgType string `json:"msg_type"` // "user_msg" | "broadcast" TargetUsers []string `json:"target_users"` EventType string `json:"event_type"` EventID uint64 `json:"event_id"` TS int64 `json:"ts"` Payload json.RawMessage `json:"payload"` } // WSClusterConfig WebSocket 集群配置 type WSClusterConfig struct { InstanceID string MsgChannel string OnlineTTL int HeartbeatInterval int } // Bus 组合 Hub + Redis Pub/Sub + OnlineTracker,实现跨实例消息路由 type Bus struct { hub *Hub rdb *redis.Client tracker *OnlineTracker instanceID string channel string sub *redis.PubSub cancel context.CancelFunc mu sync.Mutex closed bool } // NewBus 创建 Bus 实例 // 如果 rdb 为 nil(standalone 模式或 Redis 不可用),Bus 退化为纯本地 Hub 代理 func NewBus(hub *Hub, rdb *redis.Client, cfg *WSClusterConfig) *Bus { instanceID := cfg.InstanceID if instanceID == "" { instanceID = uuid.New().String() } onlineTTL := time.Duration(cfg.OnlineTTL) * time.Second heartbeatInterval := time.Duration(cfg.HeartbeatInterval) * time.Second if heartbeatInterval <= 0 { heartbeatInterval = 20 * time.Second } if onlineTTL <= 0 { onlineTTL = 60 * time.Second } tracker := NewOnlineTracker(rdb, instanceID, onlineTTL, heartbeatInterval) return &Bus{ hub: hub, rdb: rdb, tracker: tracker, instanceID: instanceID, channel: cfg.MsgChannel, } } // Hub 获取底层 Hub(供 WebSocket handler 使用 Register/Unregister 等本地操作) func (b *Bus) Hub() *Hub { return b.hub } // InstanceID 获取实例 ID func (b *Bus) InstanceID() string { return b.instanceID } // Start 启动 Bus 的 Redis 订阅和心跳 func (b *Bus) Start() error { if b.rdb == nil { zap.L().Info("WebSocket Bus: no Redis client, running in standalone mode") return nil } // 启动心跳 b.tracker.StartHeartbeat() // 订阅 Redis channel ctx, cancel := context.WithCancel(context.Background()) b.cancel = cancel sub := b.rdb.Subscribe(ctx, b.channel) b.sub = sub // 确保订阅已建立 if _, err := sub.Receive(ctx); err != nil { cancel() return err } go b.subscribeLoop(ctx, sub) zap.L().Info("WebSocket Bus started", zap.String("instance_id", b.instanceID), zap.String("channel", b.channel), ) return nil } // Register 注册客户端连接(代理 Hub.Register + tracker) func (b *Bus) Register(client *Client) error { if err := b.hub.Register(client); err != nil { return err } b.tracker.UserConnected(client.UserID) return nil } // Unregister 注销客户端连接(代理 Hub.Unregister + tracker) func (b *Bus) Unregister(client *Client) { b.hub.Unregister(client) b.tracker.UserDisconnected(client.UserID) } // Close 关闭 Bus:关闭订阅,清理在线状态 func (b *Bus) Close() error { b.mu.Lock() defer b.mu.Unlock() if b.closed { return nil } b.closed = true if b.cancel != nil { b.cancel() } if b.sub != nil { b.sub.Close() } b.tracker.Stop() b.hub.CloseAllConnections() zap.L().Info("WebSocket Bus closed", zap.String("instance_id", b.instanceID)) return nil } // --- MessagePublisher 接口实现 --- // PublishToUser 向指定用户发送事件(本地 + 远程) func (b *Bus) PublishToUser(userID string, eventType string, payload any) Event { ev := b.hub.PublishToUser(userID, eventType, payload) b.publishToRemote([]string{userID}, "user_msg", ev, payload) return ev } // PublishToUsers 向多个用户发送事件 func (b *Bus) PublishToUsers(userIDs []string, eventType string, payload any) { b.hub.PublishToUsers(userIDs, eventType, payload) b.publishToRemote(userIDs, "user_msg", Event{ Type: eventType, TS: time.Now().UnixMilli(), }, payload) } // PublishToUsersConcurrent 并发推送消息给多个用户 func (b *Bus) PublishToUsersConcurrent(userIDs []string, eventType string, payload any) { b.hub.PublishToUsersConcurrent(userIDs, eventType, payload) b.publishToRemote(userIDs, "user_msg", Event{ Type: eventType, TS: time.Now().UnixMilli(), }, payload) } // PublishToUserOnline 仅在用户在线时发送事件 func (b *Bus) PublishToUserOnline(userID string, eventType string, payload any) bool { delivered := b.hub.PublishToUserOnline(userID, eventType, payload) if !delivered { // 本地不在线,尝试远程投递 ev := Event{ Type: eventType, TS: time.Now().UnixMilli(), } b.publishToRemote([]string{userID}, "user_msg", ev, payload) } return delivered } // PublishToUserOnlineReliable 可靠地发送事件给在线用户 func (b *Bus) PublishToUserOnlineReliable(userID string, eventType string, payload any) bool { delivered := b.hub.PublishToUserOnlineReliable(userID, eventType, payload) if !delivered { ev := Event{ Type: eventType, TS: time.Now().UnixMilli(), } b.publishToRemote([]string{userID}, "user_msg", ev, payload) } return delivered } // Broadcast 广播消息给所有连接 func (b *Bus) Broadcast(eventType string, payload any) { b.hub.Broadcast(eventType, payload) b.publishToRemote(nil, "broadcast", Event{ Type: eventType, TS: time.Now().UnixMilli(), }, payload) } // HasClients 检查用户是否在线(本地优先,然后查 Redis) func (b *Bus) HasClients(userID string) bool { if b.hub.HasClients(userID) { return true } return b.tracker.IsOnline(userID) } // FilterOnline 批量检查用户在线状态 func (b *Bus) FilterOnline(userIDs []string) (online, offline []string) { localOnline, offline := b.hub.FilterOnline(userIDs) if b.rdb == nil { return localOnline, offline } // 对 offline 的列表,查 Redis remoteOnline, stillOffline := b.tracker.FilterOnline(offline) online = append(localOnline, remoteOnline...) return online, stillOffline } // SendError 向客户端发送错误消息(本地操作) func (b *Bus) SendError(client *Client, code string, message string) { b.hub.SendError(client, code, message) } // NextID 获取下一个事件 ID func (b *Bus) NextID() uint64 { return b.hub.NextID() } // OnDisconnect 注册断开连接事件处理器 func (b *Bus) OnDisconnect(handler DisconnectHandler) { b.hub.OnDisconnect(handler) } // OnConnect 注册连接事件处理器 func (b *Bus) OnConnect(handler ConnectHandler) { b.hub.OnConnect(handler) } // --- 内部方法 --- // publishToRemote 发布消息到 Redis channel func (b *Bus) publishToRemote(targetUsers []string, msgType string, ev Event, payload any) { if b.rdb == nil { return } payloadBytes, err := json.Marshal(payload) if err != nil { zap.L().Error("Bus: failed to marshal payload for remote publish", zap.String("event_type", ev.Type), zap.Error(err), ) return } msg := busMsg{ InstanceID: b.instanceID, MsgType: msgType, TargetUsers: targetUsers, EventType: ev.Type, EventID: ev.ID, TS: ev.TS, Payload: payloadBytes, } data, err := json.Marshal(msg) if err != nil { zap.L().Error("Bus: failed to marshal busMsg", zap.Error(err)) return } if err := b.rdb.Publish(context.Background(), b.channel, data).Err(); err != nil { zap.L().Error("Bus: failed to publish to Redis", zap.String("channel", b.channel), zap.Error(err), ) } } // subscribeLoop 处理 Redis 订阅消息 func (b *Bus) subscribeLoop(ctx context.Context, sub *redis.PubSub) { ch := sub.Channel() for { select { case <-ctx.Done(): return case msg, ok := <-ch: if !ok { return } b.handleRemoteMessage(msg) } } } // handleRemoteMessage 处理从 Redis 收到的远程消息 func (b *Bus) handleRemoteMessage(msg *redis.Message) { var bm busMsg if err := json.Unmarshal([]byte(msg.Payload), &bm); err != nil { zap.L().Error("Bus: failed to unmarshal remote message", zap.Error(err)) return } // 过滤自身发送的消息 if bm.InstanceID == b.instanceID { return } switch bm.MsgType { case "user_msg": b.handleRemoteUserMsg(bm) case "broadcast": b.handleRemoteBroadcast(bm) default: zap.L().Warn("Bus: unknown msg_type from remote", zap.String("msg_type", bm.MsgType), ) } } // handleRemoteUserMsg 处理远程定向消息 func (b *Bus) handleRemoteUserMsg(bm busMsg) { // 构造 ResponseMessage 用于本地投递 respMsg := ResponseMessage{ EventID: bm.EventID, Type: bm.EventType, TS: bm.TS, } // 反序列化 payload var payload any if err := json.Unmarshal(bm.Payload, &payload); err != nil { zap.L().Error("Bus: failed to unmarshal remote payload", zap.Error(err)) return } respMsg.Payload = payload data, err := json.Marshal(respMsg) if err != nil { zap.L().Error("Bus: failed to marshal response message", zap.Error(err)) return } b.hub.deliverToUsers(bm.TargetUsers, data) } // handleRemoteBroadcast 处理远程广播消息 func (b *Bus) handleRemoteBroadcast(bm busMsg) { b.hub.Broadcast(bm.EventType, bm.Payload) } // deliverToUsers 向指定用户列表投递预序列化消息 // 这是 Hub 批量投递的内部方法,Bus 需要调用它来投递远程消息 func (h *Hub) deliverToUsers(userIDs []string, data []byte) { for _, uid := range userIDs { h.mu.RLock() targets := make([]*Client, 0, len(h.clients[uid])) for _, c := range h.clients[uid] { targets = append(targets, c) } h.mu.RUnlock() for _, c := range targets { select { case <-c.Quit: case c.Send <- data: default: zap.L().Warn("WebSocket client buffer full, closing slow client", zap.String("user_id", uid), zap.Uint64("client_id", c.ID), ) close(c.Quit) } } } }