95 lines
2.8 KiB
Go
95 lines
2.8 KiB
Go
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
|
|
})
|
|
}
|