junhong_cmp_fiber/internal/polling/queue_manager.go

package polling

import (
	"context"
	"fmt"
	"strconv"
	"time"

	"github.com/redis/go-redis/v9"
	"go.uber.org/zap"

	"github.com/break/junhong_cmp_fiber/pkg/constants"
)

// allTaskTypes 轮询系统的全部任务类型（用于 RemoveFromAllQueues 遍历）
var allTaskTypes = []string{
	constants.TaskTypePollingRealname,
	constants.TaskTypePollingCarddata,
	constants.TaskTypePollingPackage,
	constants.TaskTypePollingProtect,
	constants.TaskTypePollingCardStatus,
}

// dequeueScript Lua 脚本：原子出队（ZRANGEBYSCORE + ZREM 服务端原子执行）
// 保留时间过滤语义：只取 score ≤ now 的到期卡，不触碰未来项
// 分批 ZREM：Lua unpack() 受 LUAI_MAXCSTACK 约 8000 限制，按 7000 分批避免溢出
var dequeueScript = redis.NewScript(`
    local results = redis.call('ZRANGEBYSCORE', KEYS[1], '-inf', ARGV[1], 'LIMIT', 0, tonumber(ARGV[2]))
    for i = 1, #results, 7000 do
        local j = math.min(i + 6999, #results)
        redis.call('ZREM', KEYS[1], unpack(results, i, j))
    end
    return results
`)

// CardEntry 出队卡信息
type CardEntry struct {
	CardID uint
}

// PollingQueueManager 统一 Redis 轮询队列操作
// 两个进程（API 进程和 Worker 进程）共享，仅依赖 Redis Client
// 支持分片 Sorted Set，实现千万级规模
type PollingQueueManager struct {
	redis      *redis.Client
	shardCount int
	logger     *zap.Logger
}

// NewPollingQueueManager 创建轮询队列管理器
func NewPollingQueueManager(redisClient *redis.Client, shardCount int, logger *zap.Logger) *PollingQueueManager {
	if shardCount <= 0 {
		shardCount = constants.PollingShardCount
	}
	return &PollingQueueManager{
		redis:      redisClient,
		shardCount: shardCount,
		logger:     logger,
	}
}

// DequeueReady 原子出队到期卡（Lua 脚本：ZRANGEBYSCORE + ZREM 服务端原子执行）
// 只取 score ≤ now 的到期卡，不触碰未来项
// taskType: realname | carddata | package | protect
// shardID: 0 到 shardCount-1
func (m *PollingQueueManager) DequeueReady(ctx context.Context, shardID int, taskType string, batchSize int) ([]CardEntry, error) {
	// 防御：batchSize 不超过 Lua unpack 栈限制
	if batchSize <= 0 || batchSize > constants.PollingDequeueMaxBatchSize {
		batchSize = constants.PollingDequeueMaxBatchSize
	}
	key := constants.RedisPollingShardQueueKey(shardID, taskType)
	now := time.Now().Unix()

	results, err := dequeueScript.Run(ctx, m.redis, []string{key}, now, batchSize).StringSlice()
	if err != nil && err != redis.Nil {
		return nil, err
	}

	entries := make([]CardEntry, 0, len(results))
	for _, s := range results {
		id, parseErr := strconv.ParseUint(s, 10, 64)
		if parseErr != nil {
			m.logger.Warn("解析卡ID失败", zap.String("value", s), zap.Error(parseErr))
			continue
		}
		entries = append(entries, CardEntry{CardID: uint(id)})
	}
	return entries, nil
}

// Requeue 将卡重新入队（ZADD，score 为下次检查时间戳）
func (m *PollingQueueManager) Requeue(ctx context.Context, cardID uint, taskType string, nextCheckAt time.Time) error {
	shardID := int(cardID) % m.shardCount
	key := constants.RedisPollingShardQueueKey(shardID, taskType)
	return m.redis.ZAdd(ctx, key, redis.Z{
		Score:  float64(nextCheckAt.Unix()),
		Member: fmt.Sprintf("%d", cardID),
	}).Err()
}

