Jaymr/new-api
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor: update task refund logic and remove legacy handling

Browse Source
This commit is contained in:
CaIon
2026-07-26 20:47:52 +08:00
parent 2d23cdf291
commit bc14c18f60
6 changed files with 64 additions and 310 deletions
Download Patch File Download Diff File Expand all files Collapse all files
service/task_billing.go
+3 -3
View File
@@ -162,7 +162,7 @@ func taskModelName(task *model.Task) string {
// RefundTaskQuota 统一的任务失败退款逻辑。
// 当异步任务失败时，将预扣的 quota 退还给用户（支持钱包和订阅），并退还令牌额度。
// 返回资金来源是否已成功退还；失败时保留 quota 作为后续对账标记。
// 返回资金来源是否已成功退还；失败时保留 quota，供显式重试或人工对账。
func RefundTaskQuota(ctx context.Context, task *model.Task, reason string) bool {
quota := task.Quota
if quota == 0 {
@@ -194,8 +194,8 @@ func RefundTaskQuota(ctx context.Context, task *model.Task, reason string) bool
Other: other,
})
// 4. 资金退款完成后再清除持久化标记；失败时保留非零 quota，
// 由后续对账重试。回写失败必须显式告警，避免漏掉潜在的重复退款风险。
// 4. 资金退款完成后再清除持久化标记。
// 回写失败必须显式告警，避免漏掉潜在的重复退款风险。
task.Quota = 0
if err := task.UpdateQuota(); err != nil {
logger.LogError(ctx, fmt.Sprintf("退款成功但清除 task quota 失败 task %s: %s", task.TaskID, err.Error()))
service/task_polling.go
+2 -44
View File
@@ -37,11 +37,6 @@ 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) 防止覆盖被正常轮询已推进的任务。
@@ -69,7 +64,8 @@ func sweepTimedOutTasks(ctx context.Context) {
task.FinishTime = now
if isLegacy {
task.FailReason = legacyReason
// 旧系统任务明确不退款，随终态 CAS 一并清掉 quota，避免被后续对账误判。
// 旧系统任务明确不退款，随终态 CAS 一并清掉 quota，
// 避免留下可再次退款的计费状态。
task.Quota = 0
} else {
task.FailReason = reason
@@ -95,43 +91,6 @@ 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 {
@@ -156,7 +115,6 @@ 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)
service/task_polling_test.go
+55 -56
View File
@@ -395,7 +395,7 @@ func TestUpdateSunoTasksStalePollsRefundExactlyOnce(t *testing.T) {
task.Platform = constant.TaskPlatformSuno
task.Status = model.TaskStatusInProgress
task.Progress = "50%"
task.SubmitTime = model.TaskRefundLegacyCutoff
task.SubmitTime = time.Now().Unix()
task.PrivateData.UpstreamTaskID = upstreamTaskID
require.NoError(t, model.DB.Create(task).Error)
@@ -425,74 +425,73 @@ func TestUpdateSunoTasksStalePollsRefundExactlyOnce(t *testing.T) {
assert.Equal(t, int64(1), countLogs(t))
}
func TestSweepUnrefundedFailedTasksRefundsModernTaskAndSkipsLegacy(t *testing.T) {
func TestRunTaskPollingOnceDoesNotRefundHistoricalFailedTask(t *testing.T) {
truncate(t)
const userID = 402
const initialQuota, modernTaskQuota, legacyTaskQuota = 10_000, 1_200, 1_800
const userID, initialQuota, taskQuota = 402, 10_000, 1_200
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 := makeTask(userID, 0, taskQuota, 0, BillingSourceWallet, 0)
task.TaskID = "historical_failed_already_refunded"
task.Status = model.TaskStatusFailure
task.Progress = "100%"
task.SubmitTime = model.TaskRefundLegacyCutoff
task.SubmitTime = time.Now().Add(-90 * 24 * time.Hour).Unix()
task.UpdatedAt = time.Now().Add(-time.Minute).Unix()
require.NoError(t, model.DB.Create(task).Error)
sweepUnrefundedFailedTasks(context.Background())
previousFactory := GetTaskAdaptorFunc
GetTaskAdaptorFunc = func(constant.TaskPlatform) TaskPollingAdaptor {
return &taskPollingFetchAdaptor{}
}
t.Cleanup(func() { GetTaskAdaptorFunc = previousFactory })
var afterFailedRefund model.Task
require.NoError(t, model.DB.First(&afterFailedRefund, task.ID).Error)
assert.Equal(t, taskQuota, afterFailedRefund.Quota)
summary := RunTaskPollingOnce(context.Background(), nil)
assert.Zero(t, summary.UnfinishedTasks)
assert.Equal(t, initialQuota, getUserQuota(t, userID))
assert.Equal(t, taskQuota, getTaskQuota(t, task.ID))
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)
func TestSweepTimedOutTasksHonorsRefundRolloutBoundary(t *testing.T) {
truncate(t)
sweepUnrefundedFailedTasks(context.Background())
const (
userID = 403
initialQuota = 10_000
legacyTaskQuota = 1_800
modernTaskQuota = 1_200
)
seedUser(t, userID, initialQuota)
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))
legacyTask := makeTask(userID, 0, legacyTaskQuota, 0, BillingSourceWallet, 0)
legacyTask.TaskID = "legacy_timeout_without_refund"
legacyTask.Progress = "50%"
legacyTask.SubmitTime = 1771718399 // 2026-02-21 23:59:59 UTC
require.NoError(t, model.DB.Create(legacyTask).Error)
modernTask := makeTask(userID, 0, modernTaskQuota, 0, BillingSourceWallet, 0)
modernTask.TaskID = "modern_timeout_with_refund"
modernTask.Progress = "50%"
modernTask.SubmitTime = 1771718400 // 2026-02-22 00:00:00 UTC
require.NoError(t, model.DB.Create(modernTask).Error)
previousTimeout := constant.TaskTimeoutMinutes
constant.TaskTimeoutMinutes = 1
t.Cleanup(func() { constant.TaskTimeoutMinutes = previousTimeout })
sweepTimedOutTasks(context.Background())
var reloadedLegacy model.Task
var reloadedModern model.Task
require.NoError(t, model.DB.First(&reloadedLegacy, legacyTask.ID).Error)
require.NoError(t, model.DB.First(&reloadedModern, modernTask.ID).Error)
assert.EqualValues(t, model.TaskStatusFailure, reloadedLegacy.Status)
assert.EqualValues(t, model.TaskStatusFailure, reloadedModern.Status)
assert.Zero(t, reloadedLegacy.Quota)
assert.Zero(t, reloadedModern.Quota)
assert.Contains(t, reloadedLegacy.FailReason, "旧系统遗留任务")
assert.Contains(t, reloadedModern.FailReason, "任务超时")
assert.Equal(t, initialQuota+modernTaskQuota, getUserQuota(t, userID))
assert.Equal(t, int64(1), countLogs(t))
}