feat: add persistent system task log cleanup progress

2026-06-22 19:16:53 +08:00
parent 6dc4030fdf
commit a162163b48
19 changed files with 984 additions and 51 deletions
@@ -0,0 +1,275 @@
+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
+}