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