fix quota reset reservation race
Desktop Client Build / Resolve Build Metadata (push) Successful in 28s
Frontend CI / Frontend (push) Successful in 3m57s
Backend CI / Backend (push) Successful in 17m3s
Desktop Client Build / Build Desktop Client (push) Successful in 26m8s
Desktop Client Build / Publish Client Artifacts to NAS (push) Successful in 51s
Desktop Client Build / Resolve Build Metadata (push) Successful in 28s
Frontend CI / Frontend (push) Successful in 3m57s
Backend CI / Backend (push) Successful in 17m3s
Desktop Client Build / Build Desktop Client (push) Successful in 26m8s
Desktop Client Build / Publish Client Artifacts to NAS (push) Successful in 51s
This commit is contained in:
@@ -602,7 +602,7 @@ WITH reset AS (
|
|||||||
updated_at = NOW()
|
updated_at = NOW()
|
||||||
WHERE tenant_id = $1
|
WHERE tenant_id = $1
|
||||||
AND quota_type = $2
|
AND quota_type = $2
|
||||||
AND status = 'confirmed'
|
AND status IN ('pending', 'confirmed')
|
||||||
AND created_at >= $3
|
AND created_at >= $3
|
||||||
AND created_at < $4
|
AND created_at < $4
|
||||||
RETURNING reserved_amount
|
RETURNING reserved_amount
|
||||||
@@ -627,7 +627,7 @@ WITH reset AS (
|
|||||||
updated_at = NOW()
|
updated_at = NOW()
|
||||||
WHERE tenant_id = $1
|
WHERE tenant_id = $1
|
||||||
AND quota_type = 'ai_points'
|
AND quota_type = 'ai_points'
|
||||||
AND status = 'confirmed'
|
AND status IN ('pending', 'confirmed')
|
||||||
AND created_at >= $2
|
AND created_at >= $2
|
||||||
AND created_at < $3
|
AND created_at < $3
|
||||||
RETURNING reserved_amount
|
RETURNING reserved_amount
|
||||||
|
|||||||
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -185,7 +185,7 @@ func CompleteAIPoints(
|
|||||||
if err != nil || !pending {
|
if err != nil || !pending {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
if err := confirmAIPointReservation(ctx, tx, tenantID, reservation.ReservationID); err != nil {
|
if err := confirmAIPointReservationIfPending(ctx, tx, tenantID, reservation.ReservationID); err != nil {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
var modelPtr *string
|
var modelPtr *string
|
||||||
@@ -226,29 +226,31 @@ func RefundAIPoints(
|
|||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
|
|
||||||
status, err := repository.NewQuotaRepository(tx).GetAIQuotaStatus(ctx, tenantID, time.Now().UTC())
|
refunded, err := refundAIPointReservationIfPending(ctx, tx, tenantID, reservation.ReservationID)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
|
if refunded {
|
||||||
quotaRepo := repository.NewQuotaRepository(tx)
|
status, err := repository.NewQuotaRepository(tx).GetAIQuotaStatus(ctx, tenantID, time.Now().UTC())
|
||||||
reason := "ai_points_refund"
|
if err != nil {
|
||||||
referenceType := "ai_point_usage"
|
return err
|
||||||
usageID := reservation.UsageLogID
|
}
|
||||||
if _, err := quotaRepo.InsertQuotaLedger(ctx, repository.QuotaLedgerInput{
|
quotaRepo := repository.NewQuotaRepository(tx)
|
||||||
TenantID: tenantID,
|
reason := "ai_points_refund"
|
||||||
OperatorID: operatorID,
|
referenceType := "ai_point_usage"
|
||||||
QuotaType: aiPointsQuotaType,
|
usageID := reservation.UsageLogID
|
||||||
Delta: reservation.Points,
|
if _, err := quotaRepo.InsertQuotaLedger(ctx, repository.QuotaLedgerInput{
|
||||||
BalanceAfter: status.Balance + reservation.Points,
|
TenantID: tenantID,
|
||||||
Reason: &reason,
|
OperatorID: operatorID,
|
||||||
ReferenceType: &referenceType,
|
QuotaType: aiPointsQuotaType,
|
||||||
ReferenceID: &usageID,
|
Delta: reservation.Points,
|
||||||
}); err != nil {
|
BalanceAfter: status.Balance + reservation.Points,
|
||||||
return err
|
Reason: &reason,
|
||||||
}
|
ReferenceType: &referenceType,
|
||||||
if err := refundAIPointReservation(ctx, tx, tenantID, reservation.ReservationID); err != nil {
|
ReferenceID: &usageID,
|
||||||
return err
|
}); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
}
|
}
|
||||||
if err := repository.NewAIPointUsageRepository(tx).MarkRefunded(ctx, tenantID, reservation.UsageLogID, strings.TrimSpace(errorMessage)); err != nil {
|
if err := repository.NewAIPointUsageRepository(tx).MarkRefunded(ctx, tenantID, reservation.UsageLogID, strings.TrimSpace(errorMessage)); err != nil {
|
||||||
return err
|
return err
|
||||||
@@ -375,7 +377,7 @@ func lockPendingAIPointUsage(ctx context.Context, tx pgx.Tx, tenantID, usageLogI
|
|||||||
return status == "pending", nil
|
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, `
|
tag, err := tx.Exec(ctx, `
|
||||||
UPDATE quota_reservations
|
UPDATE quota_reservations
|
||||||
SET status = 'confirmed',
|
SET status = 'confirmed',
|
||||||
@@ -389,13 +391,10 @@ func confirmAIPointReservation(ctx context.Context, tx pgx.Tx, tenantID, reserva
|
|||||||
if err != nil {
|
if err != nil {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
if tag.RowsAffected() == 0 {
|
return ignoreAIPointReservationReset(ctx, tx, tag.RowsAffected() > 0, tenantID, reservationID)
|
||||||
return errors.New("ai points reservation is not pending")
|
|
||||||
}
|
|
||||||
return nil
|
|
||||||
}
|
}
|
||||||
|
|
||||||
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, `
|
tag, err := tx.Exec(ctx, `
|
||||||
UPDATE quota_reservations
|
UPDATE quota_reservations
|
||||||
SET status = 'refunded',
|
SET status = 'refunded',
|
||||||
@@ -406,11 +405,32 @@ func refundAIPointReservation(ctx context.Context, tx pgx.Tx, tenantID, reservat
|
|||||||
AND quota_type = $3
|
AND quota_type = $3
|
||||||
AND status = 'pending'
|
AND status = 'pending'
|
||||||
`, reservationID, tenantID, aiPointsQuotaType)
|
`, 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 {
|
if err != nil {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
if tag.RowsAffected() == 0 {
|
if status == "reset" {
|
||||||
return errors.New("ai points reservation is not pending")
|
return nil
|
||||||
}
|
}
|
||||||
return nil
|
return errors.New("ai points reservation is not pending")
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -512,27 +512,29 @@ func (s *ArticleImitationService) failGeneration(ctx context.Context, job articl
|
|||||||
|
|
||||||
if job.ReservationID > 0 {
|
if job.ReservationID > 0 {
|
||||||
balance, balanceErr := quotaRepo.GetCurrentBalance(cleanupCtx, job.TenantID, "article_generation")
|
balance, balanceErr := quotaRepo.GetCurrentBalance(cleanupCtx, job.TenantID, "article_generation")
|
||||||
if balanceErr != nil {
|
refunded, err := quotaRepo.RefundPendingReservation(cleanupCtx, job.ReservationID, job.TenantID)
|
||||||
cleanupErr = errors.Join(cleanupErr, fmt.Errorf("load quota balance for refund: %w", balanceErr))
|
if err != nil {
|
||||||
} 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 {
|
|
||||||
cleanupErr = errors.Join(cleanupErr, fmt.Errorf("refund generation reservation: %w", err))
|
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))
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -79,29 +79,34 @@ func HandleClaimedGenerationTaskFailure(
|
|||||||
}
|
}
|
||||||
|
|
||||||
if task.QuotaReservationID > 0 {
|
if task.QuotaReservationID > 0 {
|
||||||
balance, balanceErr := quotaRepo.GetCurrentBalance(cleanupCtx, task.TenantID, "article_generation")
|
var balance int
|
||||||
if balanceErr != nil {
|
balanceErr := error(nil)
|
||||||
failureErr = errors.Join(failureErr, fmt.Errorf("load quota balance for refund: %w", balanceErr))
|
if task.OperatorID > 0 {
|
||||||
} else if task.OperatorID > 0 {
|
balance, balanceErr = quotaRepo.GetCurrentBalance(cleanupCtx, task.TenantID, "article_generation")
|
||||||
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))
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
refunded, err := quotaRepo.RefundPendingReservation(cleanupCtx, task.QuotaReservationID, task.TenantID)
|
||||||
if err := quotaRepo.RefundReservation(cleanupCtx, task.QuotaReservationID, task.TenantID); err != nil {
|
if err != nil {
|
||||||
failureErr = errors.Join(failureErr, fmt.Errorf("refund generation reservation: %w", err))
|
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))
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -445,28 +445,34 @@ func HandleKolGenerationTaskFailure(
|
|||||||
}
|
}
|
||||||
|
|
||||||
if task.QuotaReservationID > 0 {
|
if task.QuotaReservationID > 0 {
|
||||||
balance, balanceErr := quotaRepo.GetCurrentBalance(cleanupCtx, task.TenantID, "article_generation")
|
var balance int
|
||||||
if balanceErr != nil {
|
balanceErr := error(nil)
|
||||||
failureErr = errors.Join(failureErr, fmt.Errorf("load quota balance for refund: %w", balanceErr))
|
if task.OperatorID > 0 {
|
||||||
} else if task.OperatorID > 0 {
|
balance, balanceErr = quotaRepo.GetCurrentBalance(cleanupCtx, task.TenantID, "article_generation")
|
||||||
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))
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
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))
|
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))
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -739,28 +739,29 @@ func (s *PromptRuleGenerationService) failGeneration(ctx context.Context, job pr
|
|||||||
}
|
}
|
||||||
|
|
||||||
balance, balanceErr := quotaRepo.GetCurrentBalance(cleanupCtx, job.TenantID, "article_generation")
|
balance, balanceErr := quotaRepo.GetCurrentBalance(cleanupCtx, job.TenantID, "article_generation")
|
||||||
if balanceErr != nil {
|
refunded, err := quotaRepo.RefundPendingReservation(cleanupCtx, job.ReservationID, job.TenantID)
|
||||||
cleanupErr = errors.Join(cleanupErr, fmt.Errorf("load quota balance for refund: %w", balanceErr))
|
if err != nil {
|
||||||
} 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 {
|
|
||||||
cleanupErr = errors.Join(cleanupErr, fmt.Errorf("refund generation reservation: %w", err))
|
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)
|
invalidateArticleCaches(cleanupCtx, s.cache, job.TenantID, &job.ArticleID)
|
||||||
if cleanupErr != nil {
|
if cleanupErr != nil {
|
||||||
|
|||||||
@@ -626,28 +626,29 @@ func (s *TemplateService) failGeneration(ctx context.Context, job generationJob,
|
|||||||
}
|
}
|
||||||
|
|
||||||
balance, balanceErr := quotaRepo.GetCurrentBalance(cleanupCtx, job.TenantID, "article_generation")
|
balance, balanceErr := quotaRepo.GetCurrentBalance(cleanupCtx, job.TenantID, "article_generation")
|
||||||
if balanceErr != nil {
|
refunded, err := quotaRepo.RefundPendingReservation(cleanupCtx, job.ReservationID, job.TenantID)
|
||||||
cleanupErr = errors.Join(cleanupErr, fmt.Errorf("load quota balance for refund: %w", balanceErr))
|
if err != nil {
|
||||||
} 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 {
|
|
||||||
cleanupErr = errors.Join(cleanupErr, fmt.Errorf("refund generation reservation: %w", err))
|
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)
|
invalidateArticleCaches(cleanupCtx, s.cache, job.TenantID, &job.ArticleID)
|
||||||
if cleanupErr != nil {
|
if cleanupErr != nil {
|
||||||
|
|||||||
@@ -1417,7 +1417,7 @@ func (s *Service) completeComplianceLLMJudgePoints(ctx context.Context, job revi
|
|||||||
if err != nil || !pending {
|
if err != nil || !pending {
|
||||||
return err
|
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
|
return err
|
||||||
}
|
}
|
||||||
var modelPtr *string
|
var modelPtr *string
|
||||||
@@ -1447,28 +1447,31 @@ func (s *Service) refundComplianceLLMJudgePoints(ctx context.Context, job review
|
|||||||
if err != nil || !pending {
|
if err != nil || !pending {
|
||||||
return err
|
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 {
|
if err != nil {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
quotaRepo := tenantrepo.NewQuotaRepository(tx)
|
if refunded {
|
||||||
reason := "ai_points_refund"
|
status, err := tenantrepo.NewQuotaRepository(tx).GetAIQuotaStatus(ctx, job.TenantID, time.Now().UTC())
|
||||||
referenceType := "ai_point_usage"
|
if err != nil {
|
||||||
usageID := reservation.UsageLogID
|
return err
|
||||||
if _, err := quotaRepo.InsertQuotaLedger(ctx, tenantrepo.QuotaLedgerInput{
|
}
|
||||||
TenantID: job.TenantID,
|
quotaRepo := tenantrepo.NewQuotaRepository(tx)
|
||||||
OperatorID: job.CheckedBy,
|
reason := "ai_points_refund"
|
||||||
QuotaType: aiPointsQuotaType,
|
referenceType := "ai_point_usage"
|
||||||
Delta: reservation.Points,
|
usageID := reservation.UsageLogID
|
||||||
BalanceAfter: status.Balance + reservation.Points,
|
if _, err := quotaRepo.InsertQuotaLedger(ctx, tenantrepo.QuotaLedgerInput{
|
||||||
Reason: &reason,
|
TenantID: job.TenantID,
|
||||||
ReferenceType: &referenceType,
|
OperatorID: job.CheckedBy,
|
||||||
ReferenceID: &usageID,
|
QuotaType: aiPointsQuotaType,
|
||||||
}); err != nil {
|
Delta: reservation.Points,
|
||||||
return err
|
BalanceAfter: status.Balance + reservation.Points,
|
||||||
}
|
Reason: &reason,
|
||||||
if err := refundAIPointReservation(ctx, tx, job.TenantID, reservation.ReservationID); err != nil {
|
ReferenceType: &referenceType,
|
||||||
return err
|
ReferenceID: &usageID,
|
||||||
|
}); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
}
|
}
|
||||||
if err := tenantrepo.NewAIPointUsageRepository(tx).MarkRefunded(ctx, job.TenantID, reservation.UsageLogID, truncateRunes(errorString(jobErr), 2000)); err != nil {
|
if err := tenantrepo.NewAIPointUsageRepository(tx).MarkRefunded(ctx, job.TenantID, reservation.UsageLogID, truncateRunes(errorString(jobErr), 2000)); err != nil {
|
||||||
return err
|
return err
|
||||||
@@ -1512,7 +1515,7 @@ func lockPendingAIPointUsage(ctx context.Context, tx pgx.Tx, tenantID, usageLogI
|
|||||||
return status == "pending", nil
|
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, `
|
tag, err := tx.Exec(ctx, `
|
||||||
UPDATE quota_reservations
|
UPDATE quota_reservations
|
||||||
SET status = 'confirmed',
|
SET status = 'confirmed',
|
||||||
@@ -1526,13 +1529,10 @@ func confirmAIPointReservation(ctx context.Context, tx pgx.Tx, tenantID, reserva
|
|||||||
if err != nil {
|
if err != nil {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
if tag.RowsAffected() == 0 {
|
return ignoreAIPointReservationReset(ctx, tx, tag.RowsAffected() > 0, tenantID, reservationID)
|
||||||
return errors.New("ai points reservation is not pending")
|
|
||||||
}
|
|
||||||
return nil
|
|
||||||
}
|
}
|
||||||
|
|
||||||
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, `
|
tag, err := tx.Exec(ctx, `
|
||||||
UPDATE quota_reservations
|
UPDATE quota_reservations
|
||||||
SET status = 'refunded',
|
SET status = 'refunded',
|
||||||
@@ -1543,13 +1543,34 @@ func refundAIPointReservation(ctx context.Context, tx pgx.Tx, tenantID, reservat
|
|||||||
AND quota_type = $3
|
AND quota_type = $3
|
||||||
AND status = 'pending'
|
AND status = 'pending'
|
||||||
`, reservationID, tenantID, aiPointsQuotaType)
|
`, 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 {
|
if err != nil {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
if tag.RowsAffected() == 0 {
|
if status == "reset" {
|
||||||
return errors.New("ai points reservation is not pending")
|
return nil
|
||||||
}
|
}
|
||||||
return nil
|
return errors.New("ai points reservation is not pending")
|
||||||
}
|
}
|
||||||
|
|
||||||
func (s *Service) markReviewJobDone(ctx context.Context, jobID int64) error {
|
func (s *Service) markReviewJobDone(ctx context.Context, jobID int64) error {
|
||||||
|
|||||||
@@ -13,7 +13,7 @@ import (
|
|||||||
|
|
||||||
const confirmReservation = `-- name: ConfirmReservation :exec
|
const confirmReservation = `-- name: ConfirmReservation :exec
|
||||||
UPDATE quota_reservations SET status = 'confirmed', consumed_amount = reserved_amount, updated_at = NOW()
|
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 {
|
type ConfirmReservationParams struct {
|
||||||
@@ -114,7 +114,7 @@ func (q *Queries) InsertQuotaLedger(ctx context.Context, arg InsertQuotaLedgerPa
|
|||||||
|
|
||||||
const refundReservation = `-- name: RefundReservation :exec
|
const refundReservation = `-- name: RefundReservation :exec
|
||||||
UPDATE quota_reservations SET status = 'refunded', refunded_amount = reserved_amount, updated_at = NOW()
|
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 {
|
type RefundReservationParams struct {
|
||||||
|
|||||||
@@ -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)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -23,8 +23,8 @@ WHERE id = sqlc.arg(id) AND tenant_id = sqlc.arg(tenant_id);
|
|||||||
|
|
||||||
-- name: ConfirmReservation :exec
|
-- name: ConfirmReservation :exec
|
||||||
UPDATE quota_reservations SET status = 'confirmed', consumed_amount = reserved_amount, updated_at = NOW()
|
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
|
-- name: RefundReservation :exec
|
||||||
UPDATE quota_reservations SET status = 'refunded', refunded_amount = reserved_amount, updated_at = NOW()
|
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';
|
||||||
|
|||||||
@@ -59,6 +59,7 @@ type QuotaRepository interface {
|
|||||||
UpdateQuotaReservationResource(ctx context.Context, reservationID, tenantID, resourceID int64) error
|
UpdateQuotaReservationResource(ctx context.Context, reservationID, tenantID, resourceID int64) error
|
||||||
ConfirmReservation(ctx context.Context, reservationID, tenantID int64) error
|
ConfirmReservation(ctx context.Context, reservationID, tenantID int64) error
|
||||||
RefundReservation(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 {
|
type quotaRepository struct {
|
||||||
@@ -256,3 +257,19 @@ func (r *quotaRepository) RefundReservation(ctx context.Context, reservationID,
|
|||||||
TenantID: tenantID,
|
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
|
||||||
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user