From 2402c206dc13c161ef46075818f98f5757857c68 Mon Sep 17 00:00:00 2001 From: huang <1724659546@qq.com> Date: Wed, 17 Sep 2025 14:58:41 +0800 Subject: [PATCH] =?UTF-8?q?=E5=AE=9E=E7=8E=B0task=E8=B0=83=E5=BA=A6?= =?UTF-8?q?=E5=99=A8=E5=90=8C=E4=B8=80=E6=97=B6=E9=97=B4=E5=8F=AA=E8=83=BD?= =?UTF-8?q?=E8=BF=90=E8=A1=8C=E5=90=8C=E4=B8=80=E4=B8=AAplan=E7=9A=84?= =?UTF-8?q?=E5=90=8C=E4=B8=80=E4=B8=AAtask?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- TODO-List | 1 + .../repository/pending_task_repository.go | 54 +++--- internal/infra/task/task.go | 162 ++++++++++-------- 3 files changed, 125 insertions(+), 92 deletions(-) diff --git a/TODO-List b/TODO-List index 016fcd4..00bfc63 100644 --- a/TODO-List +++ b/TODO-List @@ -1,6 +1,7 @@ // TODO 列表 // TODO 可以实现的问题 +plan执行到一半时如果用户删掉里面的task, 后续调度器执行task时可能会找不到这个任务的细节 1. 可以用TimescaleDB代替PGSQL, 优化传感器数据存储性能 2. 系统启动时应该检查一遍执行历史库, 将所有显示为执行中的任务都修正为执行失败并报错 diff --git a/internal/infra/repository/pending_task_repository.go b/internal/infra/repository/pending_task_repository.go index 3a66c96..a5b57a9 100644 --- a/internal/infra/repository/pending_task_repository.go +++ b/internal/infra/repository/pending_task_repository.go @@ -11,7 +11,8 @@ import ( // PendingTaskRepository 定义了与待执行任务队列交互的接口。 type PendingTaskRepository interface { CreatePendingTasksInBatch(tasks []*models.PendingTask) error - ClaimNextDueTask() (*models.TaskExecutionLog, error) + ClaimNextAvailableTask(excludePlanIDs []uint) (*models.TaskExecutionLog, error) + RequeueTask(log *models.TaskExecutionLog) error } // pendingTaskRepository 是使用 GORM 的具体实现。 @@ -20,7 +21,6 @@ type pendingTaskRepository struct { } // NewPendingTaskRepository 创建一个新的待执行任务队列仓库。 -// 它接收一个 GORM DB 实例作为依赖。 func NewPendingTaskRepository(db *gorm.DB) PendingTaskRepository { return &pendingTaskRepository{db: db} } @@ -30,32 +30,28 @@ func (r *pendingTaskRepository) CreatePendingTasksInBatch(tasks []*models.Pendin return r.db.Create(&tasks).Error } -// ClaimNextDueTask 以原子方式认领下一个到期任务。 -// 它在一个事务中完成三件事： -// 1. 查找并锁定一个到期的待办任务。 -// 2. 从待办队列中将其删除。 -// 3. 将其在执行日志表中的状态更新为 'running'。 -// 最后返回更新后的执行日志对象，作为执行的唯一凭证。 -func (r *pendingTaskRepository) ClaimNextDueTask() (*models.TaskExecutionLog, error) { +// ClaimNextAvailableTask 以原子方式认领下一个可用的任务。 +func (r *pendingTaskRepository) ClaimNextAvailableTask(excludePlanIDs []uint) (*models.TaskExecutionLog, error) { var log models.TaskExecutionLog err := r.db.Transaction(func(tx *gorm.DB) error { var pendingTask models.PendingTask - - // 1. 查找并用 FOR UPDATE 锁定一个到期的待办任务 - if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}). + query := tx.Clauses(clause.Locking{Strength: "UPDATE"}). Where("execute_at <= ?", time.Now()). - Order("execute_at ASC"). - First(&pendingTask).Error; err != nil { - return err // 没找到是正常情况，事务将回滚 + Order("execute_at ASC") + + if len(excludePlanIDs) > 0 { + query = query.Where("plan_execution_log_id NOT IN ?", excludePlanIDs) + } + + if err := query.First(&pendingTask).Error; err != nil { + return err } - // 2. 在同一个事务中，立刻硬删除这个待办任务 if err := tx.Unscoped().Delete(&pendingTask).Error; err != nil { return err } - // 3. 在同一个事务中，更新其在日志表中的孪生兄弟的状态 updates := map[string]interface{}{ "status": models.ExecutionStatusStarted, "started_at": time.Now(), @@ -64,8 +60,7 @@ func (r *pendingTaskRepository) ClaimNextDueTask() (*models.TaskExecutionLog, er return err } - // 4. 获取更新后的完整日志对象，以返回给调用者 - if err := tx.First(&log, pendingTask.TaskExecutionLogID).Error; err != nil { + if err := tx.Preload("Task").First(&log, pendingTask.TaskExecutionLogID).Error; err != nil { return err } @@ -73,10 +68,27 @@ func (r *pendingTaskRepository) ClaimNextDueTask() (*models.TaskExecutionLog, er }) if err != nil { - // 如果错误是 `gorm.ErrRecordNotFound`，它仅表示当前没有到期的任务。 - // 服务层应该优雅地处理这种情况，而不是将其视为需要立即处理的严重错误。 return nil, err } return &log, nil } + +// RequeueTask 安全地将一个已被认领但无法执行的任务放回队列。 +// 它在一个事务中原子地将日志状态恢复为 'waiting'，并重新创建待办任务。 +func (r *pendingTaskRepository) RequeueTask(log *models.TaskExecutionLog) error { + return r.db.Transaction(func(tx *gorm.DB) error { + // 1. 将日志状态恢复为 waiting + if err := tx.Model(log).Update("status", models.ExecutionStatusWaiting).Error; err != nil { + return err + } + + // 2. 重新创建待办任务，立即执行 + newPendingTask := models.PendingTask{ + TaskID: log.TaskID, + TaskExecutionLogID: log.ID, + ExecuteAt: time.Now(), + } + return tx.Create(&newPendingTask).Error + }) +} diff --git a/internal/infra/task/task.go b/internal/infra/task/task.go index d58acea..7719500 100644 --- a/internal/infra/task/task.go +++ b/internal/infra/task/task.go @@ -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) }