package worker import ( "context" "encoding/json" "errors" "fmt" "log/slog" "math/rand" "net" "net/http" "strconv" "strings" "time" "juhe-factory/api/internal/billing" mediakey "juhe-factory/api/internal/media" "juhe-factory/api/internal/model" "juhe-factory/api/internal/provider" "juhe-factory/api/internal/provider/apimart" queuepkg "juhe-factory/api/internal/queue" "juhe-factory/api/internal/security" "juhe-factory/api/internal/storage" "github.com/google/uuid" "github.com/hibiken/asynq" "gorm.io/gorm" "gorm.io/gorm/clause" ) var completedStatuses = map[string]bool{"completed": true, "complete": true, "succeeded": true, "success": true, "done": true} var failedStatuses = map[string]bool{"failed": true, "failure": true, "error": true, "cancelled": true, "canceled": true} const maxSubmitAttempts = 5 type Generation struct { DB *gorm.DB Queue *asynq.Client COS *storage.COS Encryptor *security.Encryptor Provider apimart.Client Seedance2 apimart.Seedance2 DownloadHTTP *http.Client PollInterval time.Duration WorkerID string } type channelSnapshot struct { ID uuid.UUID BaseURL string APIKeyCiphertext string MaxConcurrency int MaxUserConcurrency int Enabled bool } type submitSnapshot struct { Task model.GenerationTask Channel channelSnapshot Model string } func NewGeneration(db *gorm.DB, queue *asynq.Client, cos *storage.COS, encryptor *security.Encryptor, client *http.Client, pollInterval time.Duration) *Generation { return &Generation{DB: db, Queue: queue, COS: cos, Encryptor: encryptor, Provider: apimart.NewClient(client), Seedance2: apimart.NewSeedance2(client), DownloadHTTP: client, PollInterval: pollInterval, WorkerID: uuid.NewString()} } func (w *Generation) Register(mux *asynq.ServeMux) { mux.HandleFunc(queuepkg.TypeDispatchChannel, w.dispatchChannel) mux.HandleFunc(queuepkg.TypeSubmitTask, w.submitTask) mux.HandleFunc(queuepkg.TypePollTask, w.pollTask) mux.HandleFunc(queuepkg.TypeDownloadTask, w.downloadTask) } func decodeID(task *asynq.Task) (uuid.UUID, error) { var payload queuepkg.IDPayload if err := json.Unmarshal(task.Payload(), &payload); err != nil { return uuid.Nil, err } if payload.ID == uuid.Nil { return uuid.Nil, errors.New("任务标识无效") } return payload.ID, nil } func (w *Generation) dispatchChannel(ctx context.Context, raw *asynq.Task) error { channelID, err := decodeID(raw) if err != nil { return err } selected := make([]uuid.UUID, 0) err = w.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error { var channel channelSnapshot if err := tx.Raw(`SELECT id,base_url,api_key_ciphertext,max_concurrency,max_user_concurrency,enabled FROM channels WHERE id=? AND deleted_at IS NULL FOR UPDATE`, channelID).Scan(&channel).Error; err != nil { return err } if channel.ID == uuid.Nil || !channel.Enabled { return nil } var active int64 if err := tx.Model(&model.GenerationTask{}).Where("channel_id=? AND slot_reserved=true AND slot_released_at IS NULL", channelID).Count(&active).Error; err != nil { return err } available := channel.MaxConcurrency - int(active) if available <= 0 { return nil } activeByUser := map[uuid.UUID]int{} rows := make([]struct { UserID uuid.UUID Count int }, 0) if err := tx.Model(&model.GenerationTask{}).Select("user_id,count(*) AS count").Where("channel_id=? AND slot_reserved=true AND slot_released_at IS NULL", channelID).Group("user_id").Scan(&rows).Error; err != nil { return err } for _, row := range rows { activeByUser[row.UserID] = row.Count } var candidates []model.GenerationTask now := time.Now() if err := tx.Raw(`WITH active_users AS ( SELECT user_id,count(*) AS active_count FROM generation_tasks WHERE channel_id=? AND slot_reserved=true AND slot_released_at IS NULL GROUP BY user_id ), ranked AS ( SELECT pending.id,row_number() OVER (PARTITION BY pending.user_id ORDER BY pending.created_at,pending.id) AS user_rank, (?-coalesce(active.active_count,0)) AS allowance FROM generation_tasks pending LEFT JOIN active_users active ON active.user_id=pending.user_id WHERE pending.channel_id=? AND pending.status='pending_submission' AND (pending.project_id IS NOT NULL OR (pending.task_type='image_generation' AND pending.input_data->>'product_image'='true')) AND (pending.next_submit_at IS NULL OR pending.next_submit_at<=?) AND (pending.lease_until IS NULL OR pending.lease_until0 AND ranked.user_rank<=ranked.allowance ORDER BY ranked.user_rank,task.created_at,task.id LIMIT ? FOR UPDATE OF task SKIP LOCKED`, channelID, channel.MaxUserConcurrency, channelID, now, now, available).Scan(&candidates).Error; err != nil { return err } if len(candidates) == 0 { return nil } for _, candidate := range selectFairCandidates(candidates, activeByUser, available, channel.MaxUserConcurrency) { leaseUntil := now.Add(2 * time.Minute) result := tx.Model(&model.GenerationTask{}).Where("id=? AND status='pending_submission'", candidate.ID).Updates(map[string]any{ "status": "submitting", "slot_reserved": true, "slot_reserved_at": now, "slot_released_at": nil, "lease_owner": w.WorkerID, "lease_until": leaseUntil, "submit_attempts": gorm.Expr("submit_attempts+1"), }) if result.Error != nil { return result.Error } if result.RowsAffected == 1 { selected = append(selected, candidate.ID) } } return nil }) if err != nil { return err } for _, taskID := range selected { if err := queuepkg.EnqueueID(w.Queue, queuepkg.TypeSubmitTask, taskID, 0); err != nil { slog.Error("提交任务入队失败", "task_id", taskID, "error", err) w.returnUnenqueuedTask(ctx, taskID, err.Error()) } } return nil } func selectFairCandidates(candidates []model.GenerationTask, activeByUser map[uuid.UUID]int, available, maxPerUser int) []model.GenerationTask { if available <= 0 || maxPerUser <= 0 { return nil } groups := map[uuid.UUID][]model.GenerationTask{} order := make([]uuid.UUID, 0) for _, candidate := range candidates { if _, exists := groups[candidate.UserID]; !exists { order = append(order, candidate.UserID) } groups[candidate.UserID] = append(groups[candidate.UserID], candidate) } selected := make([]model.GenerationTask, 0, available) for available > 0 { progressed := false for _, userID := range order { if available == 0 { break } if activeByUser[userID] >= maxPerUser || len(groups[userID]) == 0 { continue } selected = append(selected, groups[userID][0]) groups[userID] = groups[userID][1:] activeByUser[userID]++ available-- progressed = true } if !progressed { break } } return selected } func (w *Generation) loadSubmitSnapshot(ctx context.Context, taskID uuid.UUID) (submitSnapshot, error) { var result submitSnapshot if err := w.DB.WithContext(ctx).Where("id=?", taskID).First(&result.Task).Error; err != nil { return result, err } if result.Task.ChannelID == nil || result.Task.ModelID == nil { return result, errors.New("任务未绑定模型渠道") } if err := w.DB.WithContext(ctx).Raw(`SELECT id,base_url,api_key_ciphertext,max_concurrency,max_user_concurrency,enabled FROM channels WHERE id=? AND deleted_at IS NULL`, *result.Task.ChannelID).Scan(&result.Channel).Error; err != nil { return result, err } if !result.Channel.Enabled { return result, errors.New("任务渠道已停用") } if err := w.DB.WithContext(ctx).Table("models").Select("name").Where("id=? AND enabled=true AND deleted_at IS NULL", *result.Task.ModelID).Scan(&result.Model).Error; err != nil { return result, err } if result.Model == "" { return result, errors.New("任务模型不可用") } return result, nil } func (w *Generation) submitTask(ctx context.Context, raw *asynq.Task) error { taskID, err := decodeID(raw) if err != nil { return err } snapshot, err := w.loadSubmitSnapshot(ctx, taskID) if err != nil { // 用户可以删除卡死的活动任务;队列收到遗留消息时直接结束,避免无意义重试。 if errors.Is(err, gorm.ErrRecordNotFound) { return nil } return w.failAndRefund(ctx, taskID, "submit_configuration_error", err.Error()) } if snapshot.Task.Status != "submitting" { if snapshot.Task.Status == "pending_submission" && snapshot.Task.ChannelID != nil { return queuepkg.EnqueueID(w.Queue, queuepkg.TypeDispatchChannel, *snapshot.Task.ChannelID, 0) } return nil } apiKey, err := w.Encryptor.Decrypt(snapshot.Channel.APIKeyCiphertext) if err != nil { return w.failAndRefund(ctx, taskID, "channel_secret_error", "渠道密钥不可用") } payload, err := buildProviderPayload(snapshot) if err != nil { return w.failAndRefund(ctx, taskID, "invalid_payload", err.Error()) } requestCtx, cancel := context.WithTimeout(ctx, 90*time.Second) defer cancel() var result provider.SubmitResult if snapshot.Task.TaskType == "video_generation" { if payloadJSON, marshalErr := json.Marshal(payload); marshalErr == nil { fmt.Printf("console.log video_generation payload: %s\n", payloadJSON) } result, err = w.Seedance2.Submit(requestCtx, snapshot.Channel.BaseURL, apiKey, snapshot.Task.RequestID, payload) } else { result, err = w.Provider.SubmitImage(requestCtx, snapshot.Channel.BaseURL, apiKey, snapshot.Task.RequestID, payload) } if err != nil { var httpErr *provider.HTTPError if errors.As(err, &httpErr) { if httpErr.StatusCode == http.StatusTooManyRequests || httpErr.StatusCode >= 500 { if !canRetrySubmission(snapshot.Task.SubmitAttempts) { return w.failAndRefund(ctx, taskID, "upstream_unavailable", err.Error()) } return w.requeueRejected(ctx, snapshot.Task, err.Error(), retryDelay(snapshot.Task.SubmitAttempts, httpErr.RetryAfter)) } return w.failAndRefund(ctx, taskID, "upstream_rejected", err.Error()) } // 传输异常统一按提交失败重试;达到上限后结束任务并返还预扣积分。 if !canRetrySubmission(snapshot.Task.SubmitAttempts) { return w.failAndRefund(ctx, taskID, "upstream_unavailable", err.Error()) } return w.requeueRejected(ctx, snapshot.Task, err.Error(), retryDelay(snapshot.Task.SubmitAttempts, "")) } now := time.Now() if result.URL != "" { if err := w.markResultReady(ctx, taskID, result.URL); err != nil { return err } _ = queuepkg.EnqueueID(w.Queue, queuepkg.TypeDownloadTask, taskID, 0) w.enqueueDispatch(snapshot.Task.ChannelID) return nil } nextPoll := now.Add(w.PollInterval + jitter(3*time.Second)) update := w.DB.WithContext(ctx).Model(&model.GenerationTask{}).Where("id=? AND status='submitting'", taskID).Updates(map[string]any{ "status": "submitted", "upstream_task_id": result.TaskID, "submitted_at": now, "next_poll_at": nextPoll, "lease_owner": nil, "lease_until": nil, "error_code": nil, "error_message": nil, }) if update.Error != nil { return update.Error } if update.RowsAffected == 1 { _ = queuepkg.EnqueueID(w.Queue, queuepkg.TypePollTask, taskID, time.Until(nextPoll)) } return nil } func buildProviderPayload(snapshot submitSnapshot) (map[string]any, error) { if snapshot.Task.TaskType == "video_generation" { return apimart.BuildSeedance2Payload(snapshot.Model, snapshot.Task.InputData) } var input map[string]any if err := json.Unmarshal(snapshot.Task.InputData, &input); err != nil { return nil, err } payload := map[string]any{"model": snapshot.Model, "prompt": strings.TrimSpace(fmt.Sprint(input["prompt"]))} if payload["prompt"] == "" { return nil, errors.New("生成提示词不能为空") } if snapshot.Task.TaskType == "image_generation" { // 图片画幅由各业务模块在入队前规范化,Worker 只负责透传统一字段。 payload["size"] = stringOr(input["size"], "16:9") payload["resolution"] = stringOr(input["resolution"], "1k") } else { payload["size"] = stringOr(input["resolution"], "1k") } if n, ok := input["n"]; ok { payload["n"] = n } else { payload["n"] = 1 } if urls, ok := input["image_urls"]; ok { payload["image_urls"] = urls } return payload, nil } func (w *Generation) pollTask(ctx context.Context, raw *asynq.Task) error { taskID, err := decodeID(raw) if err != nil { return err } snapshot, err := w.loadSubmitSnapshot(ctx, taskID) if err != nil || (snapshot.Task.Status != "submitted" && snapshot.Task.Status != "processing" && snapshot.Task.Status != "cancel_requested") { return nil } apiKey, err := w.Encryptor.Decrypt(snapshot.Channel.APIKeyCiphertext) if err != nil { return err } requestCtx, cancel := context.WithTimeout(ctx, 30*time.Second) var result provider.PollResult var pollErr error if snapshot.Task.TaskType == "video_generation" { result, pollErr = w.Seedance2.Poll(requestCtx, snapshot.Channel.BaseURL, apiKey, snapshot.Task.UpstreamTaskID) } else { result, pollErr = w.Provider.PollImage(requestCtx, snapshot.Channel.BaseURL, apiKey, snapshot.Task.UpstreamTaskID) } cancel() if pollErr != nil { if snapshot.Task.Status == "cancel_requested" && snapshot.Task.PollAttempts >= 10 { if err := w.finishCancellation(ctx, taskID, true, "上游任务状态长时间不可用,已自动结束取消"); err != nil { return err } w.enqueueDispatch(snapshot.Task.ChannelID) return nil } next := time.Now().Add(w.PollInterval + jitter(4*time.Second)) w.DB.WithContext(ctx).Model(&model.GenerationTask{}).Where("id=? AND status IN ('submitted','processing','cancel_requested')", taskID).Updates(map[string]any{ "poll_attempts": gorm.Expr("poll_attempts+1"), "next_poll_at": next, "error_message": pollErr.Error(), }) _ = queuepkg.EnqueueID(w.Queue, queuepkg.TypePollTask, taskID, time.Until(next)) return nil } if completedStatuses[result.Status] { if snapshot.Task.Status == "cancel_requested" { if err := w.finishCancellation(ctx, taskID, false, "上游任务已结束"); err != nil { return err } w.enqueueDispatch(snapshot.Task.ChannelID) return nil } if result.URL == "" { next := time.Now().Add(w.PollInterval) w.DB.WithContext(ctx).Model(&model.GenerationTask{}).Where("id=?", taskID).Updates(map[string]any{"next_poll_at": next, "error_message": "中转站已完成但结果地址尚未就绪"}) _ = queuepkg.EnqueueID(w.Queue, queuepkg.TypePollTask, taskID, time.Until(next)) return nil } if err := w.markResultReady(ctx, taskID, result.URL); err != nil { return err } _ = queuepkg.EnqueueID(w.Queue, queuepkg.TypeDownloadTask, taskID, 0) w.enqueueDispatch(snapshot.Task.ChannelID) return nil } if failedStatuses[result.Status] { if snapshot.Task.Status == "cancel_requested" { if err := w.finishCancellation(ctx, taskID, true, "上游任务已取消"); err != nil { return err } } else { if err := w.failAndRefund(ctx, taskID, "upstream_failed", result.Error); err != nil { return err } } w.enqueueDispatch(snapshot.Task.ChannelID) return nil } next := time.Now().Add(w.PollInterval + jitter(4*time.Second)) updates := map[string]any{"poll_attempts": gorm.Expr("poll_attempts+1"), "next_poll_at": next, "error_code": nil, "error_message": nil} if snapshot.Task.Status != "cancel_requested" { updates["status"] = "processing" } w.DB.WithContext(ctx).Model(&model.GenerationTask{}).Where("id=? AND status IN ('submitted','processing','cancel_requested')", taskID).Updates(updates) _ = queuepkg.EnqueueID(w.Queue, queuepkg.TypePollTask, taskID, time.Until(next)) return nil } func (w *Generation) markResultReady(ctx context.Context, taskID uuid.UUID, resultURL string) error { return w.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error { var task model.GenerationTask if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("id=?", taskID).First(&task).Error; err != nil { return err } if task.Status == "cancelled" || task.Status == "cancel_requested" { return nil } updates := map[string]any{"status": "result_ready", "upstream_result_url": resultURL, "next_poll_at": nil, "lease_owner": nil, "lease_until": nil} if task.SlotReserved && task.SlotReleasedAt == nil { now := time.Now() updates["slot_released_at"] = now } return tx.Model(&task).Updates(updates).Error }) } func (w *Generation) downloadTask(ctx context.Context, raw *asynq.Task) error { taskID, err := decodeID(raw) if err != nil { return err } var row struct { model.GenerationTask SequenceNo int EpisodeNo int } if err := w.DB.WithContext(ctx).Table("generation_tasks gt").Select("gt.*,coalesce(sb.sequence_no,0) AS sequence_no,coalesce(e.episode_no,0) AS episode_no").Joins("LEFT JOIN episode_storyboards sb ON sb.id=gt.storyboard_id").Joins("LEFT JOIN project_episodes e ON e.id=gt.episode_id").Where("gt.id=?", taskID).Scan(&row).Error; err != nil { return err } if row.Status != "result_ready" && row.Status != "downloading" { return nil } if row.UpstreamResultURL == "" || w.COS == nil { return w.failWithoutRefund(ctx, taskID, "result_storage_error", "结果存储信息不完整") } w.DB.WithContext(ctx).Model(&model.GenerationTask{}).Where("id=? AND status='result_ready'", taskID).Update("status", "downloading") mediaID := uuid.New() outputType, fallbackType, maxBytes := "video", "video/mp4", w.COS.MaxVideoBytes() displayName := fmt.Sprintf("分镜%04d生成结果", row.SequenceNo) key := "" if row.TaskType == "image_generation" { outputType, fallbackType, maxBytes = "image", "image/png", w.COS.MaxImageBytes() var input map[string]any _ = json.Unmarshal(row.InputData, &input) if row.ProjectID == nil { key = mediakey.StandaloneTaskOutput(row.UserID, taskID, "image", fallbackType) displayName = "独立生成图" } else { if fmt.Sprint(input["target_type"]) == "storyboard" && row.StoryboardID != nil { if row.EpisodeID != nil { key = mediakey.TaskOutput(*row.ProjectID, *row.EpisodeID, row.SequenceNo, taskID, mediaID, "image", fallbackType) } else { key = mediakey.ProjectTaskOutput(*row.ProjectID, row.SequenceNo, taskID, mediaID, "image", fallbackType) } displayName = fmt.Sprintf("分镜%04d生成图", row.SequenceNo) } else { assetID, parseErr := uuid.Parse(fmt.Sprint(input["asset_id"])) if parseErr != nil { return w.failWithoutRefund(ctx, taskID, "result_storage_error", "资产生图任务缺少资产标识") } var asset struct{ AssetType, Name string } if err := w.DB.WithContext(ctx).Table("project_assets").Select("asset_type,name").Where("id=? AND project_id=? AND deleted_at IS NULL", assetID, *row.ProjectID).Take(&asset).Error; err != nil { return w.failWithoutRefund(ctx, taskID, "result_storage_error", "资产生图任务绑定的资产不存在") } key = mediakey.AssetOutput(*row.ProjectID, asset.AssetType, assetID, taskID, mediaID, fallbackType) displayName = asset.Name + "生成图" } } } else { if row.ProjectID != nil { if row.EpisodeID != nil { key = mediakey.TaskOutput(*row.ProjectID, *row.EpisodeID, row.SequenceNo, taskID, mediaID, outputType, fallbackType) } else { key = mediakey.ProjectTaskOutput(*row.ProjectID, row.SequenceNo, taskID, mediaID, outputType, fallbackType) } } else { key = mediakey.StandaloneTaskOutput(row.UserID, taskID, "video", fallbackType) } } requestCtx, cancel := context.WithTimeout(ctx, 5*time.Minute) url, storedKey, contentType, size, downloadErr := w.COS.PutFromURL(requestCtx, w.DownloadHTTP, row.UpstreamResultURL, key, fallbackType, maxBytes) cancel() if downloadErr != nil { attempts := row.DownloadAttempts + 1 if attempts < 3 { w.DB.WithContext(ctx).Model(&model.GenerationTask{}).Where("id=?", taskID).Updates(map[string]any{"status": "result_ready", "download_attempts": attempts, "error_message": downloadErr.Error()}) _ = queuepkg.EnqueueID(w.Queue, queuepkg.TypeDownloadTask, taskID, time.Duration(attempts*10)*time.Second) return nil } return w.failWithoutRefund(ctx, taskID, "result_download_failed", downloadErr.Error()) } key = storedKey if (outputType == "image" && !strings.HasPrefix(contentType, "image/")) || (outputType == "video" && !strings.HasPrefix(contentType, "video/")) { _ = w.COS.Delete(ctx, key) return w.failWithoutRefund(ctx, taskID, "invalid_result_media", "中转站结果媒体类型与任务不匹配") } ownerID := row.UserID mediaAsset := model.MediaAsset{ID: mediaID, OwnerUserID: &ownerID, StorageProvider: "cos", ObjectKey: key, PublicURL: url, DisplayName: displayName, MimeType: contentType, SizeBytes: size} metadata, _ := json.Marshal(map[string]any{"source": "apimart", "stored_at": time.Now()}) output := model.GenerationOutput{TaskID: taskID, MediaAssetID: mediaID, OutputType: outputType, SequenceNo: 1, Metadata: metadata} reusedOutput := false cancelledDuringDownload := false err = w.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error { var lockedTask model.GenerationTask if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Select("id", "status").Where("id=?", taskID).Take(&lockedTask).Error; err != nil { return err } if lockedTask.Status == "cancelled" || lockedTask.Status == "cancel_requested" { cancelledDuringDownload = true return nil } var existing model.GenerationOutput existingErr := tx.Where("task_id=? AND sequence_no=?", taskID, 1).Take(&existing).Error if existingErr == nil { output = existing reusedOutput = true } else if errors.Is(existingErr, gorm.ErrRecordNotFound) { if err := tx.Create(&mediaAsset).Error; err != nil { return err } if err := tx.Create(&output).Error; err != nil { return err } } else { return existingErr } now := time.Now() cost := row.PrepaidPoints if err := tx.Model(&model.GenerationTask{}).Where("id=?", taskID).Updates(map[string]any{"status": "succeeded", "actual_points": cost, "finished_at": now, "error_code": nil, "error_message": nil}).Error; err != nil { return err } if row.TaskType == "video_generation" && row.StoryboardID != nil { premiumDrama, premiumErr := isPremiumDramaGeneration(tx, row.GenerationTask) if premiumErr != nil { return premiumErr } if premiumDrama { if err := tx.Exec(`UPDATE generation_outputs SET metadata=jsonb_set(coalesce(metadata,'{}'::jsonb),'{candidate}','true'::jsonb,true) WHERE id=(SELECT active_output_id FROM episode_storyboards WHERE id=?)`, *row.StoryboardID).Error; err != nil { return err } } return tx.Model(&model.EpisodeStoryboard{}).Where("id=?", *row.StoryboardID).Updates(map[string]any{"active_output_id": output.ID, "status": "completed"}).Error } if row.TaskType == "image_generation" && row.ProjectID != nil { var input map[string]any _ = json.Unmarshal(row.InputData, &input) if fmt.Sprint(input["target_type"]) == "storyboard" && row.StoryboardID != nil { return tx.Model(&model.EpisodeStoryboard{}).Where("id=?", *row.StoryboardID).Updates(map[string]any{"thumbnail_asset_id": output.MediaAssetID, "status": "completed"}).Error } if assetID, parseErr := uuid.Parse(fmt.Sprint(input["asset_id"])); parseErr == nil { return tx.Model(&model.ProjectAsset{}).Where("id=?", assetID).Update("image_asset_id", output.MediaAssetID).Error } } return nil }) if err != nil { _ = w.COS.Delete(ctx, key) return err } if cancelledDuringDownload { _ = w.COS.Delete(ctx, key) return w.finishCancellation(ctx, taskID, false, "用户取消生成") } if reusedOutput { _ = w.COS.Delete(ctx, key) } return nil } func (w *Generation) requeueRejected(ctx context.Context, task model.GenerationTask, message string, delay time.Duration) error { if err := w.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error { var current model.GenerationTask if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("id=?", task.ID).First(¤t).Error; err != nil { return err } next := time.Now().Add(delay) updates := map[string]any{"status": "pending_submission", "next_submit_at": next, "lease_owner": nil, "lease_until": nil, "error_message": message} if current.SlotReserved && current.SlotReleasedAt == nil { now := time.Now() updates["slot_released_at"] = now } return tx.Model(¤t).Updates(updates).Error }); err != nil { return err } if task.ChannelID != nil { return queuepkg.EnqueueID(w.Queue, queuepkg.TypeDispatchChannel, *task.ChannelID, delay) } return nil } func (w *Generation) returnUnenqueuedTask(ctx context.Context, taskID uuid.UUID, message string) { var channelID *uuid.UUID _ = w.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error { var task model.GenerationTask if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("id=?", taskID).First(&task).Error; err != nil { return err } if task.Status != "submitting" { return nil } channelID = task.ChannelID now := time.Now() return tx.Model(&task).Updates(map[string]any{ "status": "pending_submission", "next_submit_at": now.Add(5 * time.Second), "slot_released_at": now, "lease_owner": nil, "lease_until": nil, "error_message": truncate(message, 1000), }).Error }) w.enqueueDispatchAfter(channelID, 5*time.Second) } func (w *Generation) finishCancellation(ctx context.Context, taskID uuid.UUID, refund bool, message string) error { return w.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error { var task model.GenerationTask if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("id=?", taskID).First(&task).Error; err != nil { return err } if task.Status == "cancelled" { return nil } now := time.Now() updates := map[string]any{ "status": "cancelled", "finished_at": now, "next_poll_at": nil, "lease_owner": nil, "lease_until": nil, "error_code": nil, "error_message": message, } if task.SlotReserved && task.SlotReleasedAt == nil { updates["slot_released_at"] = now } premiumDrama, err := isPremiumDramaGeneration(tx, task) if err != nil { return err } // Premium-drama generation keeps its prepaid charge on cancellation. refund = refund && !premiumDrama if refund { refunded, err := billing.RefundGenerationTask(tx, &task, "视频生成失败返还") if err != nil { return err } if refunded { updates["cost_refunded"] = true } updates["actual_points"] = "0.00" } else { updates["actual_points"] = task.PrepaidPoints } if err := tx.Model(&task).Updates(updates).Error; err != nil { return err } if task.TaskType == "video_generation" && task.StoryboardID != nil { return tx.Model(&model.EpisodeStoryboard{}).Where("id=?", *task.StoryboardID).Update("status", "idle").Error } return nil }) } func (w *Generation) failAndRefund(ctx context.Context, taskID uuid.UUID, code, message string) error { return w.finishFailure(ctx, taskID, code, message, true) } func (w *Generation) failWithoutRefund(ctx context.Context, taskID uuid.UUID, code, message string) error { return w.finishFailure(ctx, taskID, code, message, false) } func (w *Generation) finishFailure(ctx context.Context, taskID uuid.UUID, code, message string, refund bool) error { return w.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error { var task model.GenerationTask if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("id=?", taskID).First(&task).Error; err != nil { return err } if task.Status == "succeeded" || task.Status == "failed" || task.Status == "cancelled" { return nil } now := time.Now() updates := map[string]any{"status": "failed", "error_code": code, "error_message": truncate(message, 4000), "finished_at": now, "lease_owner": nil, "lease_until": nil} if task.SlotReserved && task.SlotReleasedAt == nil { updates["slot_released_at"] = now } premiumDrama, err := isPremiumDramaGeneration(tx, task) if err != nil { return err } if premiumDrama && code != "upstream_rejected" && code != "upstream_unavailable" && code != "upstream_failed" { refund = false } if refund { refunded, err := billing.RefundGenerationTask(tx, &task, "视频生成失败返还") if err != nil { return err } if refunded { updates["cost_refunded"] = true } updates["actual_points"] = "0.00" } else { updates["actual_points"] = task.PrepaidPoints } if err := tx.Model(&task).Updates(updates).Error; err != nil { return err } if task.TaskType == "video_generation" && task.StoryboardID != nil { return tx.Model(&model.EpisodeStoryboard{}).Where("id=?", *task.StoryboardID).Update("status", "failed").Error } return nil }) } func isPremiumDramaGeneration(tx *gorm.DB, task model.GenerationTask) (bool, error) { if task.ProjectID == nil { return false, nil } var projectType string if err := tx.Table("creative_projects").Where("id=?", *task.ProjectID).Pluck("project_type", &projectType).Error; err != nil { return false, err } return projectType == "premium_drama", nil } func (w *Generation) enqueueDispatch(channelID *uuid.UUID) { w.enqueueDispatchAfter(channelID, 0) } func (w *Generation) enqueueDispatchAfter(channelID *uuid.UUID, delay time.Duration) { if channelID != nil { _ = queuepkg.EnqueueID(w.Queue, queuepkg.TypeDispatchChannel, *channelID, delay) } } // Recover 恢复服务中断或旧版本遗留的生成任务,并重新投递待处理工作。 func (w *Generation) Recover(ctx context.Context) error { now := time.Now() if err := w.DB.WithContext(ctx).Model(&model.GenerationTask{}). Where("task_type IN ('image_generation','video_generation') AND status='submitting' AND lease_until>'product_image'='true'))").Group("channel_id").Scan(&channels).Error; err != nil { return err } for _, item := range channels { delay := time.Until(item.NextSubmitAt) if delay < 0 { delay = 0 } _ = queuepkg.EnqueueID(w.Queue, queuepkg.TypeDispatchChannel, item.ChannelID, delay) } var polls []uuid.UUID _ = w.DB.WithContext(ctx).Model(&model.GenerationTask{}).Where("(project_id IS NOT NULL OR (task_type='image_generation' AND input_data->>'product_image'='true')) AND (status IN ('submitted','processing') OR (status='cancel_requested' AND upstream_task_id IS NOT NULL AND upstream_task_id<>''))").Pluck("id", &polls).Error for _, id := range polls { _ = queuepkg.EnqueueID(w.Queue, queuepkg.TypePollTask, id, jitter(5*time.Second)) } var downloads []uuid.UUID _ = w.DB.WithContext(ctx).Model(&model.GenerationTask{}).Where("(project_id IS NOT NULL OR (task_type='image_generation' AND input_data->>'product_image'='true')) AND status IN ('result_ready','downloading')").Pluck("id", &downloads).Error for _, id := range downloads { _ = queuepkg.EnqueueID(w.Queue, queuepkg.TypeDownloadTask, id, jitter(5*time.Second)) } return nil } func stringOr(value any, fallback string) string { text := strings.TrimSpace(fmt.Sprint(value)) if text == "" || text == "" { return fallback } return text } func intOr(value any, fallback int) int { parsed, err := strconv.Atoi(fmt.Sprint(value)) if err != nil { return fallback } return parsed } func retryDelay(attempt int, retryAfter string) time.Duration { if seconds, err := strconv.Atoi(retryAfter); err == nil && seconds > 0 { return time.Duration(seconds)*time.Second + jitter(3*time.Second) } schedule := []time.Duration{15 * time.Second, 30 * time.Second, time.Minute, 2 * time.Minute, 5 * time.Minute} if attempt < 1 { attempt = 1 } if attempt > len(schedule) { attempt = len(schedule) } return schedule[attempt-1] + jitter(5*time.Second) } // canRetrySubmission 判断生成任务是否仍在允许的自动提交次数内。 func canRetrySubmission(attempt int) bool { return attempt < maxSubmitAttempts } func jitter(max time.Duration) time.Duration { if max <= 0 { return 0 } return time.Duration(rand.Int63n(int64(max))) } func truncate(value string, max int) string { if len(value) > max { return value[:max] } return value } func NewHTTPClient(maxConnections int) *http.Client { if maxConnections < 10 { maxConnections = 10 } transport := &http.Transport{ Proxy: http.ProxyFromEnvironment, DialContext: (&net.Dialer{Timeout: 15 * time.Second, KeepAlive: 30 * time.Second}).DialContext, MaxIdleConns: maxConnections, MaxIdleConnsPerHost: maxConnections, MaxConnsPerHost: maxConnections, // ResponseHeaderTimeout 设为 0 表示不限制,由各调用点的 context 超时统一控制,避免非流式 AI 推理在 60 秒内未返回响应头被误判超时 IdleConnTimeout: 90 * time.Second, TLSHandshakeTimeout: 10 * time.Second, ResponseHeaderTimeout: 0, } return &http.Client{Transport: transport} }