feat(ws): implement WebSocket cluster mode with Redis Pub/Sub

Introduce a new WebSocket messaging architecture that supports both standalone and cluster modes. This allows for horizontal scaling of WebSocket servers by using Redis Pub/Sub to synchronize messages across multiple instances. Key changes: - Added `ws.MessagePublisher` interface to abstract message distribution. - Implemented `ws.Bus` to handle cluster-mode messaging via Redis. - Added `ws.OnlineTracker` to manage user online status across the cluster. - Refactored multiple services (Chat, Group, Push, Call, etc.) to use the new `MessagePublisher` instead of a concrete `ws.Hub`. - Added WebSocket configuration options (mode, instance ID, channel, TTL, heartbeat) to `config.yaml` and `config.go`. - Updated dependency injection with Wire to support the new publisher and Redis client. - Improved logging by replacing standard `log` with `zap` in several service components.
2026-05-06 12:39:11 +08:00
parent d481742790
commit c630cbf4d0
22 changed files with 1082 additions and 120 deletions
--- a/internal/pkg/ws/bus.go
+++ b/internal/pkg/ws/bus.go
@@ -0,0 +1,395 @@
+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)
+			}
+		}
+	}
+}
--- a/internal/pkg/ws/online_tracker.go
+++ b/internal/pkg/ws/online_tracker.go
@@ -0,0 +1,257 @@
+package ws
+
+import (
+	"context"
+	"fmt"
+	"sync"
+	"time"
+
+	"github.com/redis/go-redis/v9"
+	"go.uber.org/zap"
+)
+
+// OnlineTracker 分布式在线状态追踪
+// 使用 Redis Hash 存储各实例的用户连接数，支持跨实例在线查询
+type OnlineTracker struct {
+	rdb              *redis.Client
+	instanceID       string
+	localUsers       sync.Map // 本实例在线用户集合: userID -> connCount
+	onlineTTL        time.Duration
+	heartbeatInterval time.Duration
+	cancel           context.CancelFunc
+}
+
+// NewOnlineTracker 创建在线状态追踪器
+func NewOnlineTracker(rdb *redis.Client, instanceID string, onlineTTL, heartbeatInterval time.Duration) *OnlineTracker {
+	return &OnlineTracker{
+		rdb:               rdb,
+		instanceID:        instanceID,
+		onlineTTL:         onlineTTL,
+		heartbeatInterval: heartbeatInterval,
+	}
+}
+
+// UserConnected 用户连接，更新 Redis 和本地在线状态
+func (t *OnlineTracker) UserConnected(userID string) {
+	if t.rdb == nil {
+		return
+	}
+	ctx := context.Background()
+	key := fmt.Sprintf("online:%s", userID)
+
+	// 本地计数
+	count := int64(0)
+	if v, ok := t.localUsers.Load(userID); ok {
+		count = v.(int64)
+	}
+	count++
+	t.localUsers.Store(userID, count)
+
+	// Redis: HINCRBY + EXPIRE
+	pipe := t.rdb.Pipeline()
+	pipe.HIncrBy(ctx, key, t.instanceID, 1)
+	pipe.Expire(ctx, key, t.onlineTTL)
+	if _, err := pipe.Exec(ctx); err != nil {
+		zap.L().Error("Failed to update online status on connect",
+			zap.String("user_id", userID),
+			zap.Error(err),
+		)
+	}
+}
+
+// UserDisconnected 用户断开连接，更新 Redis 和本地在线状态
+// 返回该用户剩余的连接数
+func (t *OnlineTracker) UserDisconnected(userID string) int {
+	if t.rdb == nil {
+		if v, ok := t.localUsers.Load(userID); ok {
+			count := v.(int64) - 1
+			if count <= 0 {
+				t.localUsers.Delete(userID)
+				return 0
+			}
+			t.localUsers.Store(userID, count)
+			return int(count)
+		}
+		return 0
+	}
+
+	// 本地计数递减
+	count := int64(0)
+	if v, ok := t.localUsers.Load(userID); ok {
+		count = v.(int64) - 1
+	}
+
+	if count <= 0 {
+		t.localUsers.Delete(userID)
+		// Redis: 删除该实例的 field
+		ctx := context.Background()
+		key := fmt.Sprintf("online:%s", userID)
+		t.rdb.HDel(ctx, key, t.instanceID)
+		// 检查 key 是否为空，如果为空则删除
+		exists, _ := t.rdb.HLen(ctx, key).Result()
+		if exists == 0 {
+			t.rdb.Del(ctx, key)
+		}
+		return 0
+	}
+
+	t.localUsers.Store(userID, count)
+
+	// Redis: HINCRBY -1
+	ctx := context.Background()
+	key := fmt.Sprintf("online:%s", userID)
+	newVal, err := t.rdb.HIncrBy(ctx, key, t.instanceID, -1).Result()
+	if err != nil {
+		zap.L().Error("Failed to update online status on disconnect",
+			zap.String("user_id", userID),
+			zap.Error(err),
+		)
+		return int(count)
+	}
+
+	if newVal <= 0 {
+		t.rdb.HDel(ctx, key, t.instanceID)
+		exists, _ := t.rdb.HLen(ctx, key).Result()
+		if exists == 0 {
+			t.rdb.Del(ctx, key)
+		}
+	}
+
+	return int(count)
+}
+
+// IsOnline 检查用户是否在线（查询 Redis）
+func (t *OnlineTracker) IsOnline(userID string) bool {
+	if t.rdb == nil {
+		_, ok := t.localUsers.Load(userID)
+		return ok
+	}
+	ctx := context.Background()
+	key := fmt.Sprintf("online:%s", userID)
+	exists, err := t.rdb.Exists(ctx, key).Result()
+	if err != nil {
+		zap.L().Error("Failed to check online status",
+			zap.String("user_id", userID),
+			zap.Error(err),
+		)
+		// 降级：检查本地
+		_, ok := t.localUsers.Load(userID)
+		return ok
+	}
+	return exists > 0
+}
+
+// FilterOnline 批量检查用户在线状态
+func (t *OnlineTracker) FilterOnline(userIDs []string) (online, offline []string) {
+	if t.rdb == nil {
+		// 降级：只检查本地
+		for _, uid := range userIDs {
+			if _, ok := t.localUsers.Load(uid); ok {
+				online = append(online, uid)
+			} else {
+				offline = append(offline, uid)
+			}
+		}
+		return
+	}
+
+	ctx := context.Background()
+	pipe := t.rdb.Pipeline()
+	cmds := make([]*redis.IntCmd, len(userIDs))
+	for i, uid := range userIDs {
+		key := fmt.Sprintf("online:%s", uid)
+		cmds[i] = pipe.Exists(ctx, key)
+	}
+	if _, err := pipe.Exec(ctx); err != nil {
+		zap.L().Error("Failed to batch check online status", zap.Error(err))
+		// 降级
+		for _, uid := range userIDs {
+			if _, ok := t.localUsers.Load(uid); ok {
+				online = append(online, uid)
+			} else {
+				offline = append(offline, uid)
+			}
+		}
+		return
+	}
+
+	for i, uid := range userIDs {
+		if cmds[i].Val() > 0 {
+			online = append(online, uid)
+		} else {
+			offline = append(offline, uid)
+		}
+	}
+	return
+}
+
+// StartHeartbeat 启动心跳续期
+func (t *OnlineTracker) StartHeartbeat() {
+	if t.rdb == nil {
+		return
+	}
+	ctx, cancel := context.WithCancel(context.Background())
+	t.cancel = cancel
+
+	go func() {
+		ticker := time.NewTicker(t.heartbeatInterval)
+		defer ticker.Stop()
+
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			case <-ticker.C:
+				t.heartbeat()
+			}
+		}
+	}()
+}
+
+func (t *OnlineTracker) heartbeat() {
+	ctx := context.Background()
+	pipe := t.rdb.Pipeline()
+	count := 0
+
+	t.localUsers.Range(func(key, value any) bool {
+		userID := key.(string)
+		onlineKey := fmt.Sprintf("online:%s", userID)
+		pipe.Expire(ctx, onlineKey, t.onlineTTL)
+		count++
+		return true
+	})
+
+	if count > 0 {
+		if _, err := pipe.Exec(ctx); err != nil {
+			zap.L().Error("Failed to heartbeat online keys", zap.Error(err))
+		}
+	}
+}
+
+// Stop 停止心跳并清理本实例的在线状态
+func (t *OnlineTracker) Stop() {
+	if t.cancel != nil {
+		t.cancel()
+	}
+
+	if t.rdb == nil {
+		return
+	}
+
+	ctx := context.Background()
+	pipe := t.rdb.Pipeline()
+
+	// 清理本实例在所有在线 key 中的 field
+	t.localUsers.Range(func(key, value any) bool {
+		userID := key.(string)
+		onlineKey := fmt.Sprintf("online:%s", userID)
+		pipe.HDel(ctx, onlineKey, t.instanceID)
+		// 删除后检查是否为空
+		pipe.HLen(ctx, onlineKey)
+		return true
+	})
+
+	if _, err := pipe.Exec(ctx); err != nil {
+		zap.L().Error("Failed to cleanup online status on stop", zap.Error(err))
+	}
+}
--- a/internal/pkg/ws/publisher.go
+++ b/internal/pkg/ws/publisher.go
@@ -0,0 +1,23 @@
+package ws
+
+// MessagePublisher 定义消息发布接口
+// ws.Hub 隐式满足此接口，ws.Bus 组合 Hub 并扩展跨实例能力
+type MessagePublisher interface {
+	// 消息发布
+	PublishToUser(userID string, eventType string, payload any) Event
+	PublishToUsers(userIDs []string, eventType string, payload any)
+	PublishToUsersConcurrent(userIDs []string, eventType string, payload any)
+	PublishToUserOnline(userID string, eventType string, payload any) bool
+	PublishToUserOnlineReliable(userID string, eventType string, payload any) bool
+	Broadcast(eventType string, payload any)
+
+	// 在线状态
+	HasClients(userID string) bool
+	FilterOnline(userIDs []string) (online, offline []string)
+
+	// 本地操作（不经过 Redis）
+	SendError(client *Client, code string, message string)
+	NextID() uint64
+	OnDisconnect(handler DisconnectHandler)
+	OnConnect(handler ConnectHandler)
+}