252 lines
8.5 KiB
Go
252 lines
8.5 KiB
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()
|
||
}
|
||
|
||
// 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()
|
||
}
|
||
|
||
// 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
|
||
}
|