junhong_cmp_fiber/internal/polling/scheduler.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(30*time.Second),
			asynq.Queue(constants.QueueDefault),
		)
		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))
				}
			}
		}
	}
}