From e0d5156115881780328d31fe9bce7fe25aa9c6c7 Mon Sep 17 00:00:00 2001 From: feitianbubu Date: Mon, 20 Jul 2026 22:03:13 +0800 Subject: [PATCH] fix: prevent duplicate suno task refunds via cas status update (#6074) * fix: prevent duplicate suno task refunds via cas status update * fix: reconcile failed task refunds --------- Co-authored-by: CaIon --- controller/system_task_handlers.go | 2 +- model/task.go | 81 +++++++++++++- model/task_cas_test.go | 105 ++++++++++++++++++ service/task_billing.go | 15 ++- service/task_billing_test.go | 39 ++++++- service/task_polling.go | 65 ++++++++++-- service/task_polling_test.go | 165 +++++++++++++++++++++++++++++ 7 files changed, 456 insertions(+), 16 deletions(-) diff --git a/controller/system_task_handlers.go b/controller/system_task_handlers.go index c31059d1..2168b451 100644 --- a/controller/system_task_handlers.go +++ b/controller/system_task_handlers.go @@ -140,7 +140,7 @@ type asyncTaskPollHandler struct{} func (asyncTaskPollHandler) Type() string { return model.SystemTaskTypeAsyncTaskPoll } func (asyncTaskPollHandler) Enabled() bool { - return constant.UpdateTask && model.HasUnfinishedSyncTasks() + return constant.UpdateTask && model.HasTaskPollingWork() } func (asyncTaskPollHandler) Interval() time.Duration { return 15 * time.Second } diff --git a/model/task.go b/model/task.go index cf936ed1..ecaf70f3 100644 --- a/model/task.go +++ b/model/task.go @@ -41,6 +41,10 @@ const ( TaskStatusUnknown = "UNKNOWN" ) +// TaskRefundLegacyCutoff separates legacy timeout tasks that intentionally +// do not receive automatic refunds from tasks covered by reconciliation. +const TaskRefundLegacyCutoff int64 = 1740182400 // 2025-02-22 00:00:00 UTC + type Task struct { ID int64 `json:"id" gorm:"primary_key;AUTO_INCREMENT"` CreatedAt int64 `json:"created_at" gorm:"index"` @@ -304,6 +308,28 @@ func GetTimedOutUnfinishedTasks(cutoffUnix int64, limit int) []*Task { return tasks } +// GetUnrefundedFailedTasks returns failed tasks whose non-zero quota marks a +// pending refund. Legacy timeout tasks are excluded before LIMIT is applied so +// they cannot starve refundable tasks from the reconciliation sweep. +func GetUnrefundedFailedTasks(updatedBefore int64, limit int) []*Task { + if limit <= 0 { + return nil + } + + var tasks []*Task + err := DB.Where("status = ?", TaskStatusFailure). + Where("quota != ?", 0). + Where("updated_at <= ?", updatedBefore). + Where("(submit_time <= ? OR submit_time >= ?)", 0, TaskRefundLegacyCutoff). + Order("id"). + Limit(limit). + Find(&tasks).Error + if err != nil { + return nil + } + return tasks +} + func GetAllUnFinishSyncTasks(limit int) []*Task { var tasks []*Task var err error @@ -330,6 +356,24 @@ func HasUnfinishedSyncTasks() bool { return err == nil && id != 0 } +// HasTaskPollingWork reports whether polling has either an unfinished task or +// a failed task with a pending, non-legacy refund. The latter keeps the system +// task scheduler active when reconciliation is the only work left. +func HasTaskPollingWork() bool { + if HasUnfinishedSyncTasks() { + return true + } + + var id int64 + err := DB.Model(&Task{}). + Where("status = ?", TaskStatusFailure). + Where("quota != ?", 0). + Where("(submit_time <= ? OR submit_time >= ?)", 0, TaskRefundLegacyCutoff). + Limit(1). + Pluck("id", &id).Error + return err == nil && id != 0 +} + func GetByOnlyTaskId(taskId string) (*Task, bool, error) { if taskId == "" { return nil, false, nil @@ -421,9 +465,44 @@ func (t *Task) UpdateQuota() error { return DB.Model(t).Update("quota", t.Quota).Error } +// ClaimQuotaForRefund atomically clears an expected non-zero quota. A true +// result grants the caller ownership of the corresponding refund attempt. +func ClaimQuotaForRefund(id int64, expectedQuota int) (bool, error) { + if expectedQuota == 0 { + return false, nil + } + + result := DB.Model(&Task{}). + Where("id = ? AND quota = ?", id, expectedQuota). + Update("quota", 0) + if result.Error != nil { + return false, result.Error + } + return result.RowsAffected > 0, nil +} + +// RestoreQuotaAfterFailedRefund restores a claimed quota marker only while it +// is still zero. It is used when the observable funding adjustment fails, so a +// later reconciliation pass can retry without overwriting another writer. +func RestoreQuotaAfterFailedRefund(id int64, quota int) (bool, error) { + if quota == 0 { + return false, nil + } + + result := DB.Model(&Task{}). + Where("id = ? AND quota = ?", id, 0). + Update("quota", quota) + if result.Error != nil { + return false, result.Error + } + return result.RowsAffected > 0, nil +} + // UpdateWithStatus performs a conditional UPDATE guarded by fromStatus (CAS). // Returns (true, nil) if this caller won the update, (false, nil) if -// another process already moved the task out of fromStatus. +// another process already moved the task out of fromStatus. MySQL commonly +// reports changed rows rather than matched rows, so a same-value no-op update +// can also return false even when the status predicate still matched. // // Uses Model().Select("*").Updates() instead of Save() because GORM's Save // falls back to INSERT ON CONFLICT when the WHERE-guarded UPDATE matches diff --git a/model/task_cas_test.go b/model/task_cas_test.go index a53804d3..d99a8d89 100644 --- a/model/task_cas_test.go +++ b/model/task_cas_test.go @@ -256,3 +256,108 @@ func TestUpdateWithStatus_ConcurrentWinner(t *testing.T) { } assert.Equal(t, 1, winCount, "exactly one goroutine should win the CAS") } + +func TestClaimQuotaForRefund_OnlyOneClaimSucceeds(t *testing.T) { + truncateTables(t) + + task := &Task{ + TaskID: "task_refund_claim", + Status: TaskStatusFailure, + Quota: 1000, + Data: json.RawMessage(`{}`), + } + insertTask(t, task) + + claimed, err := ClaimQuotaForRefund(task.ID, task.Quota) + require.NoError(t, err) + assert.True(t, claimed) + + claimed, err = ClaimQuotaForRefund(task.ID, task.Quota) + require.NoError(t, err) + assert.False(t, claimed) + + var reloaded Task + require.NoError(t, DB.First(&reloaded, task.ID).Error) + assert.Zero(t, reloaded.Quota) +} + +func TestGetUnrefundedFailedTasks_FiltersAndLimits(t *testing.T) { + truncateTables(t) + + tasks := []*Task{ + {TaskID: "failed_refundable_1", Status: TaskStatusFailure, Quota: 100, SubmitTime: TaskRefundLegacyCutoff, Data: json.RawMessage(`{}`)}, + {TaskID: "failed_refundable_2", Status: TaskStatusFailure, Quota: 200, SubmitTime: TaskRefundLegacyCutoff + 1, Data: json.RawMessage(`{}`)}, + {TaskID: "legacy_failed", Status: TaskStatusFailure, Quota: 400, SubmitTime: TaskRefundLegacyCutoff - 1, Data: json.RawMessage(`{}`)}, + {TaskID: "failed_without_quota", Status: TaskStatusFailure, Quota: 0, Data: json.RawMessage(`{}`)}, + {TaskID: "successful_with_quota", Status: TaskStatusSuccess, Quota: 300, Data: json.RawMessage(`{}`)}, + } + for _, task := range tasks { + insertTask(t, task) + } + + updatedBefore := time.Now().Unix() + 1 + found := GetUnrefundedFailedTasks(updatedBefore, 1) + require.Len(t, found, 1) + assert.Equal(t, tasks[0].ID, found[0].ID) + + found = GetUnrefundedFailedTasks(updatedBefore, 10) + require.Len(t, found, 2) + assert.Equal(t, []int64{tasks[0].ID, tasks[1].ID}, []int64{found[0].ID, found[1].ID}) + + assert.Empty(t, GetUnrefundedFailedTasks(updatedBefore, 0)) +} + +func TestRestoreQuotaAfterFailedRefund_OnlyRestoresClaimedMarker(t *testing.T) { + truncateTables(t) + + task := &Task{ + TaskID: "task_refund_restore", + Status: TaskStatusFailure, + Quota: 750, + Data: json.RawMessage(`{}`), + } + insertTask(t, task) + + claimed, err := ClaimQuotaForRefund(task.ID, task.Quota) + require.NoError(t, err) + require.True(t, claimed) + + restored, err := RestoreQuotaAfterFailedRefund(task.ID, task.Quota) + require.NoError(t, err) + assert.True(t, restored) + + restored, err = RestoreQuotaAfterFailedRefund(task.ID, task.Quota) + require.NoError(t, err) + assert.False(t, restored) + + var reloaded Task + require.NoError(t, DB.First(&reloaded, task.ID).Error) + assert.Equal(t, task.Quota, reloaded.Quota) +} + +func TestHasTaskPollingWork_IncludesOnlyRefundableFailedTasks(t *testing.T) { + truncateTables(t) + assert.False(t, HasTaskPollingWork()) + + legacy := &Task{ + TaskID: "legacy_failed_work", + Status: TaskStatusFailure, + Progress: "100%", + Quota: 500, + SubmitTime: TaskRefundLegacyCutoff - 1, + Data: json.RawMessage(`{}`), + } + insertTask(t, legacy) + assert.False(t, HasTaskPollingWork()) + + refundable := &Task{ + TaskID: "refundable_failed_work", + Status: TaskStatusFailure, + Progress: "100%", + Quota: 500, + SubmitTime: TaskRefundLegacyCutoff, + Data: json.RawMessage(`{}`), + } + insertTask(t, refundable) + assert.True(t, HasTaskPollingWork()) +} diff --git a/service/task_billing.go b/service/task_billing.go index 29321bd9..51fb77c3 100644 --- a/service/task_billing.go +++ b/service/task_billing.go @@ -162,16 +162,17 @@ func taskModelName(task *model.Task) string { // RefundTaskQuota 统一的任务失败退款逻辑。 // 当异步任务失败时，将预扣的 quota 退还给用户（支持钱包和订阅），并退还令牌额度。 -func RefundTaskQuota(ctx context.Context, task *model.Task, reason string) { +// 返回资金来源是否已成功退还；失败时保留 quota 作为后续对账标记。 +func RefundTaskQuota(ctx context.Context, task *model.Task, reason string) bool { quota := task.Quota if quota == 0 { - return + return true } // 1. 退还资金来源（钱包或订阅） if err := taskAdjustFunding(task, -quota); err != nil { logger.LogWarn(ctx, fmt.Sprintf("退还资金来源失败 task %s: %s", task.TaskID, err.Error())) - return + return false } // 2. 退还令牌额度 @@ -192,6 +193,14 @@ func RefundTaskQuota(ctx context.Context, task *model.Task, reason string) { Group: task.Group, Other: other, }) + + // 4. 资金退款完成后再清除持久化标记；失败时保留非零 quota， + // 由后续对账重试。回写失败必须显式告警，避免漏掉潜在的重复退款风险。 + task.Quota = 0 + if err := task.UpdateQuota(); err != nil { + logger.LogError(ctx, fmt.Sprintf("退款成功但清除 task quota 失败 task %s: %s", task.TaskID, err.Error())) + } + return true } // RecalculateTaskQuota 通用的异步差额结算。 diff --git a/service/task_billing_test.go b/service/task_billing_test.go index fb1c2915..53e3f680 100644 --- a/service/task_billing_test.go +++ b/service/task_billing_test.go @@ -270,6 +270,13 @@ func getSubscriptionUsed(t *testing.T, id int) int64 { return sub.AmountUsed } +func getTaskQuota(t *testing.T, id int64) int { + t.Helper() + var task model.Task + require.NoError(t, model.DB.Select("quota").Where("id = ?", id).First(&task).Error) + return task.Quota +} + func getLastLog(t *testing.T) *model.Log { t.Helper() var log model.Log @@ -304,8 +311,9 @@ func TestRefundTaskQuota_Wallet(t *testing.T) { seedChannel(t, channelID) task := makeTask(userID, channelID, preConsumed, tokenID, BillingSourceWallet, 0) + require.NoError(t, model.DB.Create(task).Error) - RefundTaskQuota(ctx, task, "task failed: upstream error") + assert.True(t, RefundTaskQuota(ctx, task, "task failed: upstream error")) // User quota should increase by preConsumed assert.Equal(t, initQuota+preConsumed, getUserQuota(t, userID)) @@ -320,6 +328,8 @@ func TestRefundTaskQuota_Wallet(t *testing.T) { assert.Equal(t, model.LogTypeRefund, log.Type) assert.Equal(t, preConsumed, log.Quota) assert.Equal(t, "test-model", log.ModelName) + assert.Zero(t, task.Quota) + assert.Zero(t, getTaskQuota(t, task.ID)) } func TestRefundTaskQuota_Subscription(t *testing.T) { @@ -337,8 +347,9 @@ func TestRefundTaskQuota_Subscription(t *testing.T) { seedSubscription(t, subID, userID, subTotal, subUsed) task := makeTask(userID, channelID, preConsumed, tokenID, BillingSourceSubscription, subID) + require.NoError(t, model.DB.Create(task).Error) - RefundTaskQuota(ctx, task, "subscription task failed") + assert.True(t, RefundTaskQuota(ctx, task, "subscription task failed")) // Subscription used should decrease by preConsumed assert.Equal(t, subUsed-int64(preConsumed), getSubscriptionUsed(t, subID)) @@ -349,6 +360,7 @@ func TestRefundTaskQuota_Subscription(t *testing.T) { log := getLastLog(t) require.NotNil(t, log) assert.Equal(t, model.LogTypeRefund, log.Type) + assert.Zero(t, getTaskQuota(t, task.ID)) } func TestRefundTaskQuota_ZeroQuota(t *testing.T) { @@ -360,7 +372,7 @@ func TestRefundTaskQuota_ZeroQuota(t *testing.T) { task := makeTask(userID, 0, 0, 0, BillingSourceWallet, 0) - RefundTaskQuota(ctx, task, "zero quota task") + assert.True(t, RefundTaskQuota(ctx, task, "zero quota task")) // No change to user quota assert.Equal(t, 5000, getUserQuota(t, userID)) @@ -380,8 +392,9 @@ func TestRefundTaskQuota_NoToken(t *testing.T) { seedChannel(t, channelID) task := makeTask(userID, channelID, preConsumed, 0, BillingSourceWallet, 0) // TokenId=0 + require.NoError(t, model.DB.Create(task).Error) - RefundTaskQuota(ctx, task, "no token task failed") + assert.True(t, RefundTaskQuota(ctx, task, "no token task failed")) // User quota refunded assert.Equal(t, initQuota+preConsumed, getUserQuota(t, userID)) @@ -390,6 +403,23 @@ func TestRefundTaskQuota_NoToken(t *testing.T) { log := getLastLog(t) require.NotNil(t, log) assert.Equal(t, model.LogTypeRefund, log.Type) + assert.Zero(t, getTaskQuota(t, task.ID)) +} + +func TestRefundTaskQuota_FundingFailureKeepsPendingMarker(t *testing.T) { + truncate(t) + ctx := context.Background() + + const userID, preConsumed = 5, 1200 + seedUser(t, userID, 5000) + task := makeTask(userID, 0, preConsumed, 0, BillingSourceSubscription, 9999) + task.Status = model.TaskStatusFailure + require.NoError(t, model.DB.Create(task).Error) + + assert.False(t, RefundTaskQuota(ctx, task, "subscription missing")) + assert.Equal(t, preConsumed, task.Quota) + assert.Equal(t, preConsumed, getTaskQuota(t, task.ID)) + assert.Equal(t, int64(0), countLogs(t)) } // =========================================================================== @@ -608,6 +638,7 @@ func TestCASGuardedRefund_Win(t *testing.T) { var reloaded model.Task require.NoError(t, model.DB.First(&reloaded, task.ID).Error) assert.EqualValues(t, model.TaskStatusFailure, reloaded.Status) + assert.Zero(t, reloaded.Quota) // Refund should have happened assert.Equal(t, initQuota+preConsumed, getUserQuota(t, userID)) diff --git a/service/task_polling.go b/service/task_polling.go index 179fa734..0ee3be24 100644 --- a/service/task_polling.go +++ b/service/task_polling.go @@ -37,6 +37,11 @@ type TaskPollingAdaptor interface { // 打破 service -> relay -> relay/channel -> service 的循环依赖。 var GetTaskAdaptorFunc func(platform constant.TaskPlatform) TaskPollingAdaptor +const ( + refundReconciliationLimit = 100 + refundReconciliationGracePeriod = 30 * time.Second +) + // sweepTimedOutTasks 在主轮询之前独立清理超时任务。 // 每次最多处理 100 条，剩余的下个周期继续处理。 // 使用 per-task CAS (UpdateWithStatus) 防止覆盖被正常轮询已推进的任务。 @@ -50,14 +55,13 @@ func sweepTimedOutTasks(ctx context.Context) { return } - const legacyTaskCutoff int64 = 1740182400 // 2026-02-22 00:00:00 UTC reason := fmt.Sprintf("任务超时（%d分钟）", constant.TaskTimeoutMinutes) legacyReason := "任务超时（旧系统遗留任务，不进行退款，请联系管理员）" now := time.Now().Unix() timedOutCount := 0 for _, task := range tasks { - isLegacy := task.SubmitTime > 0 && task.SubmitTime < legacyTaskCutoff + isLegacy := task.SubmitTime > 0 && task.SubmitTime < model.TaskRefundLegacyCutoff oldStatus := task.Status task.Status = model.TaskStatusFailure @@ -65,6 +69,8 @@ func sweepTimedOutTasks(ctx context.Context) { task.FinishTime = now if isLegacy { task.FailReason = legacyReason + // 旧系统任务明确不退款，随终态 CAS 一并清掉 quota，避免被后续对账误判。 + task.Quota = 0 } else { task.FailReason = reason } @@ -89,6 +95,43 @@ func sweepTimedOutTasks(ctx context.Context) { } } +// sweepUnrefundedFailedTasks 重试已落 FAILURE 终态但仍保留 quota 的欠退款任务。 +// 先等待一个短暂宽限期，让终态 CAS 的胜出者完成主路径即时退款，避免正常 +// 轮询与对账同时处理刚失败的任务。 +func sweepUnrefundedFailedTasks(ctx context.Context) { + updatedBefore := time.Now().Add(-refundReconciliationGracePeriod).Unix() + tasks := model.GetUnrefundedFailedTasks(updatedBefore, refundReconciliationLimit) + for _, task := range tasks { + if ctx.Err() != nil { + return + } + + quota := task.Quota + claimed, err := model.ClaimQuotaForRefund(task.ID, quota) + if err != nil { + logger.LogError(ctx, fmt.Sprintf("sweepUnrefundedFailedTasks claim error for task %s: %v", task.TaskID, err)) + continue + } + if !claimed { + logger.LogDebug(ctx, "sweepUnrefundedFailedTasks: task %s claim lost, skip refund", task.TaskID) + continue + } + + // 对账先清 marker 再退款，确保并发 sweep 只有一个实际退款者。若进程在 + // claim 后、退款前崩溃，会偏向漏退而不是双退，需由人工账务对账兜底。 + if RefundTaskQuota(ctx, task, task.FailReason) { + continue + } + + restored, restoreErr := model.RestoreQuotaAfterFailedRefund(task.ID, quota) + if restoreErr != nil { + logger.LogError(ctx, fmt.Sprintf("sweepUnrefundedFailedTasks restore quota error for task %s: %v", task.TaskID, restoreErr)) + } else if !restored { + logger.LogError(ctx, fmt.Sprintf("sweepUnrefundedFailedTasks could not restore quota marker for task %s", task.TaskID)) + } + } +} + // TaskPollSummary is the result recorded on an async_task_poll system task row, // summarizing one polling pass. type TaskPollSummary struct { @@ -113,6 +156,7 @@ func RunTaskPollingOnce(ctx context.Context, report func(processed, total int)) common.SysLog("任务进度轮询开始") sweepTimedOutTasks(ctx) + sweepUnrefundedFailedTasks(ctx) allTasks := model.GetAllUnFinishSyncTasks(constant.TaskQueryLimit) summary.UnfinishedTasks = len(allTasks) platformTask := make(map[constant.TaskPlatform][]*model.Task) @@ -277,24 +321,31 @@ func updateSunoTasks(ctx context.Context, channelId int, taskIds []string, taskM continue } + prevStatus := task.Status task.Status = lo.If(model.TaskStatus(responseItem.Status) != "", model.TaskStatus(responseItem.Status)).Else(task.Status) task.FailReason = lo.If(responseItem.FailReason != "", responseItem.FailReason).Else(task.FailReason) task.SubmitTime = lo.If(responseItem.SubmitTime != 0, responseItem.SubmitTime).Else(task.SubmitTime) task.StartTime = lo.If(responseItem.StartTime != 0, responseItem.StartTime).Else(task.StartTime) task.FinishTime = lo.If(responseItem.FinishTime != 0, responseItem.FinishTime).Else(task.FinishTime) - if responseItem.FailReason != "" || task.Status == model.TaskStatusFailure { + isFailure := responseItem.FailReason != "" || task.Status == model.TaskStatusFailure + if isFailure { logger.LogInfo(ctx, task.TaskID+" 构建失败，"+task.FailReason) + task.Status = model.TaskStatusFailure task.Progress = "100%" - RefundTaskQuota(ctx, task, task.FailReason) } if responseItem.Status == model.TaskStatusSuccess { task.Progress = "100%" } task.Data = responseItem.Data - err = task.Update() + // 持久化走 CAS，防止重叠轮询/sweep/多实例/持久化失败重试导致重复退款或覆盖终态。 + won, err := task.UpdateWithStatus(prevStatus) if err != nil { - common.SysLog("UpdateSunoTask task error: " + err.Error()) + logger.LogError(ctx, fmt.Sprintf("UpdateSunoTask task %s error: %v", task.TaskID, err)) + } else if !won { + logger.LogWarn(ctx, fmt.Sprintf("Task %s CAS lost or no-op update, skip billing", task.TaskID)) + } else if isFailure && prevStatus != model.TaskStatusFailure && task.Quota != 0 { + RefundTaskQuota(ctx, task, task.FailReason) } } return nil @@ -568,7 +619,7 @@ func updateVideoSingleTask(ctx context.Context, adaptor TaskPollingAdaptor, ch * shouldRefund = false shouldSettle = false } else if !won { - logger.LogWarn(ctx, fmt.Sprintf("Task %s already transitioned by another process, skip billing", task.TaskID)) + logger.LogWarn(ctx, fmt.Sprintf("Task %s CAS lost or no-op update, skip billing", task.TaskID)) shouldRefund = false shouldSettle = false } diff --git a/service/task_polling_test.go b/service/task_polling_test.go index 3164d0a9..dd70bf81 100644 --- a/service/task_polling_test.go +++ b/service/task_polling_test.go @@ -29,6 +29,45 @@ type taskPollingFetchAdaptor struct { blockOnce sync.Once } +type sunoFailurePollingAdaptor struct { + failReason string +} + +func (a *sunoFailurePollingAdaptor) Init(_ *relaycommon.RelayInfo) {} + +func (a *sunoFailurePollingAdaptor) FetchTask(_ string, _ string, body map[string]any, _ string) (*http.Response, error) { + taskIDs, _ := body["ids"].([]string) + items := make([]dto.SunoDataResponse, 0, len(taskIDs)) + for _, taskID := range taskIDs { + items = append(items, dto.SunoDataResponse{ + TaskID: taskID, + Status: string(model.TaskStatusFailure), + FailReason: a.failReason, + FinishTime: time.Now().Unix(), + }) + } + + responseBody, err := common.Marshal(dto.TaskResponse[[]dto.SunoDataResponse]{ + Code: dto.TaskSuccessCode, + Data: items, + }) + if err != nil { + return nil, err + } + return &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(bytes.NewReader(responseBody)), + }, nil +} + +func (a *sunoFailurePollingAdaptor) ParseTaskResult([]byte) (*relaycommon.TaskInfo, error) { + return nil, nil +} + +func (a *sunoFailurePollingAdaptor) AdjustBillingOnComplete(_ *model.Task, _ *relaycommon.TaskInfo) int { + return 0 +} + func (a *taskPollingFetchAdaptor) Init(_ *relaycommon.RelayInfo) {} func (a *taskPollingFetchAdaptor) FetchTask(_ string, _ string, body map[string]any, _ string) (*http.Response, error) { @@ -331,3 +370,129 @@ func TestUpdateVideoTasksMixedChannelSleepSettings(t *testing.T) { require.ErrorIs(t, err, context.DeadlineExceeded) assert.ElementsMatch(t, []string{"upstream_sleepy_1", "upstream_fast_1", "upstream_fast_2"}, adaptor.fetchedTaskIDs()) } + +func TestUpdateSunoTasksStalePollsRefundExactlyOnce(t *testing.T) { + truncate(t) + + const userID, tokenID, channelID = 401, 401, 401 + const initialUserQuota, initialTokenQuota, taskQuota = 10_000, 6_000, 2_500 + const publicTaskID, upstreamTaskID = "suno_public_refund_once", "suno_upstream_refund_once" + + seedUser(t, userID, initialUserQuota) + seedToken(t, tokenID, userID, "sk-suno-refund-once", initialTokenQuota) + baseURL := "https://suno.invalid" + require.NoError(t, model.DB.Create(&model.Channel{ + Id: channelID, + Type: constant.ChannelTypeSunoAPI, + Name: "suno_refund_once", + Key: "sk-suno-channel", + Status: common.ChannelStatusEnabled, + BaseURL: &baseURL, + }).Error) + + task := makeTask(userID, channelID, taskQuota, tokenID, BillingSourceWallet, 0) + task.TaskID = publicTaskID + task.Platform = constant.TaskPlatformSuno + task.Status = model.TaskStatusInProgress + task.Progress = "50%" + task.SubmitTime = model.TaskRefundLegacyCutoff + task.PrivateData.UpstreamTaskID = upstreamTaskID + require.NoError(t, model.DB.Create(task).Error) + + var firstPollTask model.Task + var staleSecondPollTask model.Task + require.NoError(t, model.DB.First(&firstPollTask, task.ID).Error) + require.NoError(t, model.DB.First(&staleSecondPollTask, task.ID).Error) + + adaptor := &sunoFailurePollingAdaptor{failReason: "upstream failed"} + previousFactory := GetTaskAdaptorFunc + GetTaskAdaptorFunc = func(constant.TaskPlatform) TaskPollingAdaptor { return adaptor } + t.Cleanup(func() { GetTaskAdaptorFunc = previousFactory }) + + require.NoError(t, updateSunoTasks(context.Background(), channelID, []string{upstreamTaskID}, map[string]*model.Task{ + upstreamTaskID: &firstPollTask, + })) + require.NoError(t, updateSunoTasks(context.Background(), channelID, []string{upstreamTaskID}, map[string]*model.Task{ + upstreamTaskID: &staleSecondPollTask, + })) + + var reloaded model.Task + require.NoError(t, model.DB.First(&reloaded, task.ID).Error) + assert.EqualValues(t, model.TaskStatusFailure, reloaded.Status) + assert.Zero(t, reloaded.Quota) + assert.Equal(t, initialUserQuota+taskQuota, getUserQuota(t, userID)) + assert.Equal(t, initialTokenQuota+taskQuota, getTokenRemainQuota(t, tokenID)) + assert.Equal(t, int64(1), countLogs(t)) +} + +func TestSweepUnrefundedFailedTasksRefundsModernTaskAndSkipsLegacy(t *testing.T) { + truncate(t) + + const userID = 402 + const initialQuota, modernTaskQuota, legacyTaskQuota = 10_000, 1_200, 1_800 + seedUser(t, userID, initialQuota) + + modernTask := makeTask(userID, 0, modernTaskQuota, 0, BillingSourceWallet, 0) + modernTask.TaskID = "modern_failed_pending_refund" + modernTask.Status = model.TaskStatusFailure + modernTask.Progress = "100%" + modernTask.SubmitTime = model.TaskRefundLegacyCutoff + modernTask.UpdatedAt = time.Now().Add(-time.Minute).Unix() + require.NoError(t, model.DB.Create(modernTask).Error) + + legacyTask := makeTask(userID, 0, legacyTaskQuota, 0, BillingSourceWallet, 0) + legacyTask.TaskID = "legacy_failed_without_refund" + legacyTask.Status = model.TaskStatusFailure + legacyTask.Progress = "100%" + legacyTask.SubmitTime = model.TaskRefundLegacyCutoff - 1 + legacyTask.UpdatedAt = time.Now().Add(-time.Minute).Unix() + require.NoError(t, model.DB.Create(legacyTask).Error) + + sweepUnrefundedFailedTasks(context.Background()) + sweepUnrefundedFailedTasks(context.Background()) + + var reloadedModern model.Task + var reloadedLegacy model.Task + require.NoError(t, model.DB.First(&reloadedModern, modernTask.ID).Error) + require.NoError(t, model.DB.First(&reloadedLegacy, legacyTask.ID).Error) + assert.Zero(t, reloadedModern.Quota) + assert.Equal(t, legacyTaskQuota, reloadedLegacy.Quota) + assert.Equal(t, initialQuota+modernTaskQuota, getUserQuota(t, userID)) + assert.Equal(t, int64(1), countLogs(t)) +} + +func TestSweepUnrefundedFailedTasksRestoresMarkerAfterFundingFailure(t *testing.T) { + truncate(t) + + const userID, subscriptionID, taskQuota = 404, 404, 900 + const subscriptionUsed int64 = 5_000 + seedUser(t, userID, 0) + + task := makeTask(userID, 0, taskQuota, 0, BillingSourceSubscription, subscriptionID) + task.TaskID = "subscription_failed_pending_refund" + task.Status = model.TaskStatusFailure + task.Progress = "100%" + task.SubmitTime = model.TaskRefundLegacyCutoff + task.UpdatedAt = time.Now().Add(-time.Minute).Unix() + require.NoError(t, model.DB.Create(task).Error) + + sweepUnrefundedFailedTasks(context.Background()) + + var afterFailedRefund model.Task + require.NoError(t, model.DB.First(&afterFailedRefund, task.ID).Error) + assert.Equal(t, taskQuota, afterFailedRefund.Quota) + assert.Equal(t, int64(0), countLogs(t)) + + seedSubscription(t, subscriptionID, userID, 10_000, subscriptionUsed) + require.NoError(t, model.DB.Model(&model.Task{}). + Where("id = ?", task.ID). + UpdateColumn("updated_at", time.Now().Add(-time.Minute).Unix()).Error) + + sweepUnrefundedFailedTasks(context.Background()) + + var afterSuccessfulRetry model.Task + require.NoError(t, model.DB.First(&afterSuccessfulRetry, task.ID).Error) + assert.Zero(t, afterSuccessfulRetry.Quota) + assert.Equal(t, subscriptionUsed-int64(taskQuota), getSubscriptionUsed(t, subscriptionID)) + assert.Equal(t, int64(1), countLogs(t)) +}