OpenFlare/backend/plugins/drivers/driver_asynq_cron/scheduler.go

// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0

package driver_asynq_cron

import (
	"Wavelet/core/contracts"
	"Wavelet/pkg/logger"
	"context"
	"fmt"
	"os/signal"
	"sync"
	"syscall"
	"time"

	"github.com/hibiken/asynq"
)

var (
	activeScheduler *asynq.Scheduler
	schedulerMutex  sync.Mutex
	quitChan        chan struct{}
	schedulerOnce   sync.Once
	taskSvcMu       sync.RWMutex
	taskSvcInstance contracts.TaskService
	// RedisOpt is the Redis connection option for the scheduler.
	RedisOpt asynq.RedisConnOpt
)

func setTaskService(s contracts.TaskService) {
	taskSvcMu.Lock()
	defer taskSvcMu.Unlock()
	taskSvcInstance = s
}

func getTaskService() contracts.TaskService {
	taskSvcMu.RLock()
	defer taskSvcMu.RUnlock()
	return taskSvcInstance
}

// StartScheduler 启动调度器 (该函数阻塞，直到调度器退出)
func StartScheduler() error {
	var err error
	schedulerOnce.Do(func() {
		quitChan = make(chan struct{})
		done := quitChan

		// 初始化并运行首次调度
		if err = ReloadScheduler(); err != nil {
			err = fmt.Errorf("initial reload failed: %w", err)
			return
		}

		signalCtx, stopSignals := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
		defer stopSignals()

		if waitForStop(done, signalCtx.Done()) {
			StopScheduler()
		}
	})
	return err
}

// StopScheduler 停止调度服务并解除 StartScheduler 阻塞
func StopScheduler() {
	schedulerMutex.Lock()
	defer schedulerMutex.Unlock()

	if activeScheduler != nil {
		activeScheduler.Shutdown()
		activeScheduler = nil
	}

	if quitChan != nil {
		close(quitChan)
		quitChan = nil
	}
}

// ReloadScheduler 重载调度器配置 (线程安全)
func ReloadScheduler() error {
	schedulerMutex.Lock()
	defer schedulerMutex.Unlock()

	// 1. 如果有运行中的调度器，先关闭它
	if activeScheduler != nil {
		activeScheduler.Shutdown()
		activeScheduler = nil
	}

	// 2. 从数据库载入启用的定时任务配置
	schedules, err := ListActiveSchedules(context.Background())
	if err != nil {
		return fmt.Errorf("load schedules from db failed: %w", err)
	}

	location, err := time.LoadLocation("Asia/Shanghai")
	if err != nil {
		return fmt.Errorf(errLoadLocationFailed, err)
	}

	// 3. 实例化新的调度器
	newScheduler := asynq.NewScheduler(
		RedisOpt,
		&asynq.SchedulerOpts{
			Location: location,
		},
	)

	// 4. 遍历并注册任务
	taskSvc := getTaskService()
	for _, s := range schedules {
		taskName, maxRetry, queue := resolveTaskScheduleMeta(taskSvc, s.TaskType)

		// 构造 Asynq 载荷
		t := asynq.NewTask(taskName, []byte(s.Payload))

		if _, err := newScheduler.Register(
			s.Cron,
			t,
			asynq.MaxRetry(maxRetry),
			asynq.Queue(queue),
		); err != nil {
			// 定时任务配置可能有误（如 Cron 格式不被 Asynq 识别），记录日志并跳过
			logger.ErrorF(context.Background(), "[Scheduler] 注册定时任务失败 id=%d name=%s: %v", s.ID, s.Name, err)
			continue
		}
	}

	// 5. 启动并替换全局调度器。进程信号由 StartScheduler 统一处理。
	if err := newScheduler.Start(); err != nil {
		return fmt.Errorf("start scheduler failed: %w", err)
	}
	activeScheduler = newScheduler

	logger.InfoF(context.Background(), "[Scheduler] 成功重新加载定时任务，共注册 %d 个活动任务", len(schedules))
	return nil
}

func waitForStop(done, signals <-chan struct{}) bool {
	select {
	case <-done:
		return false
	case <-signals:
		return true
	}
}

func resolveTaskScheduleMeta(taskSvc contracts.TaskService, taskType string) (name string, maxRetry int, queue string) {
	name = taskType
	maxRetry = 3
	queue = "default"
	if taskSvc == nil {
		return
	}
	meta, ok := taskSvc.GetTaskMeta(taskType)
	if !ok {
		return
	}
	name = meta.Name
	if meta.MaxRetry > 0 {
		maxRetry = meta.MaxRetry
	}
	if meta.Queue != "" {
		queue = meta.Queue
	}
	return
}