实现task调度器同一时间只能运行同一个plan的同一个task

TODO-List

@@ -1,6 +1,7 @@
 // TODO 列表
 
 // TODO 可以实现的问题
+plan执行到一半时如果用户删掉里面的task, 后续调度器执行task时可能会找不到这个任务的细节
 1. 可以用TimescaleDB代替PGSQL, 优化传感器数据存储性能
 2. 系统启动时应该检查一遍执行历史库, 将所有显示为执行中的任务都修正为执行失败并报错
 

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
+	})
+}

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
+	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 方法会阻塞，直到有协程空闲出来，这自然地实现了背压。
+	// 尝试获取内存执行锁
+	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)
-		// 可以在这里添加逻辑，将任务状态恢复为 pending
+			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)
 }