280 lines
8.7 KiB
Go
280 lines
8.7 KiB
Go
package polling
|
||
|
||
import (
|
||
"context"
|
||
"sync"
|
||
"time"
|
||
|
||
"github.com/hibiken/asynq"
|
||
"github.com/redis/go-redis/v9"
|
||
"go.uber.org/zap"
|
||
|
||
packagepkg "github.com/break/junhong_cmp_fiber/internal/service/package"
|
||
"github.com/break/junhong_cmp_fiber/pkg/constants"
|
||
)
|
||
|
||
// InitProgress 已迁移到 initializer.go
|
||
|
||
// Scheduler 轮询调度器
|
||
// 职责：读取分片 Sorted Set 中到期的卡，生成 Asynq 任务
|
||
// 不再负责：配置加载、卡初始化（分别由 PollingConfigManager、PollingInitializer 负责）
|
||
type Scheduler struct {
|
||
redis *redis.Client
|
||
queueClient *asynq.Client
|
||
logger *zap.Logger
|
||
queueMgr *PollingQueueManager
|
||
configMgr *PollingConfigManager
|
||
cfg *SchedulerConfig // 启动时固定，避免每次调度重新创建
|
||
|
||
packageActivationHandler *PackageActivationHandler
|
||
dataResetHandler *DataResetHandler
|
||
initializer *PollingInitializer // 可选，nil 表示不守卫 Init 完成
|
||
|
||
stopChan chan struct{}
|
||
wg sync.WaitGroup
|
||
}
|
||
|
||
// SchedulerConfig 调度器配置
|
||
type SchedulerConfig struct {
|
||
ScheduleInterval time.Duration
|
||
MaxManualBatchSize int
|
||
ScheduleBatchSize int
|
||
}
|
||
|
||
// DefaultSchedulerConfig 默认调度器配置
|
||
func DefaultSchedulerConfig() *SchedulerConfig {
|
||
return &SchedulerConfig{
|
||
ScheduleInterval: 1 * time.Second,
|
||
MaxManualBatchSize: 1000,
|
||
ScheduleBatchSize: constants.PollingDequeueMaxBatchSize,
|
||
}
|
||
}
|
||
|
||
// NewScheduler 创建调度器
|
||
func NewScheduler(
|
||
redisClient *redis.Client,
|
||
queueClient *asynq.Client,
|
||
queueMgr *PollingQueueManager,
|
||
configMgr *PollingConfigManager,
|
||
logger *zap.Logger,
|
||
packageActivationHandler *PackageActivationHandler,
|
||
dataResetHandler *DataResetHandler,
|
||
) *Scheduler {
|
||
return &Scheduler{
|
||
redis: redisClient,
|
||
queueClient: queueClient,
|
||
queueMgr: queueMgr,
|
||
configMgr: configMgr,
|
||
logger: logger,
|
||
cfg: DefaultSchedulerConfig(),
|
||
stopChan: make(chan struct{}),
|
||
packageActivationHandler: packageActivationHandler,
|
||
dataResetHandler: dataResetHandler,
|
||
}
|
||
}
|
||
|
||
// Start 启动调度循环（快速启动，配置加载和初始化由外部完成）
|
||
func (s *Scheduler) Start(ctx context.Context) error {
|
||
s.wg.Add(1)
|
||
go s.scheduleLoop(ctx)
|
||
s.logger.Info("轮询调度器已启动")
|
||
return nil
|
||
}
|
||
|
||
// Stop 停止调度器
|
||
func (s *Scheduler) Stop() {
|
||
s.logger.Info("正在停止轮询调度器...")
|
||
close(s.stopChan)
|
||
s.wg.Wait()
|
||
s.logger.Info("轮询调度器已停止")
|
||
}
|
||
|
||
// RefreshConfigs 刷新配置缓存
|
||
func (s *Scheduler) RefreshConfigs(ctx context.Context) error {
|
||
if s.configMgr == nil {
|
||
return nil
|
||
}
|
||
return s.configMgr.Load(ctx)
|
||
}
|
||
|
||
// SetStopResumeCallback 注入停复机回调（在 Start 前调用）
|
||
func (s *Scheduler) SetStopResumeCallback(callback packagepkg.StopResumeCallback) {
|
||
if s.packageActivationHandler != nil {
|
||
s.packageActivationHandler.stopResumeCallback = callback
|
||
}
|
||
}
|
||
|
||
// SetInitializer 注入初始化器（可选，在 Start 前调用）
|
||
// Init 未完成时调度器跳过分片出队，消除启动期无效轮询噪音
|
||
func (s *Scheduler) SetInitializer(init *PollingInitializer) {
|
||
s.initializer = init
|
||
}
|
||
|
||
// writeHeartbeat 写入调度器心跳，表明调度器存活
|
||
func (s *Scheduler) writeHeartbeat(ctx context.Context) {
|
||
if err := s.redis.Set(ctx, constants.RedisPollingSchedulerHeartbeatKey(),
|
||
time.Now().Unix(), 2*s.cfg.ScheduleInterval).Err(); err != nil {
|
||
s.logger.Warn("写入调度器心跳失败", zap.Error(err))
|
||
}
|
||
}
|
||
|
||
// scheduleLoop 调度循环
|
||
func (s *Scheduler) scheduleLoop(ctx context.Context) {
|
||
defer s.wg.Done()
|
||
defer func() {
|
||
if r := recover(); r != nil {
|
||
s.logger.Error("调度主循环发生 panic，调度已停止，需重启 Worker",
|
||
zap.Any("panic", r), zap.Stack("stack"))
|
||
}
|
||
}()
|
||
|
||
ticker := time.NewTicker(s.cfg.ScheduleInterval)
|
||
activationTicker := time.NewTicker(10 * time.Second)
|
||
defer ticker.Stop()
|
||
defer activationTicker.Stop()
|
||
|
||
s.logger.Info("调度循环已启动", zap.Duration("interval", s.cfg.ScheduleInterval))
|
||
|
||
for {
|
||
select {
|
||
case <-s.stopChan:
|
||
s.logger.Info("调度循环收到停止信号")
|
||
return
|
||
case <-ctx.Done():
|
||
s.logger.Info("调度循环收到 ctx 取消信号")
|
||
return
|
||
case <-ticker.C:
|
||
s.writeHeartbeat(ctx)
|
||
s.processShardSchedule(ctx)
|
||
case <-activationTicker.C:
|
||
s.processActivationTasks(ctx)
|
||
}
|
||
}
|
||
}
|
||
|
||
// processShardSchedule 处理手动队列和分片定时队列（每 1 秒触发）
|
||
// 使用 90% 的 tick 间隔作为超时，确保单分片 Redis 挂起时不阻塞下一个 tick
|
||
func (s *Scheduler) processShardSchedule(ctx context.Context) {
|
||
// 手动队列不受 Init 影响（ConfigManager 已就绪即可）
|
||
for _, taskType := range allTaskTypes {
|
||
s.processManualQueue(ctx, taskType, s.cfg.MaxManualBatchSize)
|
||
}
|
||
|
||
// Init 未完成时跳过分片扫描，避免空轮询噪音
|
||
if s.initializer != nil && !s.initializer.IsCompleted() {
|
||
return
|
||
}
|
||
|
||
if s.queueMgr == nil {
|
||
return
|
||
}
|
||
|
||
timeout := s.cfg.ScheduleInterval * 9 / 10
|
||
tickCtx, cancel := context.WithTimeout(ctx, timeout)
|
||
defer cancel()
|
||
|
||
var wg sync.WaitGroup
|
||
for shardID := 0; shardID < s.queueMgr.shardCount; shardID++ {
|
||
wg.Add(1)
|
||
go func(sid int) {
|
||
defer wg.Done()
|
||
defer func() {
|
||
if r := recover(); r != nil {
|
||
s.logger.Error("分片处理 panic，已恢复",
|
||
zap.Int("shard_id", sid), zap.Any("panic", r))
|
||
}
|
||
}()
|
||
s.processOneShard(tickCtx, sid)
|
||
}(shardID)
|
||
}
|
||
wg.Wait()
|
||
}
|
||
|
||
// processOneShard 处理单个分片的所有任务类型出队并推入 Asynq
|
||
func (s *Scheduler) processOneShard(ctx context.Context, shardID int) {
|
||
for _, taskType := range allTaskTypes {
|
||
depth, err := s.queueMgr.GetQueueDepth(ctx, shardID, taskType)
|
||
if err != nil {
|
||
s.logger.Warn("获取分片队列深度失败",
|
||
zap.Int("shard_id", shardID), zap.String("task_type", taskType), zap.Error(err))
|
||
} else if depth > constants.PollingBackpressureThreshold {
|
||
s.logger.Debug("背压：分片队列积压过深，跳过本轮出队",
|
||
zap.Int("shard_id", shardID), zap.String("task_type", taskType), zap.Int64("depth", depth))
|
||
continue
|
||
}
|
||
|
||
entries, err := s.queueMgr.DequeueReady(ctx, shardID, taskType, s.cfg.ScheduleBatchSize)
|
||
if err != nil {
|
||
s.logger.Error("分片出队失败",
|
||
zap.Int("shard_id", shardID), zap.String("task_type", taskType), zap.Error(err))
|
||
continue
|
||
}
|
||
if len(entries) > 0 {
|
||
cardIDs := make([]string, len(entries))
|
||
for i, e := range entries {
|
||
cardIDs[i] = formatUint(e.CardID)
|
||
}
|
||
s.logger.Info("分片出队",
|
||
zap.Int("shard_id", shardID), zap.String("task_type", taskType),
|
||
zap.Int("count", len(entries)))
|
||
s.enqueueBatch(ctx, taskType, cardIDs)
|
||
}
|
||
}
|
||
}
|
||
|
||
// processActivationTasks 套餐激活检查和流量重置调度（每 10 秒触发）
|
||
func (s *Scheduler) processActivationTasks(ctx context.Context) {
|
||
if s.packageActivationHandler != nil {
|
||
if err := s.packageActivationHandler.HandlePackageActivationCheck(ctx); err != nil {
|
||
s.logger.Warn("套餐激活检查失败", zap.Error(err))
|
||
}
|
||
}
|
||
if s.dataResetHandler != nil {
|
||
if err := s.dataResetHandler.HandleDataReset(ctx); err != nil {
|
||
s.logger.Warn("流量重置调度失败", zap.Error(err))
|
||
}
|
||
}
|
||
}
|
||
|
||
// processManualQueue 处理手动触发队列
|
||
func (s *Scheduler) processManualQueue(ctx context.Context, taskType string, maxBatch int) {
|
||
key := constants.RedisPollingManualQueueKey(taskType)
|
||
cardIDs, err := s.redis.LPopCount(ctx, key, maxBatch).Result()
|
||
if err != nil || len(cardIDs) == 0 {
|
||
return
|
||
}
|
||
s.enqueueBatch(ctx, taskType, cardIDs)
|
||
}
|
||
|
||
// enqueueBatch 批量提交任务到 Asynq 队列；入队失败时回退至分片队列防止卡永久丢失
|
||
func (s *Scheduler) enqueueBatch(ctx context.Context, taskType string, cardIDs []string) {
|
||
for _, cardID := range cardIDs {
|
||
payload := map[string]interface{}{
|
||
"card_id": cardID,
|
||
"is_manual": false,
|
||
"timestamp": time.Now().Unix(),
|
||
}
|
||
payloadBytes, marshalErr := marshalJSON(payload)
|
||
if marshalErr != nil {
|
||
s.logger.Error("序列化任务载荷失败，跳过该卡",
|
||
zap.String("task_type", taskType), zap.String("card_id", cardID), zap.Error(marshalErr))
|
||
continue
|
||
}
|
||
task := asynq.NewTask(taskType, payloadBytes,
|
||
asynq.MaxRetry(0),
|
||
asynq.Timeout(60*time.Second),
|
||
asynq.Queue(constants.QueueForTaskType(taskType)),
|
||
)
|
||
if _, err := s.queueClient.Enqueue(task); err != nil {
|
||
s.logger.Error("提交任务失败，回退至分片队列防止卡永久丢失",
|
||
zap.String("task_type", taskType), zap.String("card_id", cardID), zap.Error(err))
|
||
if id, parseErr := parseUint(cardID); parseErr == nil {
|
||
if reqErr := s.queueMgr.Requeue(ctx, id, taskType, time.Now()); reqErr != nil {
|
||
s.logger.Error("回退入队失败，卡可能永久丢失",
|
||
zap.String("card_id", cardID), zap.Error(reqErr))
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|