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.
This commit is contained in:
257
internal/pkg/ws/online_tracker.go
Normal file
257
internal/pkg/ws/online_tracker.go
Normal file
@@ -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))
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user