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 <i@caion.me>
2026-07-20 22:03:13 +08:00
parent 4aa08f917e
commit e0d5156115
7 changed files with 456 additions and 16 deletions
@@ -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
@@ -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())
+}