NianAIGC/backend/internal/jobs/service.go

package jobs

import (
	"context"
	"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)
	ClaimJobs(context.Context, string, int, int) ([]Job, error)
}

type Patch struct {
	Status              *Status
	Error               *JobError
	Attempts            *int
	ScheduledAt         *time.Time
	CompletedAt         *time.Time
	ProviderTaskID      *string
	ClearProviderTaskID bool
	ClearLease          bool
	WebhookAttempts     *int
	WebhookLastStatus   []byte
	SetWebhookStatus    bool
}

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.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
	job, err = service.store.UpdateJob(ctx, job.ID, Patch{Status: &status, CompletedAt: &now})
	if err != nil {
		return Job{}, err
	}
	if refunds != nil {
		job, err = refunds.Refund(ctx, job, "任务已取消")
		if err != nil {
			return Job{}, err
		}
	}
	return service.store.UpdateJob(ctx, job.ID, Patch{ClearLease: true})
}

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
}