junhong_cmp_fiber/internal/polling/recalculate.go

package polling

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

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

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

// recalculateBatchSize 是重算到期时间每批读取的卡数。
const recalculateBatchSize = 500

// RecalculateProgress 是重算到期时间任务的进度快照。
type RecalculateProgress struct {
	Total     int
	Processed int
	Advanced  int // 被提前的检查项数
}

// DueRecalculator 按当前轮询配置重算已在队列中的卡的下次到期时间：新到期 = min(当前到期, 上次检查 + 新间隔)，只提前不推后。
// 只更新已在分片队列中的成员（ZADD XX LT），不新增成员，避免把执行中的卡重复放回队列；不调用网关。
type DueRecalculator struct {
	db        *gorm.DB
	redis     *redis.Client
	configMgr *PollingConfigManager
	queueMgr  *PollingQueueManager
	logger    *zap.Logger
}

// NewDueRecalculator 创建重算器。
func NewDueRecalculator(db *gorm.DB, redisClient *redis.Client, configMgr *PollingConfigManager, queueMgr *PollingQueueManager, logger *zap.Logger) *DueRecalculator {
	if logger == nil {
		logger = zap.NewNop()
	}
	return &DueRecalculator{db: db, redis: redisClient, configMgr: configMgr, queueMgr: queueMgr, logger: logger}
}

// Run 执行重算；carrierID 非空时只处理该运营商账户的卡。每处理一批回调一次进度。
func (r *DueRecalculator) Run(ctx context.Context, carrierID *uint, report func(RecalculateProgress)) (RecalculateProgress, error) {
	var progress RecalculateProgress
	if r == nil || r.db == nil || r.redis == nil || r.configMgr == nil || r.queueMgr == nil {
		return progress, fmt.Errorf("重算到期时间能力未配置")
	}
	if err := r.configMgr.Load(ctx); err != nil {
		return progress, err
	}
	base := r.db.WithContext(ctx).Model(&model.IotCard{}).Where("enable_polling = true AND deleted_at IS NULL")
	if carrierID != nil {
		base = base.Where("carrier_id = ?", *carrierID)
	}
	var total int64
	if err := base.Count(&total).Error; err != nil {
		return progress, err
	}
	progress.Total = int(total)
	if report != nil {
		report(progress)
	}
	var lastID uint
	for {
		if err := ctx.Err(); err != nil {
			return progress, err
		}
		var cards []*model.IotCard
		query := r.db.WithContext(ctx).Where("id > ? AND enable_polling = true AND deleted_at IS NULL", lastID)
		if carrierID != nil {
			query = query.Where("carrier_id = ?", *carrierID)
		}
		if err := query.Order("id ASC").Limit(recalculateBatchSize).Find(&cards).Error; err != nil {
			return progress, err
		}
		if len(cards) == 0 {
			break
		}
		lastID = cards[len(cards)-1].ID
		advanced, err := r.recalculateBatch(ctx, cards)
		if err != nil {
			return progress, err
		}
		progress.Processed += len(cards)
		progress.Advanced += advanced
		if report != nil {
			report(progress)
		}
	}
	return progress, nil
}

func (r *DueRecalculator) recalculateBatch(ctx context.Context, cards []*model.IotCard) (int, error) {
	if tier := r.configMgr.TierResolver(); tier != nil {
		tier.Prefill(ctx, cards)
	}
	now := time.Now()
	pipe := r.redis.Pipeline()
	var cmds []*redis.IntCmd
	for _, card := range cards {
		shardID := int(card.ID) % r.queueMgr.shardCount
		member := strconv.FormatUint(uint64(card.ID), 10)
		for taskType, info := range r.configMgr.MergedTaskIntervals(card) {
			due := recalculatedDue(lastCheckAtByTaskType(card, taskType), info.Interval, now)
			if windowDue := r.configMgr.NextCheckAt(ctx, card, taskType, info, now); windowDue.Before(due) && windowDue.After(now) {
				due = windowDue
			}
			cmds = append(cmds, pipe.ZAddArgs(ctx, constants.RedisPollingShardQueueKey(shardID, taskType), redis.ZAddArgs{
				XX: true, LT: true, Ch: true, Members: []redis.Z{{Score: float64(due.Unix()), Member: member}},
			}))
		}
	}
	if len(cmds) == 0 {
		return 0, nil
	}
	if _, err := pipe.Exec(ctx); err != nil && err != redis.Nil {
		return 0, err
	}
	advanced := 0
	for _, cmd := range cmds {
		if n, err := cmd.Result(); err == nil {
			advanced += int(n)
		}
	}
	return advanced, nil
}

// recalculatedDue 返回重算的候选到期时间。上次检查时间只在查询成功时更新，查询失败或被限流延后的卡按
// "上次检查 + 新间隔"会算成已过期；此时候选改为"现在 + 新间隔"，只提前排队超过一个新间隔的成员，
// 不抵消失败退避与限流延后，也避免重算后集中触发调用。
func recalculatedDue(lastCheckAt *time.Time, intervalSeconds int, now time.Time) time.Time {
	interval := time.Duration(intervalSeconds) * time.Second
	if lastCheckAt == nil {
		return now.Add(interval)
	}
	if due := lastCheckAt.Add(interval); due.After(now) {
		return due
	}
	return now.Add(interval)
}