851 lines
33 KiB
Go
851 lines
33 KiB
Go
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_until<?)
|
|
)
|
|
SELECT task.* FROM generation_tasks task JOIN ranked ON ranked.id=task.id
|
|
WHERE ranked.allowance>0 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<?", now).
|
|
Updates(map[string]any{
|
|
"status": "pending_submission", "next_submit_at": now, "slot_released_at": now,
|
|
"error_code": nil, "error_message": "提交中断，正在自动重试", "lease_owner": nil, "lease_until": nil,
|
|
}).Error; err != nil {
|
|
return err
|
|
}
|
|
channels := make([]struct {
|
|
ChannelID uuid.UUID
|
|
NextSubmitAt time.Time
|
|
}, 0)
|
|
if err := w.DB.WithContext(ctx).Model(&model.GenerationTask{}).
|
|
Select("channel_id,min(coalesce(next_submit_at,?)) AS next_submit_at", now).
|
|
Where("status='pending_submission' AND channel_id IS NOT NULL AND (project_id IS NOT NULL OR (task_type='image_generation' AND input_data->>'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 == "<nil>" {
|
|
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}
|
|
}
|