pig-farm-controller/internal/infra/task/task.go

package task

import (
	"context"
	"errors"
	"sync"
	"time"

	"git.huangwc.com/pig/pig-farm-controller/internal/infra/models"
	"git.huangwc.com/pig/pig-farm-controller/internal/infra/repository"
	"github.com/panjf2000/ants/v2"
	"gorm.io/gorm"
)

// Logger 定义了调度器期望的日志接口，方便替换为项目中的日志组件
type Logger interface {
	Printf(format string, v ...interface{})
}

// ProgressTracker 在内存中跟踪正在运行的计划的完成进度
type ProgressTracker struct {
	mu             sync.Mutex
	totalTasks     map[uint]int // key: planExecutionLogID, value: total tasks
	completedTasks map[uint]int // key: planExecutionLogID, value: completed tasks
}

func NewProgressTracker() *ProgressTracker {
	return &ProgressTracker{
		totalTasks:     make(map[uint]int),
		completedTasks: make(map[uint]int),
	}
}

// Scheduler 是核心的、持久化的任务调度器
type Scheduler struct {
	logger          Logger
	pollingInterval time.Duration
	workers         int
	pendingTaskRepo repository.PendingTaskRepository
	progressTracker *ProgressTracker

	pool   *ants.Pool // 使用 ants 协程池来管理并发
	wg     sync.WaitGroup
	ctx    context.Context
	cancel context.CancelFunc
}

// NewScheduler 创建一个新的调度器实例
func NewScheduler(pendingTaskRepo repository.PendingTaskRepository, logger Logger, interval time.Duration, numWorkers int) *Scheduler {
	ctx, cancel := context.WithCancel(context.Background())
	return &Scheduler{
		pendingTaskRepo: pendingTaskRepo,
		logger:          logger,
		pollingInterval: interval,
		workers:         numWorkers,
		progressTracker: NewProgressTracker(),
		ctx:             ctx,
		cancel:          cancel,
	}
}

// Start 启动调度器，包括初始化协程池和启动主轮询循环
func (s *Scheduler) Start() {
	s.logger.Printf("任务调度器正在启动，工作协程数: %d...", s.workers)

	// 初始化 ants 协程池
	pool, err := ants.NewPool(s.workers, ants.WithPanicHandler(func(err interface{}) {
		s.logger.Printf("[严重] 任务执行时发生 panic: %v", err)
	}))
	if err != nil {
		panic("初始化协程池失败: " + err.Error())
	}
	s.pool = pool

	// 启动主轮询循环
	s.wg.Add(1)
	go s.run()

	s.logger.Printf("任务调度器已成功启动")
}

// Stop 优雅地停止调度器
func (s *Scheduler) Stop() {
	s.logger.Printf("正在停止任务调度器...")
	s.cancel()       // 1. 发出取消信号，停止主循环
	s.wg.Wait()      // 2. 等待主循环完成
	s.pool.Release() // 3. 释放 ants 池 (等待所有已提交的任务执行完毕)
	s.logger.Printf("任务调度器已安全停止")
}

// run 是主轮询循环，负责从数据库认领任务并提交到协程池
func (s *Scheduler) run() {
	defer s.wg.Done()
	ticker := time.NewTicker(s.pollingInterval)
	defer ticker.Stop()

	for {
		select {
		case <-s.ctx.Done():
			return
		case <-ticker.C:
			s.claimAndSubmit()
		}
	}
}

// claimAndSubmit 认领一个任务并将其提交到 ants 协程池
func (s *Scheduler) claimAndSubmit() {
	// ants 池的 Running() 数量可以用来提前判断是否繁忙，但这只是一个快照，
	// 真正的阻塞和背压由 Submit() 方法保证。
	if s.pool.Running() >= s.workers {
		// 可选：如果所有 worker 都在忙，可以跳过本次数据库查询，以减轻数据库压力
		return
	}

	claimedLog, err := s.pendingTaskRepo.ClaimNextDueTask()
	if err != nil {
		if !errors.Is(err, gorm.ErrRecordNotFound) {
			s.logger.Printf("认领任务时发生错误: %v", err)
		}
		return
	}

	// 将任务处理逻辑作为一个函数提交给 ants 池。
	// 如果池已满，Submit 方法会阻塞，直到有协程空闲出来，这自然地实现了背压。
	err = s.pool.Submit(func() {
		s.processTask(claimedLog)
	})
	if err != nil {
		// 如果在调度器停止期间提交任务，可能会发生此错误
		s.logger.Printf("向协程池提交任务失败: %v", err)
		// 可以在这里添加逻辑，将任务状态恢复为 pending
	}
}

// processTask 包含了处理单个任务的完整逻辑
func (s *Scheduler) processTask(claimedLog *models.TaskExecutionLog) {
	s.logger.Printf("开始处理任务, 日志ID: %d, 任务ID: %d, 任务名称: %s",
		claimedLog.ID, claimedLog.TaskID, claimedLog.Task.Name)

	// 在这里，我们将根据 claimedLog.TaskID 或未来的 Task.Kind 来分发给不同的处理器
	// 现在，我们只做一个模拟执行
	time.Sleep(2 * time.Second) // 模拟任务执行耗时

	// 任务执行完毕后，更新日志和进度
	s.logger.Printf("完成任务, 日志ID: %d", claimedLog.ID)

	// ----------------------------------------------------
	// 未来的逻辑将在这里展开：
	//
	// 1. 调用 handler.Handle(claimedLog)
	// 2. 根据 handler 返回的 error 更新日志为 'completed' 或 'failed'
	//    execLogRepo.UpdateTaskExecutionLog(...)
	// 3. 如果成功，则 s.progressTracker.Increment(claimedLog.PlanExecutionLogID)
	// 4. 检查 s.progressTracker.IsComplete(...)，如果完成则执行计划收尾工作
	//
	// ----------------------------------------------------
}

// ProgressTracker 的方法实现
func (t *ProgressTracker) StartTracking(planLogID uint, total int) {
	t.mu.Lock()
	defer t.mu.Unlock()
	t.totalTasks[planLogID] = total
	t.completedTasks[planLogID] = 0
}

func (t *ProgressTracker) Increment(planLogID uint) {
	t.mu.Lock()
	defer t.mu.Unlock()
	t.completedTasks[planLogID]++
}

func (t *ProgressTracker) IsComplete(planLogID uint) bool {
	t.mu.Lock()
	defer t.mu.Unlock()
	return t.completedTasks[planLogID] >= t.totalTasks[planLogID]
}

func (t *ProgressTracker) StopTracking(planLogID uint) {
	t.mu.Lock()
	defer t.mu.Unlock()
	delete(t.totalTasks, planLogID)
	delete(t.completedTasks, planLogID)
}