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)) } }