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

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

// RemoveFromAllQueues 从所有分片的所有4个队列（含protect）移除指定卡
// 修复 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
}

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

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

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