task调度器该用ants线程池
This commit is contained in:
@@ -8,6 +8,7 @@ import (
|
||||
|
||||
"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"
|
||||
)
|
||||
|
||||
@@ -30,36 +31,6 @@ func NewProgressTracker() *ProgressTracker {
|
||||
}
|
||||
}
|
||||
|
||||
// StartTracking 开始跟踪一个新的计划执行
|
||||
func (t *ProgressTracker) StartTracking(planLogID uint, total int) {
|
||||
t.mu.Lock()
|
||||
defer t.mu.Unlock()
|
||||
t.totalTasks[planLogID] = total
|
||||
t.completedTasks[planLogID] = 0
|
||||
}
|
||||
|
||||
// Increment 将指定计划的完成计数加一
|
||||
func (t *ProgressTracker) Increment(planLogID uint) {
|
||||
t.mu.Lock()
|
||||
defer t.mu.Unlock()
|
||||
t.completedTasks[planLogID]++
|
||||
}
|
||||
|
||||
// IsComplete 检查指定计划是否已完成所有任务
|
||||
func (t *ProgressTracker) IsComplete(planLogID uint) bool {
|
||||
t.mu.Lock()
|
||||
defer t.mu.Unlock()
|
||||
return t.completedTasks[planLogID] >= t.totalTasks[planLogID]
|
||||
}
|
||||
|
||||
// StopTracking 停止跟踪一个计划，清理内存
|
||||
func (t *ProgressTracker) StopTracking(planLogID uint) {
|
||||
t.mu.Lock()
|
||||
defer t.mu.Unlock()
|
||||
delete(t.totalTasks, planLogID)
|
||||
delete(t.completedTasks, planLogID)
|
||||
}
|
||||
|
||||
// Scheduler 是核心的、持久化的任务调度器
|
||||
type Scheduler struct {
|
||||
logger Logger
|
||||
@@ -68,10 +39,10 @@ type Scheduler struct {
|
||||
pendingTaskRepo repository.PendingTaskRepository
|
||||
progressTracker *ProgressTracker
|
||||
|
||||
taskChannel chan *models.TaskExecutionLog // 用于向 workers 派发任务的缓冲 channel
|
||||
wg sync.WaitGroup
|
||||
ctx context.Context
|
||||
cancel context.CancelFunc
|
||||
pool *ants.Pool // 使用 ants 协程池来管理并发
|
||||
wg sync.WaitGroup
|
||||
ctx context.Context
|
||||
cancel context.CancelFunc
|
||||
}
|
||||
|
||||
// NewScheduler 创建一个新的调度器实例
|
||||
@@ -83,21 +54,23 @@ func NewScheduler(pendingTaskRepo repository.PendingTaskRepository, logger Logge
|
||||
pollingInterval: interval,
|
||||
workers: numWorkers,
|
||||
progressTracker: NewProgressTracker(),
|
||||
taskChannel: make(chan *models.TaskExecutionLog, numWorkers), // 缓冲大小与 worker 数量一致
|
||||
ctx: ctx,
|
||||
cancel: cancel,
|
||||
}
|
||||
}
|
||||
|
||||
// Start 启动调度器，包括主轮询循环和所有工作协程
|
||||
// Start 启动调度器，包括初始化协程池和启动主轮询循环
|
||||
func (s *Scheduler) Start() {
|
||||
s.logger.Printf("任务调度器正在启动，工作协程数: %d...", s.workers)
|
||||
|
||||
// 启动工作协程池
|
||||
s.wg.Add(s.workers)
|
||||
for i := 0; i < s.workers; i++ {
|
||||
go s.worker(i)
|
||||
// 初始化 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)
|
||||
@@ -106,15 +79,16 @@ func (s *Scheduler) Start() {
|
||||
s.logger.Printf("任务调度器已成功启动")
|
||||
}
|
||||
|
||||
// Stop 优雅地停止调度器和所有工作协程
|
||||
// Stop 优雅地停止调度器
|
||||
func (s *Scheduler) Stop() {
|
||||
s.logger.Printf("正在停止任务调度器...")
|
||||
s.cancel() // 发出取消信号
|
||||
s.wg.Wait() // 等待所有协程完成
|
||||
s.cancel() // 1. 发出取消信号，停止主循环
|
||||
s.wg.Wait() // 2. 等待主循环完成
|
||||
s.pool.Release() // 3. 释放 ants 池 (等待所有已提交的任务执行完毕)
|
||||
s.logger.Printf("任务调度器已安全停止")
|
||||
}
|
||||
|
||||
// run 是主轮询循环，负责从数据库认领任务并派发
|
||||
// run 是主轮询循环，负责从数据库认领任务并提交到协程池
|
||||
func (s *Scheduler) run() {
|
||||
defer s.wg.Done()
|
||||
ticker := time.NewTicker(s.pollingInterval)
|
||||
@@ -123,16 +97,22 @@ func (s *Scheduler) run() {
|
||||
for {
|
||||
select {
|
||||
case <-s.ctx.Done():
|
||||
close(s.taskChannel) // 关闭 channel，让 workers 退出循环
|
||||
return
|
||||
case <-ticker.C:
|
||||
s.claimAndDispatch()
|
||||
s.claimAndSubmit()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// claimAndDispatch 认领一个任务并将其发送到派发通道
|
||||
func (s *Scheduler) claimAndDispatch() {
|
||||
// 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) {
|
||||
@@ -141,39 +121,29 @@ func (s *Scheduler) claimAndDispatch() {
|
||||
return
|
||||
}
|
||||
|
||||
// 将认领到的任务发送到派发通道
|
||||
// 如果所有 worker 都在忙，这里会阻塞，从而实现背压，防止队列无限增长
|
||||
select {
|
||||
case s.taskChannel <- claimedLog:
|
||||
s.logger.Printf("成功认领并派发任务, 日志ID: %d, 任务ID: %d", claimedLog.ID, claimedLog.TaskID)
|
||||
case <-s.ctx.Done():
|
||||
// 如果在等待派发时调度器被停止，需要处理这个未派发的任务
|
||||
// 简单的处理方式是忽略它，让清理器进程后续来处理这个 'running' 状态的任务
|
||||
s.logger.Printf("在派发任务时调度器被停止, 日志ID: %d", claimedLog.ID)
|
||||
// 将任务处理逻辑作为一个函数提交给 ants 池。
|
||||
// 如果池已满，Submit 方法会阻塞，直到有协程空闲出来，这自然地实现了背压。
|
||||
err = s.pool.Submit(func() {
|
||||
s.processTask(claimedLog)
|
||||
})
|
||||
if err != nil {
|
||||
// 如果在调度器停止期间提交任务，可能会发生此错误
|
||||
s.logger.Printf("向协程池提交任务失败: %v", err)
|
||||
// 可以在这里添加逻辑，将任务状态恢复为 pending
|
||||
}
|
||||
}
|
||||
|
||||
// worker 是工作协程的实现
|
||||
func (s *Scheduler) worker(id int) {
|
||||
defer s.wg.Done()
|
||||
s.logger.Printf("工作协程 #%d 已启动", id)
|
||||
for claimedLog := range s.taskChannel {
|
||||
s.processTask(id, claimedLog)
|
||||
}
|
||||
s.logger.Printf("工作协程 #%d 已停止", id)
|
||||
}
|
||||
|
||||
// processTask 包含了处理单个任务的完整逻辑
|
||||
func (s *Scheduler) processTask(workerID int, claimedLog *models.TaskExecutionLog) {
|
||||
s.logger.Printf("工作协程 #%d 正在处理任务, 日志ID: %d, 任务ID: %d, 任务名称: %s",
|
||||
workerID, claimedLog.ID, claimedLog.TaskID, claimedLog.Task.Name)
|
||||
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("工作协程 #%d 已完成任务, 日志ID: %d", workerID, claimedLog.ID)
|
||||
s.logger.Printf("完成任务, 日志ID: %d", claimedLog.ID)
|
||||
|
||||
// ----------------------------------------------------
|
||||
// 未来的逻辑将在这里展开：
|
||||
@@ -186,3 +156,30 @@ func (s *Scheduler) processTask(workerID int, claimedLog *models.TaskExecutionLo
|
||||
//
|
||||
// ----------------------------------------------------
|
||||
}
|
||||
|
||||
// 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)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user