From 01fa2b830972d4bfe29d9e32ee48820dcb63eb9c Mon Sep 17 00:00:00 2001 From: liangxu Date: Tue, 12 May 2026 02:24:57 +0800 Subject: [PATCH] fix quota reset reservation race --- server/internal/ops/repository/admin_user.go | 4 +- .../repository/admin_user_reset_sql_test.go | 77 ++++++++++++++++++ server/internal/tenant/app/ai_point_usage.go | 80 ++++++++++++------- .../tenant/app/article_imitation_service.go | 42 +++++----- .../tenant/app/generation_task_runtime.go | 47 ++++++----- .../tenant/app/kol_generation_service.go | 46 ++++++----- .../tenant/app/prompt_generate_service.go | 43 +++++----- .../internal/tenant/app/template_service.go | 43 +++++----- server/internal/tenant/compliance/service.go | 79 +++++++++++------- .../tenant/repository/generated/quota.sql.go | 4 +- .../repository/generated/quota_sql_test.go | 19 +++++ .../tenant/repository/queries/quota.sql | 4 +- .../internal/tenant/repository/quota_repo.go | 17 ++++ 13 files changed, 337 insertions(+), 168 deletions(-) create mode 100644 server/internal/ops/repository/admin_user_reset_sql_test.go create mode 100644 server/internal/tenant/repository/generated/quota_sql_test.go diff --git a/server/internal/ops/repository/admin_user.go b/server/internal/ops/repository/admin_user.go index c6ea88c..179f009 100644 --- a/server/internal/ops/repository/admin_user.go +++ b/server/internal/ops/repository/admin_user.go @@ -602,7 +602,7 @@ WITH reset AS ( updated_at = NOW() WHERE tenant_id = $1 AND quota_type = $2 - AND status = 'confirmed' + AND status IN ('pending', 'confirmed') AND created_at >= $3 AND created_at < $4 RETURNING reserved_amount @@ -627,7 +627,7 @@ WITH reset AS ( updated_at = NOW() WHERE tenant_id = $1 AND quota_type = 'ai_points' - AND status = 'confirmed' + AND status IN ('pending', 'confirmed') AND created_at >= $2 AND created_at < $3 RETURNING reserved_amount diff --git a/server/internal/ops/repository/admin_user_reset_sql_test.go b/server/internal/ops/repository/admin_user_reset_sql_test.go new file mode 100644 index 0000000..d2102fc --- /dev/null +++ b/server/internal/ops/repository/admin_user_reset_sql_test.go @@ -0,0 +1,77 @@ +package repository + +import ( + "context" + "errors" + "strings" + "testing" + "time" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgconn" +) + +var ( + testStart = time.Unix(0, 0).UTC() + testEnd = time.Unix(3600, 0).UTC() +) + +type captureResetSQLTx struct { + sql string +} + +func (tx *captureResetSQLTx) QueryRow(_ context.Context, sql string, _ ...any) pgx.Row { + tx.sql = sql + return failingRow{} +} + +func (tx *captureResetSQLTx) Begin(context.Context) (pgx.Tx, error) { return nil, errors.New("unused") } +func (tx *captureResetSQLTx) Commit(context.Context) error { return nil } +func (tx *captureResetSQLTx) Rollback(context.Context) error { return nil } +func (tx *captureResetSQLTx) CopyFrom(context.Context, pgx.Identifier, []string, pgx.CopyFromSource) (int64, error) { + return 0, errors.New("unused") +} +func (tx *captureResetSQLTx) SendBatch(context.Context, *pgx.Batch) pgx.BatchResults { + return nil +} +func (tx *captureResetSQLTx) LargeObjects() pgx.LargeObjects { return pgx.LargeObjects{} } +func (tx *captureResetSQLTx) Prepare(context.Context, string, string) (*pgconn.StatementDescription, error) { + return nil, errors.New("unused") +} +func (tx *captureResetSQLTx) Exec(context.Context, string, ...any) (pgconn.CommandTag, error) { + return pgconn.CommandTag{}, errors.New("unused") +} +func (tx *captureResetSQLTx) Query(context.Context, string, ...any) (pgx.Rows, error) { + return nil, errors.New("unused") +} +func (tx *captureResetSQLTx) Conn() *pgx.Conn { return nil } + +type failingRow struct{} + +func (failingRow) Scan(...any) error { return errors.New("stop after capturing sql") } + +func TestResetQuotaReservationsSQLMatchesEffectiveUsage(t *testing.T) { + tx := &captureResetSQLTx{} + + _, _, _ = resetQuotaReservations(context.Background(), tx, 1, "article_generation", testStart, testEnd) + + if !strings.Contains(tx.sql, "status IN ('pending', 'confirmed')") { + t.Fatalf("reset must clear every status counted as used; query:\n%s", tx.sql) + } + if strings.Contains(tx.sql, "ai_point_usage_logs") { + t.Fatalf("reset must not mutate AI point usage logs; query:\n%s", tx.sql) + } +} + +func TestResetAIPointReservationsSQLDoesNotTouchUsageLogs(t *testing.T) { + tx := &captureResetSQLTx{} + + _, _, _ = resetAIPointReservations(context.Background(), tx, 1, testStart, testEnd) + + if !strings.Contains(tx.sql, "status IN ('pending', 'confirmed')") { + t.Fatalf("reset must clear every status counted as used; query:\n%s", tx.sql) + } + if strings.Contains(tx.sql, "ai_point_usage_logs") { + t.Fatalf("reset must not mutate AI point usage logs; query:\n%s", tx.sql) + } +} diff --git a/server/internal/tenant/app/ai_point_usage.go b/server/internal/tenant/app/ai_point_usage.go index 8823700..cf5ef56 100644 --- a/server/internal/tenant/app/ai_point_usage.go +++ b/server/internal/tenant/app/ai_point_usage.go @@ -185,7 +185,7 @@ func CompleteAIPoints( if err != nil || !pending { return err } - if err := confirmAIPointReservation(ctx, tx, tenantID, reservation.ReservationID); err != nil { + if err := confirmAIPointReservationIfPending(ctx, tx, tenantID, reservation.ReservationID); err != nil { return err } var modelPtr *string @@ -226,29 +226,31 @@ func RefundAIPoints( return err } - status, err := repository.NewQuotaRepository(tx).GetAIQuotaStatus(ctx, tenantID, time.Now().UTC()) + refunded, err := refundAIPointReservationIfPending(ctx, tx, tenantID, reservation.ReservationID) if err != nil { return err } - - quotaRepo := repository.NewQuotaRepository(tx) - reason := "ai_points_refund" - referenceType := "ai_point_usage" - usageID := reservation.UsageLogID - if _, err := quotaRepo.InsertQuotaLedger(ctx, repository.QuotaLedgerInput{ - TenantID: tenantID, - OperatorID: operatorID, - QuotaType: aiPointsQuotaType, - Delta: reservation.Points, - BalanceAfter: status.Balance + reservation.Points, - Reason: &reason, - ReferenceType: &referenceType, - ReferenceID: &usageID, - }); err != nil { - return err - } - if err := refundAIPointReservation(ctx, tx, tenantID, reservation.ReservationID); err != nil { - return err + if refunded { + status, err := repository.NewQuotaRepository(tx).GetAIQuotaStatus(ctx, tenantID, time.Now().UTC()) + if err != nil { + return err + } + quotaRepo := repository.NewQuotaRepository(tx) + reason := "ai_points_refund" + referenceType := "ai_point_usage" + usageID := reservation.UsageLogID + if _, err := quotaRepo.InsertQuotaLedger(ctx, repository.QuotaLedgerInput{ + TenantID: tenantID, + OperatorID: operatorID, + QuotaType: aiPointsQuotaType, + Delta: reservation.Points, + BalanceAfter: status.Balance + reservation.Points, + Reason: &reason, + ReferenceType: &referenceType, + ReferenceID: &usageID, + }); err != nil { + return err + } } if err := repository.NewAIPointUsageRepository(tx).MarkRefunded(ctx, tenantID, reservation.UsageLogID, strings.TrimSpace(errorMessage)); err != nil { return err @@ -375,7 +377,7 @@ func lockPendingAIPointUsage(ctx context.Context, tx pgx.Tx, tenantID, usageLogI return status == "pending", nil } -func confirmAIPointReservation(ctx context.Context, tx pgx.Tx, tenantID, reservationID int64) error { +func confirmAIPointReservationIfPending(ctx context.Context, tx pgx.Tx, tenantID, reservationID int64) error { tag, err := tx.Exec(ctx, ` UPDATE quota_reservations SET status = 'confirmed', @@ -389,13 +391,10 @@ func confirmAIPointReservation(ctx context.Context, tx pgx.Tx, tenantID, reserva if err != nil { return err } - if tag.RowsAffected() == 0 { - return errors.New("ai points reservation is not pending") - } - return nil + return ignoreAIPointReservationReset(ctx, tx, tag.RowsAffected() > 0, tenantID, reservationID) } -func refundAIPointReservation(ctx context.Context, tx pgx.Tx, tenantID, reservationID int64) error { +func refundAIPointReservationIfPending(ctx context.Context, tx pgx.Tx, tenantID, reservationID int64) (bool, error) { tag, err := tx.Exec(ctx, ` UPDATE quota_reservations SET status = 'refunded', @@ -406,11 +405,32 @@ func refundAIPointReservation(ctx context.Context, tx pgx.Tx, tenantID, reservat AND quota_type = $3 AND status = 'pending' `, reservationID, tenantID, aiPointsQuotaType) + if err != nil { + return false, err + } + if tag.RowsAffected() > 0 { + return true, nil + } + return false, ignoreAIPointReservationReset(ctx, tx, false, tenantID, reservationID) +} + +func ignoreAIPointReservationReset(ctx context.Context, tx pgx.Tx, updated bool, tenantID, reservationID int64) error { + if updated { + return nil + } + var status string + err := tx.QueryRow(ctx, ` + SELECT status + FROM quota_reservations + WHERE id = $1 + AND tenant_id = $2 + AND quota_type = $3 + `, reservationID, tenantID, aiPointsQuotaType).Scan(&status) if err != nil { return err } - if tag.RowsAffected() == 0 { - return errors.New("ai points reservation is not pending") + if status == "reset" { + return nil } - return nil + return errors.New("ai points reservation is not pending") } diff --git a/server/internal/tenant/app/article_imitation_service.go b/server/internal/tenant/app/article_imitation_service.go index 115150d..4a833e9 100644 --- a/server/internal/tenant/app/article_imitation_service.go +++ b/server/internal/tenant/app/article_imitation_service.go @@ -512,27 +512,29 @@ func (s *ArticleImitationService) failGeneration(ctx context.Context, job articl if job.ReservationID > 0 { balance, balanceErr := quotaRepo.GetCurrentBalance(cleanupCtx, job.TenantID, "article_generation") - if balanceErr != nil { - cleanupErr = errors.Join(cleanupErr, fmt.Errorf("load quota balance for refund: %w", balanceErr)) - } else { - reason := "generation_refund" - referenceType := "generation_task" - taskReferenceID := job.TaskID - if _, err := quotaRepo.InsertQuotaLedger(cleanupCtx, repository.QuotaLedgerInput{ - TenantID: job.TenantID, - OperatorID: job.OperatorID, - QuotaType: "article_generation", - Delta: 1, - BalanceAfter: balance + 1, - Reason: &reason, - ReferenceType: &referenceType, - ReferenceID: &taskReferenceID, - }); err != nil { - cleanupErr = errors.Join(cleanupErr, fmt.Errorf("insert generation refund ledger: %w", err)) - } - } - if err := quotaRepo.RefundReservation(cleanupCtx, job.ReservationID, job.TenantID); err != nil { + refunded, err := quotaRepo.RefundPendingReservation(cleanupCtx, job.ReservationID, job.TenantID) + if err != nil { cleanupErr = errors.Join(cleanupErr, fmt.Errorf("refund generation reservation: %w", err)) + } else if refunded { + if balanceErr != nil { + cleanupErr = errors.Join(cleanupErr, fmt.Errorf("load quota balance for refund: %w", balanceErr)) + } else { + reason := "generation_refund" + referenceType := "generation_task" + taskReferenceID := job.TaskID + if _, err := quotaRepo.InsertQuotaLedger(cleanupCtx, repository.QuotaLedgerInput{ + TenantID: job.TenantID, + OperatorID: job.OperatorID, + QuotaType: "article_generation", + Delta: 1, + BalanceAfter: balance + 1, + Reason: &reason, + ReferenceType: &referenceType, + ReferenceID: &taskReferenceID, + }); err != nil { + cleanupErr = errors.Join(cleanupErr, fmt.Errorf("insert generation refund ledger: %w", err)) + } + } } } diff --git a/server/internal/tenant/app/generation_task_runtime.go b/server/internal/tenant/app/generation_task_runtime.go index c742fcd..fe441a1 100644 --- a/server/internal/tenant/app/generation_task_runtime.go +++ b/server/internal/tenant/app/generation_task_runtime.go @@ -79,29 +79,34 @@ func HandleClaimedGenerationTaskFailure( } if task.QuotaReservationID > 0 { - balance, balanceErr := quotaRepo.GetCurrentBalance(cleanupCtx, task.TenantID, "article_generation") - if balanceErr != nil { - failureErr = errors.Join(failureErr, fmt.Errorf("load quota balance for refund: %w", balanceErr)) - } else if task.OperatorID > 0 { - reason := "generation_refund" - referenceType := "generation_task" - taskReferenceID := task.ID - if _, err := quotaRepo.InsertQuotaLedger(cleanupCtx, repository.QuotaLedgerInput{ - TenantID: task.TenantID, - OperatorID: task.OperatorID, - QuotaType: "article_generation", - Delta: 1, - BalanceAfter: balance + 1, - Reason: &reason, - ReferenceType: &referenceType, - ReferenceID: &taskReferenceID, - }); err != nil { - failureErr = errors.Join(failureErr, fmt.Errorf("insert generation refund ledger: %w", err)) - } + var balance int + balanceErr := error(nil) + if task.OperatorID > 0 { + balance, balanceErr = quotaRepo.GetCurrentBalance(cleanupCtx, task.TenantID, "article_generation") } - - if err := quotaRepo.RefundReservation(cleanupCtx, task.QuotaReservationID, task.TenantID); err != nil { + refunded, err := quotaRepo.RefundPendingReservation(cleanupCtx, task.QuotaReservationID, task.TenantID) + if err != nil { failureErr = errors.Join(failureErr, fmt.Errorf("refund generation reservation: %w", err)) + } else if refunded { + if balanceErr != nil { + failureErr = errors.Join(failureErr, fmt.Errorf("load quota balance for refund: %w", balanceErr)) + } else if task.OperatorID > 0 { + reason := "generation_refund" + referenceType := "generation_task" + taskReferenceID := task.ID + if _, err := quotaRepo.InsertQuotaLedger(cleanupCtx, repository.QuotaLedgerInput{ + TenantID: task.TenantID, + OperatorID: task.OperatorID, + QuotaType: "article_generation", + Delta: 1, + BalanceAfter: balance + 1, + Reason: &reason, + ReferenceType: &referenceType, + ReferenceID: &taskReferenceID, + }); err != nil { + failureErr = errors.Join(failureErr, fmt.Errorf("insert generation refund ledger: %w", err)) + } + } } } diff --git a/server/internal/tenant/app/kol_generation_service.go b/server/internal/tenant/app/kol_generation_service.go index 7faca85..0879af3 100644 --- a/server/internal/tenant/app/kol_generation_service.go +++ b/server/internal/tenant/app/kol_generation_service.go @@ -445,28 +445,34 @@ func HandleKolGenerationTaskFailure( } if task.QuotaReservationID > 0 { - balance, balanceErr := quotaRepo.GetCurrentBalance(cleanupCtx, task.TenantID, "article_generation") - if balanceErr != nil { - failureErr = errors.Join(failureErr, fmt.Errorf("load quota balance for refund: %w", balanceErr)) - } else if task.OperatorID > 0 { - reason := "generation_refund" - referenceType := "generation_task" - taskReferenceID := task.ID - if _, err := quotaRepo.InsertQuotaLedger(cleanupCtx, repository.QuotaLedgerInput{ - TenantID: task.TenantID, - OperatorID: task.OperatorID, - QuotaType: "article_generation", - Delta: 1, - BalanceAfter: balance + 1, - Reason: &reason, - ReferenceType: &referenceType, - ReferenceID: &taskReferenceID, - }); err != nil { - failureErr = errors.Join(failureErr, fmt.Errorf("insert generation refund ledger: %w", err)) - } + var balance int + balanceErr := error(nil) + if task.OperatorID > 0 { + balance, balanceErr = quotaRepo.GetCurrentBalance(cleanupCtx, task.TenantID, "article_generation") } - if err := quotaRepo.RefundReservation(cleanupCtx, task.QuotaReservationID, task.TenantID); err != nil { + refunded, err := quotaRepo.RefundPendingReservation(cleanupCtx, task.QuotaReservationID, task.TenantID) + if err != nil { failureErr = errors.Join(failureErr, fmt.Errorf("refund generation reservation: %w", err)) + } else if refunded { + if balanceErr != nil { + failureErr = errors.Join(failureErr, fmt.Errorf("load quota balance for refund: %w", balanceErr)) + } else if task.OperatorID > 0 { + reason := "generation_refund" + referenceType := "generation_task" + taskReferenceID := task.ID + if _, err := quotaRepo.InsertQuotaLedger(cleanupCtx, repository.QuotaLedgerInput{ + TenantID: task.TenantID, + OperatorID: task.OperatorID, + QuotaType: "article_generation", + Delta: 1, + BalanceAfter: balance + 1, + Reason: &reason, + ReferenceType: &referenceType, + ReferenceID: &taskReferenceID, + }); err != nil { + failureErr = errors.Join(failureErr, fmt.Errorf("insert generation refund ledger: %w", err)) + } + } } } diff --git a/server/internal/tenant/app/prompt_generate_service.go b/server/internal/tenant/app/prompt_generate_service.go index 54991d2..b398bb5 100644 --- a/server/internal/tenant/app/prompt_generate_service.go +++ b/server/internal/tenant/app/prompt_generate_service.go @@ -739,28 +739,29 @@ func (s *PromptRuleGenerationService) failGeneration(ctx context.Context, job pr } balance, balanceErr := quotaRepo.GetCurrentBalance(cleanupCtx, job.TenantID, "article_generation") - if balanceErr != nil { - cleanupErr = errors.Join(cleanupErr, fmt.Errorf("load quota balance for refund: %w", balanceErr)) - } else { - reason := "generation_refund" - referenceType := "generation_task" - taskReferenceID := job.TaskID - if _, err := quotaRepo.InsertQuotaLedger(cleanupCtx, repository.QuotaLedgerInput{ - TenantID: job.TenantID, - OperatorID: job.OperatorID, - QuotaType: "article_generation", - Delta: 1, - BalanceAfter: balance + 1, - Reason: &reason, - ReferenceType: &referenceType, - ReferenceID: &taskReferenceID, - }); err != nil { - cleanupErr = errors.Join(cleanupErr, fmt.Errorf("insert generation refund ledger: %w", err)) - } - } - - if err := quotaRepo.RefundReservation(cleanupCtx, job.ReservationID, job.TenantID); err != nil { + refunded, err := quotaRepo.RefundPendingReservation(cleanupCtx, job.ReservationID, job.TenantID) + if err != nil { cleanupErr = errors.Join(cleanupErr, fmt.Errorf("refund generation reservation: %w", err)) + } else if refunded { + if balanceErr != nil { + cleanupErr = errors.Join(cleanupErr, fmt.Errorf("load quota balance for refund: %w", balanceErr)) + } else { + reason := "generation_refund" + referenceType := "generation_task" + taskReferenceID := job.TaskID + if _, err := quotaRepo.InsertQuotaLedger(cleanupCtx, repository.QuotaLedgerInput{ + TenantID: job.TenantID, + OperatorID: job.OperatorID, + QuotaType: "article_generation", + Delta: 1, + BalanceAfter: balance + 1, + Reason: &reason, + ReferenceType: &referenceType, + ReferenceID: &taskReferenceID, + }); err != nil { + cleanupErr = errors.Join(cleanupErr, fmt.Errorf("insert generation refund ledger: %w", err)) + } + } } invalidateArticleCaches(cleanupCtx, s.cache, job.TenantID, &job.ArticleID) if cleanupErr != nil { diff --git a/server/internal/tenant/app/template_service.go b/server/internal/tenant/app/template_service.go index af0e4e8..3317a21 100644 --- a/server/internal/tenant/app/template_service.go +++ b/server/internal/tenant/app/template_service.go @@ -626,28 +626,29 @@ func (s *TemplateService) failGeneration(ctx context.Context, job generationJob, } balance, balanceErr := quotaRepo.GetCurrentBalance(cleanupCtx, job.TenantID, "article_generation") - if balanceErr != nil { - cleanupErr = errors.Join(cleanupErr, fmt.Errorf("load quota balance for refund: %w", balanceErr)) - } else { - reason := "generation_refund" - referenceType := "generation_task" - taskReferenceID := job.TaskID - if _, err := quotaRepo.InsertQuotaLedger(cleanupCtx, repository.QuotaLedgerInput{ - TenantID: job.TenantID, - OperatorID: job.OperatorID, - QuotaType: "article_generation", - Delta: 1, - BalanceAfter: balance + 1, - Reason: &reason, - ReferenceType: &referenceType, - ReferenceID: &taskReferenceID, - }); err != nil { - cleanupErr = errors.Join(cleanupErr, fmt.Errorf("insert generation refund ledger: %w", err)) - } - } - - if err := quotaRepo.RefundReservation(cleanupCtx, job.ReservationID, job.TenantID); err != nil { + refunded, err := quotaRepo.RefundPendingReservation(cleanupCtx, job.ReservationID, job.TenantID) + if err != nil { cleanupErr = errors.Join(cleanupErr, fmt.Errorf("refund generation reservation: %w", err)) + } else if refunded { + if balanceErr != nil { + cleanupErr = errors.Join(cleanupErr, fmt.Errorf("load quota balance for refund: %w", balanceErr)) + } else { + reason := "generation_refund" + referenceType := "generation_task" + taskReferenceID := job.TaskID + if _, err := quotaRepo.InsertQuotaLedger(cleanupCtx, repository.QuotaLedgerInput{ + TenantID: job.TenantID, + OperatorID: job.OperatorID, + QuotaType: "article_generation", + Delta: 1, + BalanceAfter: balance + 1, + Reason: &reason, + ReferenceType: &referenceType, + ReferenceID: &taskReferenceID, + }); err != nil { + cleanupErr = errors.Join(cleanupErr, fmt.Errorf("insert generation refund ledger: %w", err)) + } + } } invalidateArticleCaches(cleanupCtx, s.cache, job.TenantID, &job.ArticleID) if cleanupErr != nil { diff --git a/server/internal/tenant/compliance/service.go b/server/internal/tenant/compliance/service.go index df64ecf..8ef67b4 100644 --- a/server/internal/tenant/compliance/service.go +++ b/server/internal/tenant/compliance/service.go @@ -1417,7 +1417,7 @@ func (s *Service) completeComplianceLLMJudgePoints(ctx context.Context, job revi if err != nil || !pending { return err } - if err := confirmAIPointReservation(ctx, tx, job.TenantID, reservation.ReservationID); err != nil { + if err := confirmAIPointReservationIfPending(ctx, tx, job.TenantID, reservation.ReservationID); err != nil { return err } var modelPtr *string @@ -1447,28 +1447,31 @@ func (s *Service) refundComplianceLLMJudgePoints(ctx context.Context, job review if err != nil || !pending { return err } - status, err := tenantrepo.NewQuotaRepository(tx).GetAIQuotaStatus(ctx, job.TenantID, time.Now().UTC()) + refunded, err := refundAIPointReservationIfPending(ctx, tx, job.TenantID, reservation.ReservationID) if err != nil { return err } - quotaRepo := tenantrepo.NewQuotaRepository(tx) - reason := "ai_points_refund" - referenceType := "ai_point_usage" - usageID := reservation.UsageLogID - if _, err := quotaRepo.InsertQuotaLedger(ctx, tenantrepo.QuotaLedgerInput{ - TenantID: job.TenantID, - OperatorID: job.CheckedBy, - QuotaType: aiPointsQuotaType, - Delta: reservation.Points, - BalanceAfter: status.Balance + reservation.Points, - Reason: &reason, - ReferenceType: &referenceType, - ReferenceID: &usageID, - }); err != nil { - return err - } - if err := refundAIPointReservation(ctx, tx, job.TenantID, reservation.ReservationID); err != nil { - return err + if refunded { + status, err := tenantrepo.NewQuotaRepository(tx).GetAIQuotaStatus(ctx, job.TenantID, time.Now().UTC()) + if err != nil { + return err + } + quotaRepo := tenantrepo.NewQuotaRepository(tx) + reason := "ai_points_refund" + referenceType := "ai_point_usage" + usageID := reservation.UsageLogID + if _, err := quotaRepo.InsertQuotaLedger(ctx, tenantrepo.QuotaLedgerInput{ + TenantID: job.TenantID, + OperatorID: job.CheckedBy, + QuotaType: aiPointsQuotaType, + Delta: reservation.Points, + BalanceAfter: status.Balance + reservation.Points, + Reason: &reason, + ReferenceType: &referenceType, + ReferenceID: &usageID, + }); err != nil { + return err + } } if err := tenantrepo.NewAIPointUsageRepository(tx).MarkRefunded(ctx, job.TenantID, reservation.UsageLogID, truncateRunes(errorString(jobErr), 2000)); err != nil { return err @@ -1512,7 +1515,7 @@ func lockPendingAIPointUsage(ctx context.Context, tx pgx.Tx, tenantID, usageLogI return status == "pending", nil } -func confirmAIPointReservation(ctx context.Context, tx pgx.Tx, tenantID, reservationID int64) error { +func confirmAIPointReservationIfPending(ctx context.Context, tx pgx.Tx, tenantID, reservationID int64) error { tag, err := tx.Exec(ctx, ` UPDATE quota_reservations SET status = 'confirmed', @@ -1526,13 +1529,10 @@ func confirmAIPointReservation(ctx context.Context, tx pgx.Tx, tenantID, reserva if err != nil { return err } - if tag.RowsAffected() == 0 { - return errors.New("ai points reservation is not pending") - } - return nil + return ignoreAIPointReservationReset(ctx, tx, tag.RowsAffected() > 0, tenantID, reservationID) } -func refundAIPointReservation(ctx context.Context, tx pgx.Tx, tenantID, reservationID int64) error { +func refundAIPointReservationIfPending(ctx context.Context, tx pgx.Tx, tenantID, reservationID int64) (bool, error) { tag, err := tx.Exec(ctx, ` UPDATE quota_reservations SET status = 'refunded', @@ -1543,13 +1543,34 @@ func refundAIPointReservation(ctx context.Context, tx pgx.Tx, tenantID, reservat AND quota_type = $3 AND status = 'pending' `, reservationID, tenantID, aiPointsQuotaType) + if err != nil { + return false, err + } + if tag.RowsAffected() > 0 { + return true, nil + } + return false, ignoreAIPointReservationReset(ctx, tx, false, tenantID, reservationID) +} + +func ignoreAIPointReservationReset(ctx context.Context, tx pgx.Tx, updated bool, tenantID, reservationID int64) error { + if updated { + return nil + } + var status string + err := tx.QueryRow(ctx, ` + SELECT status + FROM quota_reservations + WHERE id = $1 + AND tenant_id = $2 + AND quota_type = $3 + `, reservationID, tenantID, aiPointsQuotaType).Scan(&status) if err != nil { return err } - if tag.RowsAffected() == 0 { - return errors.New("ai points reservation is not pending") + if status == "reset" { + return nil } - return nil + return errors.New("ai points reservation is not pending") } func (s *Service) markReviewJobDone(ctx context.Context, jobID int64) error { diff --git a/server/internal/tenant/repository/generated/quota.sql.go b/server/internal/tenant/repository/generated/quota.sql.go index debcc08..f7ab64f 100644 --- a/server/internal/tenant/repository/generated/quota.sql.go +++ b/server/internal/tenant/repository/generated/quota.sql.go @@ -13,7 +13,7 @@ import ( const confirmReservation = `-- name: ConfirmReservation :exec UPDATE quota_reservations SET status = 'confirmed', consumed_amount = reserved_amount, updated_at = NOW() -WHERE id = $1 AND tenant_id = $2 +WHERE id = $1 AND tenant_id = $2 AND status = 'pending' ` type ConfirmReservationParams struct { @@ -114,7 +114,7 @@ func (q *Queries) InsertQuotaLedger(ctx context.Context, arg InsertQuotaLedgerPa const refundReservation = `-- name: RefundReservation :exec UPDATE quota_reservations SET status = 'refunded', refunded_amount = reserved_amount, updated_at = NOW() -WHERE id = $1 AND tenant_id = $2 +WHERE id = $1 AND tenant_id = $2 AND status = 'pending' ` type RefundReservationParams struct { diff --git a/server/internal/tenant/repository/generated/quota_sql_test.go b/server/internal/tenant/repository/generated/quota_sql_test.go new file mode 100644 index 0000000..5af2afc --- /dev/null +++ b/server/internal/tenant/repository/generated/quota_sql_test.go @@ -0,0 +1,19 @@ +package generated + +import ( + "strings" + "testing" +) + +func TestQuotaReservationTransitionsOnlyUpdatePending(t *testing.T) { + for name, query := range map[string]string{ + "confirm": confirmReservation, + "refund": refundReservation, + } { + t.Run(name, func(t *testing.T) { + if !strings.Contains(query, "status = 'pending'") { + t.Fatalf("%s reservation query must only update pending reservations; query:\n%s", name, query) + } + }) + } +} diff --git a/server/internal/tenant/repository/queries/quota.sql b/server/internal/tenant/repository/queries/quota.sql index 612c0e3..40eed0f 100644 --- a/server/internal/tenant/repository/queries/quota.sql +++ b/server/internal/tenant/repository/queries/quota.sql @@ -23,8 +23,8 @@ WHERE id = sqlc.arg(id) AND tenant_id = sqlc.arg(tenant_id); -- name: ConfirmReservation :exec UPDATE quota_reservations SET status = 'confirmed', consumed_amount = reserved_amount, updated_at = NOW() -WHERE id = sqlc.arg(id) AND tenant_id = sqlc.arg(tenant_id); +WHERE id = sqlc.arg(id) AND tenant_id = sqlc.arg(tenant_id) AND status = 'pending'; -- name: RefundReservation :exec UPDATE quota_reservations SET status = 'refunded', refunded_amount = reserved_amount, updated_at = NOW() -WHERE id = sqlc.arg(id) AND tenant_id = sqlc.arg(tenant_id); +WHERE id = sqlc.arg(id) AND tenant_id = sqlc.arg(tenant_id) AND status = 'pending'; diff --git a/server/internal/tenant/repository/quota_repo.go b/server/internal/tenant/repository/quota_repo.go index 2fc60a5..6a305ef 100644 --- a/server/internal/tenant/repository/quota_repo.go +++ b/server/internal/tenant/repository/quota_repo.go @@ -59,6 +59,7 @@ type QuotaRepository interface { UpdateQuotaReservationResource(ctx context.Context, reservationID, tenantID, resourceID int64) error ConfirmReservation(ctx context.Context, reservationID, tenantID int64) error RefundReservation(ctx context.Context, reservationID, tenantID int64) error + RefundPendingReservation(ctx context.Context, reservationID, tenantID int64) (bool, error) } type quotaRepository struct { @@ -256,3 +257,19 @@ func (r *quotaRepository) RefundReservation(ctx context.Context, reservationID, TenantID: tenantID, }) } + +func (r *quotaRepository) RefundPendingReservation(ctx context.Context, reservationID, tenantID int64) (bool, error) { + tag, err := r.db.Exec(ctx, ` +UPDATE quota_reservations +SET status = 'refunded', + refunded_amount = reserved_amount, + updated_at = NOW() +WHERE id = $1 + AND tenant_id = $2 + AND status = 'pending' +`, reservationID, tenantID) + if err != nil { + return false, err + } + return tag.RowsAffected() > 0, nil +}