NianAIGC/backend/internal/jobs/service.go

package jobs

import (
	"context"
	"errors"
	"fmt"
	"time"
)

type Store interface {
	ListJobs(context.Context, ListFilter) ([]Job, error)
	FindJob(context.Context, Scope, string) (Job, bool, error)
	FindIdempotentJob(context.Context, Scope, string) (Job, bool, error)
	CreateJob(context.Context, Job) (Job, error)
	UpdateJob(context.Context, string, Patch) (Job, error)
	DeleteJob(context.Context, string) error
	ClaimJobs(context.Context, string, int, int) ([]Job, error)
}

type ArtifactDeleter interface {
	DeleteOutputs(context.Context, Job) ([]string, error)
}

type Patch struct {
	Status                    *Status
	Error                     *JobError
	ClearError                bool
	Attempts                  *int
	ScheduledAt               *time.Time
	CompletedAt               *time.Time
	ProviderTaskID            *string
	ProviderDispatchStartedAt *time.Time
	ResponsePayload           []byte
	SetResponsePayload        bool
	ClearProviderTaskID       bool
	ClearProviderDispatch     bool
	ClearLease                bool
	WebhookAttempts           *int
	WebhookLastStatus         []byte
	SetWebhookStatus          bool
	FinalizedAt               *time.Time
	DispatchReadyAt           *time.Time
	ExpectedStatuses          []Status
	ExpectedLockedBy          *string
}

type Service struct {
	store Store
	now   func() time.Time
}

func NewService(store Store, now func() time.Time) *Service {
	if now == nil {
		now = time.Now
	}
	return &Service{store: store, now: now}
}

func (service *Service) List(ctx context.Context, filter ListFilter) ([]Job, error) {
	if filter.Scope.OwnerID == "" {
		return nil, &Error{Kind: ErrorInvalid, Status: 400, Message: "owner is required"}
	}
	if filter.Status != "" && !filter.Status.Valid() {
		return nil, &Error{Kind: ErrorInvalid, Status: 400, Message: fmt.Sprintf("Unsupported status filter: %s", filter.Status)}
	}
	if filter.Capability != "" && filter.Capability != "image.generate" && filter.Capability != "video.generate" {
		return nil, &Error{Kind: ErrorInvalid, Status: 400, Message: fmt.Sprintf("Unsupported capability filter: %s", filter.Capability)}
	}
	filter.Limit = NormalizePublicLimit(filter.Limit)
	return service.store.ListJobs(ctx, filter)
}

func (service *Service) Get(ctx context.Context, scope Scope, id string) (Job, error) {
	job, found, err := service.store.FindJob(ctx, scope, id)
	if err != nil {
		return Job{}, err
	}
	if !found {
		return Job{}, &Error{Kind: ErrorNotFound, Status: 404, Message: "Job not found."}
	}
	return job, nil
}

func (service *Service) Create(ctx context.Context, command CreateCommand) (Job, bool, error) {
	job := command.Job
	if job.OwnerID == "" || job.ID == "" || !job.Status.Valid() || job.Capability == "" || job.Provider == "" || job.ReqKey == "" {
		return Job{}, false, &Error{Kind: ErrorInvalid, Status: 400, Message: "invalid generation job"}
	}
	job.Priority = NormalizePriority(job.Priority)
	if job.MaxAttempts <= 0 {
		job.MaxAttempts = 3
	}
	now := service.now().UTC()
	if job.CreatedAt.IsZero() {
		job.CreatedAt = now
	}
	if job.UpdatedAt.IsZero() {
		job.UpdatedAt = now
	}
	if job.ScheduledAt.IsZero() {
		job.ScheduledAt = now
	}
	if job.DispatchReadyAt == nil && !command.HoldDispatch {
		ready := now
		job.DispatchReadyAt = &ready
	}
	if job.InputAssetIDs == nil {
		job.InputAssetIDs = []string{}
	}
	if job.InputURLs == nil {
		job.InputURLs = []string{}
	}
	if job.OutputAssetIDs == nil {
		job.OutputAssetIDs = []string{}
	}
	if len(job.RequestPayload) == 0 {
		job.RequestPayload = []byte(`{}`)
	}

	scope := Scope{OwnerID: job.OwnerID, ExternalClientID: job.ExternalClientID}
	if job.IdempotencyKey != "" && job.ExternalClientID != "" {
		fingerprint, err := Fingerprint(command.IdempotencyBody)
		if err != nil {
			return Job{}, false, err
		}
		job.IdempotencyFingerprint = fingerprint
		if existing, found, err := service.store.FindIdempotentJob(ctx, scope, job.IdempotencyKey); err != nil {
			return Job{}, false, err
		} else if found {
			return compareIdempotent(existing, fingerprint)
		}
	}

	created, err := service.store.CreateJob(ctx, job)
	if err == ErrUniqueIdempotency && job.IdempotencyKey != "" {
		existing, found, lookupErr := service.store.FindIdempotentJob(ctx, scope, job.IdempotencyKey)
		if lookupErr != nil {
			return Job{}, false, lookupErr
		}
		if found {
			return compareIdempotent(existing, job.IdempotencyFingerprint)
		}
	}
	return created, false, err
}

