276 lines
7.5 KiB
Go
276 lines
7.5 KiB
Go
package model
|
|
|
|
import (
|
|
"errors"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
|
|
"gorm.io/gorm"
|
|
)
|
|
|
|
type SystemTaskStatus string
|
|
|
|
const (
|
|
SystemTaskStatusPending SystemTaskStatus = "pending"
|
|
SystemTaskStatusRunning SystemTaskStatus = "running"
|
|
SystemTaskStatusSucceeded SystemTaskStatus = "succeeded"
|
|
SystemTaskStatusFailed SystemTaskStatus = "failed"
|
|
|
|
SystemTaskTypeLogCleanup = "log_cleanup"
|
|
)
|
|
|
|
var ErrSystemTaskLockLost = errors.New("system task lock lost")
|
|
|
|
type SystemTask struct {
|
|
ID int64 `json:"id" gorm:"primary_key;AUTO_INCREMENT"`
|
|
TaskID string `json:"task_id" gorm:"type:varchar(64);uniqueIndex"`
|
|
Type string `json:"type" gorm:"type:varchar(64);index"`
|
|
Status SystemTaskStatus `json:"status" gorm:"type:varchar(32);index"`
|
|
ActiveKey *string `json:"active_key,omitempty" gorm:"type:varchar(64);uniqueIndex"`
|
|
Payload string `json:"payload" gorm:"type:text"`
|
|
State string `json:"state" gorm:"type:text"`
|
|
Result string `json:"result" gorm:"type:text"`
|
|
Error string `json:"error" gorm:"type:text"`
|
|
LockedBy string `json:"locked_by" gorm:"type:varchar(128);index"`
|
|
LockedUntil int64 `json:"locked_until" gorm:"bigint;index"`
|
|
CreatedAt int64 `json:"created_at" gorm:"bigint;index"`
|
|
UpdatedAt int64 `json:"updated_at" gorm:"bigint;index"`
|
|
}
|
|
|
|
type SystemTaskResponse struct {
|
|
ID int64 `json:"id"`
|
|
TaskID string `json:"task_id"`
|
|
Type string `json:"type"`
|
|
Status SystemTaskStatus `json:"status"`
|
|
ActiveKey string `json:"active_key,omitempty"`
|
|
Payload any `json:"payload"`
|
|
State any `json:"state"`
|
|
Result any `json:"result"`
|
|
Error string `json:"error"`
|
|
LockedBy string `json:"locked_by"`
|
|
LockedUntil int64 `json:"locked_until"`
|
|
CreatedAt int64 `json:"created_at"`
|
|
UpdatedAt int64 `json:"updated_at"`
|
|
}
|
|
|
|
func (task *SystemTask) BeforeCreate(_ *gorm.DB) error {
|
|
now := common.GetTimestamp()
|
|
if task.CreatedAt == 0 {
|
|
task.CreatedAt = now
|
|
}
|
|
if task.UpdatedAt == 0 {
|
|
task.UpdatedAt = now
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func GenerateSystemTaskID() (string, error) {
|
|
key, err := common.GenerateRandomCharsKey(32)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
return "systask_" + key, nil
|
|
}
|
|
|
|
func CreateSystemTask(taskType string, activeKey string, payload any, state any) (*SystemTask, error) {
|
|
taskID, err := GenerateSystemTaskID()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
payloadText, err := marshalSystemTaskJSON(payload)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
stateText, err := marshalSystemTaskJSON(state)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
task := &SystemTask{
|
|
TaskID: taskID,
|
|
Type: taskType,
|
|
Status: SystemTaskStatusPending,
|
|
Payload: payloadText,
|
|
State: stateText,
|
|
}
|
|
if activeKey != "" {
|
|
task.ActiveKey = &activeKey
|
|
}
|
|
|
|
if err := DB.Create(task).Error; err != nil {
|
|
return nil, err
|
|
}
|
|
return task, nil
|
|
}
|
|
|
|
func GetSystemTaskByTaskID(taskID string) (*SystemTask, error) {
|
|
var task SystemTask
|
|
if err := DB.Where("task_id = ?", taskID).First(&task).Error; err != nil {
|
|
if errors.Is(err, gorm.ErrRecordNotFound) {
|
|
return nil, nil
|
|
}
|
|
return nil, err
|
|
}
|
|
return &task, nil
|
|
}
|
|
|
|
func GetActiveSystemTask(taskType string) (*SystemTask, error) {
|
|
var task SystemTask
|
|
err := DB.Where("type = ? AND active_key IS NOT NULL", taskType).
|
|
Where("status IN ?", activeSystemTaskStatuses()).
|
|
Order("id desc").
|
|
First(&task).Error
|
|
if err != nil {
|
|
if errors.Is(err, gorm.ErrRecordNotFound) {
|
|
return nil, nil
|
|
}
|
|
return nil, err
|
|
}
|
|
return &task, nil
|
|
}
|
|
|
|
func FindRunnableSystemTasks(taskType string, now int64, limit int) ([]*SystemTask, error) {
|
|
var tasks []*SystemTask
|
|
if limit <= 0 {
|
|
limit = 1
|
|
}
|
|
err := DB.Where("type = ? AND status IN ? AND (locked_until = 0 OR locked_until < ?)", taskType, activeSystemTaskStatuses(), now).
|
|
Order("id asc").
|
|
Limit(limit).
|
|
Find(&tasks).Error
|
|
return tasks, err
|
|
}
|
|
|
|
func ClaimSystemTask(id int64, taskType string, runnerID string, lockUntil int64) (*SystemTask, bool, error) {
|
|
now := common.GetTimestamp()
|
|
result := DB.Model(&SystemTask{}).
|
|
Where("id = ? AND type = ? AND status IN ? AND (locked_until = 0 OR locked_until < ? OR locked_by = ?)", id, taskType, activeSystemTaskStatuses(), now, runnerID).
|
|
Updates(map[string]any{
|
|
"status": SystemTaskStatusRunning,
|
|
"locked_by": runnerID,
|
|
"locked_until": lockUntil,
|
|
"updated_at": now,
|
|
})
|
|
if result.Error != nil {
|
|
return nil, false, result.Error
|
|
}
|
|
if result.RowsAffected == 0 {
|
|
return nil, false, nil
|
|
}
|
|
|
|
var task SystemTask
|
|
if err := DB.Where("id = ?", id).First(&task).Error; err != nil {
|
|
return nil, false, err
|
|
}
|
|
return &task, true, nil
|
|
}
|
|
|
|
func UpdateSystemTaskState(taskID string, lockedBy string, state any, lockUntil int64) error {
|
|
stateText, err := marshalSystemTaskJSON(state)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
result := DB.Model(&SystemTask{}).
|
|
Where("task_id = ? AND status = ? AND locked_by = ?", taskID, SystemTaskStatusRunning, lockedBy).
|
|
Updates(map[string]any{
|
|
"state": stateText,
|
|
"locked_until": lockUntil,
|
|
"updated_at": common.GetTimestamp(),
|
|
})
|
|
if result.Error != nil {
|
|
return result.Error
|
|
}
|
|
if result.RowsAffected == 0 {
|
|
return ErrSystemTaskLockLost
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func FinishSystemTask(taskID string, lockedBy string, status SystemTaskStatus, resultPayload any, errorMessage string) error {
|
|
resultText, err := marshalSystemTaskJSON(resultPayload)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
result := DB.Model(&SystemTask{}).
|
|
Where("task_id = ? AND status = ? AND locked_by = ?", taskID, SystemTaskStatusRunning, lockedBy).
|
|
Updates(map[string]any{
|
|
"status": status,
|
|
"active_key": nil,
|
|
"result": resultText,
|
|
"error": errorMessage,
|
|
"locked_by": "",
|
|
"locked_until": 0,
|
|
"updated_at": common.GetTimestamp(),
|
|
})
|
|
if result.Error != nil {
|
|
return result.Error
|
|
}
|
|
if result.RowsAffected == 0 {
|
|
return ErrSystemTaskLockLost
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (task *SystemTask) DecodePayload(v any) error {
|
|
return decodeSystemTaskJSONString(task.Payload, v)
|
|
}
|
|
|
|
func (task *SystemTask) DecodeState(v any) error {
|
|
return decodeSystemTaskJSONString(task.State, v)
|
|
}
|
|
|
|
func (task *SystemTask) ToResponse() SystemTaskResponse {
|
|
activeKey := ""
|
|
if task.ActiveKey != nil {
|
|
activeKey = *task.ActiveKey
|
|
}
|
|
return SystemTaskResponse{
|
|
ID: task.ID,
|
|
TaskID: task.TaskID,
|
|
Type: task.Type,
|
|
Status: task.Status,
|
|
ActiveKey: activeKey,
|
|
Payload: decodeSystemTaskJSONValue(task.Payload),
|
|
State: decodeSystemTaskJSONValue(task.State),
|
|
Result: decodeSystemTaskJSONValue(task.Result),
|
|
Error: task.Error,
|
|
LockedBy: task.LockedBy,
|
|
LockedUntil: task.LockedUntil,
|
|
CreatedAt: task.CreatedAt,
|
|
UpdatedAt: task.UpdatedAt,
|
|
}
|
|
}
|
|
|
|
func activeSystemTaskStatuses() []string {
|
|
return []string{string(SystemTaskStatusPending), string(SystemTaskStatusRunning)}
|
|
}
|
|
|
|
func marshalSystemTaskJSON(v any) (string, error) {
|
|
if v == nil {
|
|
return "", nil
|
|
}
|
|
data, err := common.Marshal(v)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
return string(data), nil
|
|
}
|
|
|
|
func decodeSystemTaskJSONString(data string, v any) error {
|
|
if data == "" {
|
|
return nil
|
|
}
|
|
return common.UnmarshalJsonStr(data, v)
|
|
}
|
|
|
|
func decodeSystemTaskJSONValue(data string) any {
|
|
if data == "" {
|
|
return nil
|
|
}
|
|
var value any
|
|
if err := common.UnmarshalJsonStr(data, &value); err != nil {
|
|
return data
|
|
}
|
|
return value
|
|
}
|