// EnsureQueued 仅在任务当前不在分片队列中时补入任务，不覆盖已有任务的执行时间。
func (m *PollingQueueManager) EnsureQueued(ctx context.Context, cardID uint, taskType string, nextCheckAt time.Time) (bool, error) {
	shardID := int(cardID) % m.shardCount
	key := constants.RedisPollingShardQueueKey(shardID, taskType)
	added, err := m.redis.ZAddArgs(ctx, key, redis.ZAddArgs{
		NX: true,
		Members: []redis.Z{{
			Score:  float64(nextCheckAt.Unix()),
			Member: fmt.Sprintf("%d", cardID),
		}},
	}).Result()
	return added > 0, err
}

// RemoveFromAllQueues 从所有分片的所有5个队列（realname/carddata/package/protect/card_status）移除指定卡
// 修复 Bug3：旧实现漏掉 protect 队列
func (m *PollingQueueManager) RemoveFromAllQueues(ctx context.Context, cardID uint) error {
	member := fmt.Sprintf("%d", cardID)
	pipe := m.redis.Pipeline()
	for i := 0; i < m.shardCount; i++ {
		for _, taskType := range allTaskTypes {
			key := constants.RedisPollingShardQueueKey(i, taskType)
			pipe.ZRem(ctx, key, member)
		}
	}
	_, err := pipe.Exec(ctx)
	return err
}

// RemoveBatchFromCurrentShardQueues 批量移除卡在当前分片队列中的旧成员，并清理卡缓存。
// 迁移收尾重建队列前使用，避免重复执行时残留旧任务类型或旧缓存。
func (m *PollingQueueManager) RemoveBatchFromCurrentShardQueues(ctx context.Context, cardIDs []uint) error {
	if len(cardIDs) == 0 {
		return nil
	}
	pipe := m.redis.Pipeline()
	for _, cardID := range cardIDs {
		member := fmt.Sprintf("%d", cardID)
		shardID := int(cardID) % m.shardCount
		for _, taskType := range allTaskTypes {
			key := constants.RedisPollingShardQueueKey(shardID, taskType)
			pipe.ZRem(ctx, key, member)
		}
		pipe.Del(ctx, constants.RedisPollingCardInfoKey(cardID))
	}
	_, err := pipe.Exec(ctx)
	return err
}

// EnqueueManual 手动触发入队（List RPUSH，调度器优先消费）
func (m *PollingQueueManager) EnqueueManual(ctx context.Context, cardID uint, taskType string) error {
	key := constants.RedisPollingManualQueueKey(taskType)
	return m.redis.RPush(ctx, key, fmt.Sprintf("%d", cardID)).Err()
}

// EnqueuePriority 优先轮询提示下发（List RPUSH，与手动触发队列分离）。
// 使用 RPUSH 与调度器的 LPopCount 组成先进先出；提示通道不是权威，丢失或缓存服务重启时
// 库内活动优先项仍在，后续普通轮询执行按条件认领仍会执行，只退化为延迟一个普通轮询周期。
func (m *PollingQueueManager) EnqueuePriority(ctx context.Context, cardID uint, taskType string) error {
	key := constants.RedisPollingPriorityQueueKey(taskType)
	return m.redis.RPush(ctx, key, fmt.Sprintf("%d", cardID)).Err()
}

// OnCardDeleted 卡删除事件处理（移除所有队列 + 清理卡信息缓存）
func (m *PollingQueueManager) OnCardDeleted(ctx context.Context, cardID uint) error {
	// 从所有分片队列移除
	if err := m.RemoveFromAllQueues(ctx, cardID); err != nil {
		return err
	}
	// 清理轮询卡信息缓存
	cacheKey := constants.RedisPollingCardInfoKey(cardID)
	return m.redis.Del(ctx, cacheKey).Err()
}