func compareIdempotent(existing Job, fingerprint string) (Job, bool, error) {
	if existing.IdempotencyFingerprint != fingerprint {
		return Job{}, false, &Error{Kind: ErrorConflict, Status: 409, Message: "Idempotency key was already used with a different request body."}
	}
	return existing, true, nil
}

type RefundPort interface {
	Refund(context.Context, Job, string) (Job, error)
}

func (service *Service) Cancel(ctx context.Context, scope Scope, id string, refunds RefundPort) (Job, error) {
	job, err := service.Get(ctx, scope, id)
	if err != nil || job.Status.Terminal() {
		return job, err
	}
	now := service.now().UTC()
	status := StatusCancelled
	patch := Patch{Status: &status, CompletedAt: &now, ClearLease: true, ExpectedStatuses: []Status{job.Status}}
	if job.LockedBy != "" {
		owner := job.LockedBy
		patch.ExpectedLockedBy = &owner
	}
	job, err = service.store.UpdateJob(ctx, job.ID, patch)
	if errors.Is(err, ErrTransitionConflict) {
		current, getErr := service.Get(ctx, scope, id)
		if getErr == nil && current.Status.Terminal() {
			return current, nil
		}
	}
	if err != nil {
		return Job{}, err
	}
	if refunds != nil {
		job, err = refunds.Refund(ctx, job, "任务已取消")
		if err != nil {
			return Job{}, err
		}
	}
	return job, nil
}

func (service *Service) Retry(ctx context.Context, scope Scope, id, newID string) (Job, error) {
	original, err := service.Get(ctx, scope, id)
	if err != nil {
		return Job{}, err
	}
	if original.Capability != "image.generate" || (original.Status != StatusFailed && original.Status != StatusExpired && original.Status != StatusCancelled) {
		return Job{}, &Error{Kind: ErrorInvalid, Status: 400, Message: "任务当前不可重试"}
	}
	now := service.now().UTC()
	retry := original
	retry.ID, retry.RetryOf, retry.Status = newID, original.ID, StatusQueued
	retry.ProviderTaskID, retry.Error, retry.IdempotencyKey, retry.IdempotencyFingerprint = "", nil, "", ""
	retry.OutputAssetIDs, retry.Attempts, retry.LockedAt, retry.LockedBy = []string{}, 0, nil, ""
	retry.StartedAt, retry.CompletedAt = nil, nil
	retry.ScheduledAt, retry.CreatedAt, retry.UpdatedAt = now, now, now
	return service.store.CreateJob(ctx, retry)
}

func (service *Service) Delete(ctx context.Context, scope Scope, id string, artifacts ArtifactDeleter) (Job, error) {
	job, err := service.Get(ctx, scope, id)
	if err != nil {
		return Job{}, err
	}
	if artifacts != nil {
		if _, err := artifacts.DeleteOutputs(ctx, job); err != nil {
			return Job{}, err
		}
	}
	if err := service.store.DeleteJob(ctx, job.ID); err != nil {
		return Job{}, err
	}
	return job, nil
}

func RetryDelay(attempts int, base, maximum time.Duration) time.Duration {
	if base <= 0 {
		base = 10 * time.Second
	}
	if maximum <= 0 {
		maximum = 5 * time.Minute
	}
	delay := base
	for index := 1; index < attempts && delay < maximum; index++ {
		delay *= 2
		if delay > maximum {
			return maximum
		}
	}
	return delay
}