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 }