实现task调度器同一时间只能运行同一个plan的同一个task
This commit is contained in:
@@ -17,18 +17,67 @@ type Logger interface {
|
||||
Printf(format string, v ...interface{})
|
||||
}
|
||||
|
||||
// ProgressTracker 在内存中跟踪正在运行的计划的完成进度
|
||||
// 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
|
||||
cond *sync.Cond // 用于实现阻塞锁
|
||||
totalTasks map[uint]int // key: planExecutionLogID, value: total tasks
|
||||
completedTasks map[uint]int // key: planExecutionLogID, value: completed tasks
|
||||
runningPlans map[uint]bool // key: planExecutionLogID, value: true (用作内存锁)
|
||||
}
|
||||
|
||||
func NewProgressTracker() *ProgressTracker {
|
||||
return &ProgressTracker{
|
||||
t := &ProgressTracker{
|
||||
totalTasks: make(map[uint]int),
|
||||
completedTasks: make(map[uint]int),
|
||||
runningPlans: make(map[uint]bool),
|
||||
}
|
||||
t.cond = sync.NewCond(&t.mu)
|
||||
return t
|
||||
}
|
||||
|
||||
// TryLock (非阻塞) 尝试锁定一个计划。如果计划未被锁定，则锁定并返回 true。
|
||||
func (t *ProgressTracker) TryLock(planLogID uint) bool {
|
||||
t.mu.Lock()
|
||||
defer t.mu.Unlock()
|
||||
if t.runningPlans[planLogID] {
|
||||
return false // 已被锁定
|
||||
}
|
||||
t.runningPlans[planLogID] = true
|
||||
return true
|
||||
}
|
||||
|
||||
// Lock (阻塞) 获取一个计划的执行锁。如果锁已被占用，则会一直等待直到锁被释放。
|
||||
func (t *ProgressTracker) Lock(planLogID uint) {
|
||||
t.mu.Lock()
|
||||
// 当计划正在运行时，调用 t.cond.Wait() 会原子地解锁 mu 并挂起当前协程。
|
||||
// 当被唤醒时，它会重新锁定 mu 并再次检查循环条件。
|
||||
for t.runningPlans[planLogID] {
|
||||
t.cond.Wait()
|
||||
}
|
||||
// 获取到锁
|
||||
t.runningPlans[planLogID] = true
|
||||
t.mu.Unlock()
|
||||
}
|
||||
|
||||
// Unlock 解锁一个计划，并唤醒所有正在等待此锁的协程。
|
||||
func (t *ProgressTracker) Unlock(planLogID uint) {
|
||||
t.mu.Lock()
|
||||
defer t.mu.Unlock()
|
||||
delete(t.runningPlans, planLogID)
|
||||
// 唤醒所有在此条件上等待的协程
|
||||
t.cond.Broadcast()
|
||||
}
|
||||
|
||||
// GetRunningPlanIDs 获取当前所有正在执行的计划ID列表
|
||||
func (t *ProgressTracker) GetRunningPlanIDs() []uint {
|
||||
t.mu.Lock()
|
||||
defer t.mu.Unlock()
|
||||
ids := make([]uint, 0, len(t.runningPlans))
|
||||
for id := range t.runningPlans {
|
||||
ids = append(ids, id)
|
||||
}
|
||||
return ids
|
||||
}
|
||||
|
||||
// Scheduler 是核心的、持久化的任务调度器
|
||||
@@ -62,8 +111,6 @@ func NewScheduler(pendingTaskRepo repository.PendingTaskRepository, logger Logge
|
||||
// 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)
|
||||
}))
|
||||
@@ -72,10 +119,8 @@ func (s *Scheduler) Start() {
|
||||
}
|
||||
s.pool = pool
|
||||
|
||||
// 启动主轮询循环
|
||||
s.wg.Add(1)
|
||||
go s.run()
|
||||
|
||||
s.logger.Printf("任务调度器已成功启动")
|
||||
}
|
||||
|
||||
@@ -99,87 +144,62 @@ func (s *Scheduler) run() {
|
||||
case <-s.ctx.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
s.claimAndSubmit()
|
||||
go s.claimAndSubmit()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// claimAndSubmit 认领一个任务并将其提交到 ants 协程池
|
||||
// claimAndSubmit 实现了“认领-锁定-执行 或 等待-放回”的健壮逻辑
|
||||
func (s *Scheduler) claimAndSubmit() {
|
||||
// ants 池的 Running() 数量可以用来提前判断是否繁忙，但这只是一个快照，
|
||||
// 真正的阻塞和背压由 Submit() 方法保证。
|
||||
if s.pool.Running() >= s.workers {
|
||||
// 可选：如果所有 worker 都在忙，可以跳过本次数据库查询，以减轻数据库压力
|
||||
return
|
||||
}
|
||||
runningPlanIDs := s.progressTracker.GetRunningPlanIDs()
|
||||
|
||||
claimedLog, err := s.pendingTaskRepo.ClaimNextDueTask()
|
||||
claimedLog, err := s.pendingTaskRepo.ClaimNextAvailableTask(runningPlanIDs)
|
||||
if err != nil {
|
||||
if !errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
s.logger.Printf("认领任务时发生错误: %v", err)
|
||||
}
|
||||
// gorm.ErrRecordNotFound 说明没任务要执行
|
||||
return
|
||||
}
|
||||
|
||||
// 将任务处理逻辑作为一个函数提交给 ants 池。
|
||||
// 如果池已满，Submit 方法会阻塞，直到有协程空闲出来，这自然地实现了背压。
|
||||
err = s.pool.Submit(func() {
|
||||
s.processTask(claimedLog)
|
||||
})
|
||||
if err != nil {
|
||||
// 如果在调度器停止期间提交任务，可能会发生此错误
|
||||
s.logger.Printf("向协程池提交任务失败: %v", err)
|
||||
// 可以在这里添加逻辑，将任务状态恢复为 pending
|
||||
// 尝试获取内存执行锁
|
||||
if s.progressTracker.TryLock(claimedLog.PlanExecutionLogID) {
|
||||
// 成功获取锁，正常派发任务
|
||||
err = s.pool.Submit(func() {
|
||||
defer s.progressTracker.Unlock(claimedLog.PlanExecutionLogID)
|
||||
s.processTask(claimedLog)
|
||||
})
|
||||
if err != nil {
|
||||
s.logger.Printf("向协程池提交任务失败: %v", err)
|
||||
s.progressTracker.Unlock(claimedLog.PlanExecutionLogID) // 提交失败，必须释放刚刚获取的锁
|
||||
// 同样需要将任务安全放回
|
||||
s.handleRequeue(claimedLog)
|
||||
}
|
||||
} else {
|
||||
// 获取锁失败，说明有“兄弟”任务正在执行。执行“锁定并安全放回”逻辑。
|
||||
s.handleRequeue(claimedLog)
|
||||
}
|
||||
}
|
||||
|
||||
// processTask 包含了处理单个任务的完整逻辑
|
||||
// handleRequeue 处理需要被安全放回队列的任务
|
||||
func (s *Scheduler) handleRequeue(log *models.TaskExecutionLog) {
|
||||
s.logger.Printf("计划 %d 正在执行，任务 %d 将等待并重新入队...", log.PlanExecutionLogID, log.ID)
|
||||
|
||||
// 1. 阻塞式地等待，直到可以获取到该计划的锁
|
||||
s.progressTracker.Lock(log.PlanExecutionLogID)
|
||||
// 2. 在持有锁的情况下，将任务安全地放回队列
|
||||
// 增加一个小的延迟（例如1秒），以避免与主循环发生过于频繁的竞争
|
||||
if err := s.pendingTaskRepo.RequeueTask(log); err != nil {
|
||||
s.logger.Printf("任务重新入队失败, 日志ID: %d, 错误: %v", log.ID, err)
|
||||
}
|
||||
// 3. 释放锁，让其他等待的任务（或主循环）可以继续
|
||||
s.progressTracker.Unlock(log.PlanExecutionLogID)
|
||||
}
|
||||
|
||||
// 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) // 模拟任务执行耗时
|
||||
|
||||
// 任务执行完毕后，更新日志和进度
|
||||
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)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user