backend/internal/pkg/ws/online_tracker.go

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