package repository import ( "time" "git.huangwc.com/pig/pig-farm-controller/internal/infra/models" "gorm.io/gorm" "gorm.io/gorm/clause" ) // PendingTaskRepository 定义了与待执行任务队列交互的接口。 type PendingTaskRepository interface { CreatePendingTasksInBatch(tasks []*models.PendingTask) error ClaimNextAvailableTask(excludePlanIDs []uint) (*models.TaskExecutionLog, error) RequeueTask(log *models.TaskExecutionLog) error } // pendingTaskRepository 是使用 GORM 的具体实现。 type pendingTaskRepository struct { db *gorm.DB } // NewPendingTaskRepository 创建一个新的待执行任务队列仓库。 func NewPendingTaskRepository(db *gorm.DB) PendingTaskRepository { return &pendingTaskRepository{db: db} } // CreatePendingTasksInBatch 在一次数据库调用中创建多个待执行任务条目。 func (r *pendingTaskRepository) CreatePendingTasksInBatch(tasks []*models.PendingTask) error { return r.db.Create(&tasks).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 query := tx.Clauses(clause.Locking{Strength: "UPDATE"}). Where("execute_at <= ?", time.Now()). 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 } if err := tx.Unscoped().Delete(&pendingTask).Error; err != nil { return err } updates := map[string]interface{}{ "status": models.ExecutionStatusStarted, "started_at": time.Now(), } if err := tx.Model(&models.TaskExecutionLog{}).Where("id = ?", pendingTask.TaskExecutionLogID).Updates(updates).Error; err != nil { return err } if err := tx.Preload("Task").First(&log, pendingTask.TaskExecutionLogID).Error; err != nil { return err } return nil }) if err != nil { 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 }) }