// InvalidateCardCache 清理轮询卡信息缓存，强制下次轮询从 DB 重建
func (m *PollingQueueManager) InvalidateCardCache(ctx context.Context, cardID uint) {
	cacheKey := constants.RedisPollingCardInfoKey(cardID)
	if err := m.redis.Del(ctx, cacheKey).Err(); err != nil {
		m.logger.Warn("清理轮询卡缓存失败", zap.Uint("card_id", cardID), zap.Error(err))
	}
}

// GetQueueDepth 获取分片队列深度（用于背压检测）
func (m *PollingQueueManager) GetQueueDepth(ctx context.Context, shardID int, taskType string) (int64, error) {
	key := constants.RedisPollingShardQueueKey(shardID, taskType)
	return m.redis.ZCard(ctx, key).Result()
}

// GetTotalQueueDepth 获取指定任务类型的总队列深度（聚合所有分片）
// 供 MonitoringService 使用，替代直接读取旧的非分片 Redis Key
// 若任意分片查询失败，返回已累计的部分总量和第一个错误，调用方可据此判断数据完整性
func (m *PollingQueueManager) GetTotalQueueDepth(ctx context.Context, taskType string) (int64, error) {
	var total int64
	var firstErr error
	for i := 0; i < m.shardCount; i++ {
		depth, err := m.GetQueueDepth(ctx, i, taskType)
		if err != nil {
			m.logger.Warn("获取分片队列深度失败",
				zap.Int("shard_id", i),
				zap.String("task_type", taskType),
				zap.Error(err))
			if firstErr == nil {
				firstErr = err
			}
			continue
		}
		total += depth
	}
	return total, firstErr
}

// GetTotalDueCount 获取指定任务类型所有分片中已到期的队列数量。
// 只统计 score <= now 的成员，用于后台监控展示真实待调度积压。
func (m *PollingQueueManager) GetTotalDueCount(ctx context.Context, taskType string, now time.Time) (int64, error) {
	var total int64
	var firstErr error
	for i := 0; i < m.shardCount; i++ {
		key := constants.RedisPollingShardQueueKey(i, taskType)
		count, err := m.redis.ZCount(ctx, key, "-inf", fmt.Sprintf("%d", now.Unix())).Result()
		if err != nil {
			m.logger.Warn("获取分片到期数量失败",
				zap.Int("shard_id", i),
				zap.String("task_type", taskType),
				zap.Error(err))
			if firstErr == nil {
				firstErr = err
			}
			continue
		}
		total += count
	}
	return total, firstErr
}

// GetAverageWaitTime 获取指定任务类型所有分片最早若干任务的平均等待秒数。
// 每个分片最多取 samplePerShard 条，避免监控接口扫描完整队列。
func (m *PollingQueueManager) GetAverageWaitTime(ctx context.Context, taskType string, now time.Time, samplePerShard int64) (float64, error) {
	if samplePerShard <= 0 {
		samplePerShard = 10
	}

	var totalWait float64
	var sampleCount int64
	var firstErr error
	nowUnix := now.Unix()

	for i := 0; i < m.shardCount; i++ {
		key := constants.RedisPollingShardQueueKey(i, taskType)
		earliest, err := m.redis.ZRangeWithScores(ctx, key, 0, samplePerShard-1).Result()
		if err != nil {
			m.logger.Warn("获取分片最早任务失败",
				zap.Int("shard_id", i),
				zap.String("task_type", taskType),
				zap.Error(err))
			if firstErr == nil {
				firstErr = err
			}
			continue
		}
		for _, z := range earliest {
			waitTime := float64(nowUnix) - z.Score
			if waitTime > 0 {
				totalWait += waitTime
			}
			sampleCount++
		}
	}

	if sampleCount == 0 {
		return 0, firstErr
	}
	return totalWait / float64(sampleCount), firstErr
}