wangbo
c8d04ca731
平台额度跨窗口重排时保留任务已有的钱包预留，后续 Worker 领取复用同一幂等键，避免重复 reserve/release 事务。\n\n补齐排队任务取消后的原子结算释放，并允许任务 attempt 持久化软负载、满载原因和选择原因等有界路由快照。\n\n验证：go test ./...；go vet ./...；钱包预留复用与排队取消 PostgreSQL 集成测试。
302 lines
11 KiB
Go
302 lines
11 KiB
Go
package store
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"os"
|
|
"strconv"
|
|
"strings"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
|
|
)
|
|
|
|
func TestReserveTaskBillingSerializesConcurrentWalletReservations(t *testing.T) {
|
|
databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL"))
|
|
if databaseURL == "" {
|
|
t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run the wallet reservation integration test")
|
|
}
|
|
|
|
ctx := context.Background()
|
|
db, err := Connect(ctx, databaseURL)
|
|
if err != nil {
|
|
t.Fatalf("connect store: %v", err)
|
|
}
|
|
defer db.Close()
|
|
|
|
tenantID, userID := seedWalletReservationUser(t, ctx, db)
|
|
if _, err := db.SetUserWalletBalance(ctx, WalletBalanceAdjustmentInput{
|
|
GatewayUserID: userID,
|
|
Currency: "resource",
|
|
Balance: 10,
|
|
Reason: "seed wallet reservation test",
|
|
}); err != nil {
|
|
t.Fatalf("seed wallet balance: %v", err)
|
|
}
|
|
|
|
firstTaskID := newWalletReservationTestUUID(t, ctx, db)
|
|
secondTaskID := newWalletReservationTestUUID(t, ctx, db)
|
|
billings := []any{map[string]any{"currency": "resource", "amount": float64(10)}}
|
|
user := &auth.User{GatewayUserID: userID, GatewayTenantID: tenantID}
|
|
tasks := []GatewayTask{
|
|
{ID: firstTaskID, GatewayUserID: userID, GatewayTenantID: tenantID, Kind: "images.generations", Model: "mock-image"},
|
|
{ID: secondTaskID, GatewayUserID: userID, GatewayTenantID: tenantID, Kind: "videos.generations", Model: "mock-video"},
|
|
}
|
|
|
|
type reserveResult struct {
|
|
reservations []WalletBillingReservation
|
|
err error
|
|
}
|
|
results := make(chan reserveResult, len(tasks))
|
|
var wg sync.WaitGroup
|
|
for _, task := range tasks {
|
|
task := task
|
|
wg.Add(1)
|
|
go func() {
|
|
defer wg.Done()
|
|
reservations, err := db.ReserveTaskBilling(ctx, task, user, billings)
|
|
results <- reserveResult{reservations: reservations, err: err}
|
|
}()
|
|
}
|
|
wg.Wait()
|
|
close(results)
|
|
|
|
var successReservations []WalletBillingReservation
|
|
successCount := 0
|
|
insufficientCount := 0
|
|
for result := range results {
|
|
if result.err == nil {
|
|
successCount++
|
|
successReservations = result.reservations
|
|
continue
|
|
}
|
|
if errors.Is(result.err, ErrInsufficientWalletBalance) {
|
|
insufficientCount++
|
|
continue
|
|
}
|
|
t.Fatalf("unexpected reservation error: %v", result.err)
|
|
}
|
|
if successCount != 1 || insufficientCount != 1 {
|
|
t.Fatalf("expected one successful reservation and one insufficient balance rejection, got success=%d insufficient=%d", successCount, insufficientCount)
|
|
}
|
|
if len(successReservations) != 1 || !walletFloatNear(successReservations[0].Amount, 10) {
|
|
t.Fatalf("unexpected successful reservations: %+v", successReservations)
|
|
}
|
|
reusedReservations, err := db.ReserveTaskBilling(ctx, GatewayTask{
|
|
ID: successReservations[0].TaskID,
|
|
GatewayUserID: userID,
|
|
GatewayTenantID: tenantID,
|
|
Kind: "images.generations",
|
|
Model: "mock-image",
|
|
}, user, billings)
|
|
if err != nil || len(reusedReservations) != 1 || reusedReservations[0].IdempotencyKey != successReservations[0].IdempotencyKey {
|
|
t.Fatalf("active reservation was not reused: reservations=%+v err=%v", reusedReservations, err)
|
|
}
|
|
var reserveTransactions int
|
|
if err := db.pool.QueryRow(ctx, `
|
|
SELECT count(*)
|
|
FROM gateway_wallet_transactions
|
|
WHERE reference_type='gateway_task'
|
|
AND reference_id=$1
|
|
AND transaction_type='reserve'`, successReservations[0].TaskID).Scan(&reserveTransactions); err != nil {
|
|
t.Fatalf("count reused wallet reservations: %v", err)
|
|
}
|
|
if reserveTransactions != 1 {
|
|
t.Fatalf("active reservation reuse created %d reserve transactions, want 1", reserveTransactions)
|
|
}
|
|
|
|
balance, frozen, spent := readWalletReservationAccount(t, ctx, db, userID)
|
|
if !walletFloatNear(balance, 10) || !walletFloatNear(frozen, 10) || !walletFloatNear(spent, 0) {
|
|
t.Fatalf("reservation should freeze balance without spending it, balance=%f frozen=%f spent=%f", balance, frozen, spent)
|
|
}
|
|
|
|
settleTask := GatewayTask{
|
|
ID: successReservations[0].TaskID,
|
|
GatewayUserID: userID,
|
|
GatewayTenantID: tenantID,
|
|
Kind: "images.generations",
|
|
Model: "mock-image",
|
|
ResolvedModel: "mock-image",
|
|
Billings: billings,
|
|
BillingSummary: map[string]any{"currency": "resource", "totalAmount": float64(10)},
|
|
FinalChargeAmount: 10,
|
|
}
|
|
if err := db.SettleTaskBilling(ctx, settleTask); err != nil {
|
|
t.Fatalf("settle reserved task billing: %v", err)
|
|
}
|
|
if err := db.SettleTaskBilling(ctx, settleTask); err != nil {
|
|
t.Fatalf("settle reserved task billing should be idempotent: %v", err)
|
|
}
|
|
balance, frozen, spent = readWalletReservationAccount(t, ctx, db, userID)
|
|
if !walletFloatNear(balance, 0) || !walletFloatNear(frozen, 0) || !walletFloatNear(spent, 10) {
|
|
t.Fatalf("settlement should release reservation and debit once, balance=%f frozen=%f spent=%f", balance, frozen, spent)
|
|
}
|
|
}
|
|
|
|
func TestWalletReservationLockSetSerializesSameWallet(t *testing.T) {
|
|
var locks walletReservationLockSet
|
|
unlock := locks.lock("wallet-user")
|
|
acquired := make(chan struct{})
|
|
released := make(chan struct{})
|
|
go func() {
|
|
defer close(released)
|
|
unlockSecond := locks.lock("wallet-user")
|
|
close(acquired)
|
|
unlockSecond()
|
|
}()
|
|
select {
|
|
case <-acquired:
|
|
t.Fatal("same-wallet reservation lock was acquired concurrently")
|
|
case <-time.After(20 * time.Millisecond):
|
|
}
|
|
unlock()
|
|
select {
|
|
case <-acquired:
|
|
case <-time.After(time.Second):
|
|
t.Fatal("same-wallet reservation lock did not unblock")
|
|
}
|
|
<-released
|
|
}
|
|
|
|
func TestCancelQueuedTaskReleasesRetainedBillingReservation(t *testing.T) {
|
|
databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL"))
|
|
if databaseURL == "" {
|
|
t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run the queued cancellation billing integration test")
|
|
}
|
|
|
|
ctx := context.Background()
|
|
db, err := Connect(ctx, databaseURL)
|
|
if err != nil {
|
|
t.Fatalf("connect store: %v", err)
|
|
}
|
|
defer db.Close()
|
|
|
|
tenantID, userID := seedWalletReservationUser(t, ctx, db)
|
|
if _, err := db.SetUserWalletBalance(ctx, WalletBalanceAdjustmentInput{
|
|
GatewayUserID: userID,
|
|
Currency: "resource",
|
|
Balance: 10,
|
|
Reason: "seed queued cancellation reservation test",
|
|
}); err != nil {
|
|
t.Fatalf("seed wallet balance: %v", err)
|
|
}
|
|
user := &auth.User{GatewayUserID: userID, GatewayTenantID: tenantID}
|
|
task, err := db.CreateTask(ctx, CreateTaskInput{
|
|
Kind: "videos.generations", Model: "queued-cancellation", RunMode: "production", Async: true,
|
|
Request: map[string]any{"model": "queued-cancellation"},
|
|
}, user)
|
|
if err != nil {
|
|
t.Fatalf("create queued task: %v", err)
|
|
}
|
|
t.Cleanup(func() {
|
|
cleanupCtx := context.Background()
|
|
_, _ = db.pool.Exec(cleanupCtx, `DELETE FROM settlement_outbox WHERE task_id=$1::uuid`, task.ID)
|
|
_, _ = db.pool.Exec(cleanupCtx, `DELETE FROM gateway_tasks WHERE id=$1::uuid`, task.ID)
|
|
_, _ = db.pool.Exec(cleanupCtx, `DELETE FROM gateway_wallet_transactions WHERE gateway_user_id=$1::uuid`, userID)
|
|
_, _ = db.pool.Exec(cleanupCtx, `
|
|
DELETE FROM gateway_wallet_account_audit_guards
|
|
WHERE account_id IN (SELECT id FROM gateway_wallet_accounts WHERE gateway_user_id=$1::uuid)`, userID)
|
|
_, _ = db.pool.Exec(cleanupCtx, `DELETE FROM gateway_wallet_accounts WHERE gateway_user_id=$1::uuid`, userID)
|
|
})
|
|
|
|
reservations, err := db.ReserveTaskBilling(ctx, task, user, []any{map[string]any{"currency": "resource", "amount": float64(1)}})
|
|
if err != nil || len(reservations) != 1 {
|
|
t.Fatalf("reserve queued task billing: reservations=%+v err=%v", reservations, err)
|
|
}
|
|
if _, changed, err := db.CancelQueuedTask(ctx, task.ID, "cancel retained reservation"); err != nil || !changed {
|
|
t.Fatalf("cancel queued task: changed=%v err=%v", changed, err)
|
|
}
|
|
|
|
claims, err := db.ClaimBillingSettlements(ctx, "queued-cancellation-test", BillingSettlementBatchSize, BillingSettlementLockTimeout)
|
|
if err != nil {
|
|
t.Fatalf("claim cancellation release: %v", err)
|
|
}
|
|
settlement := settlementForTask(t, claims, task.ID)
|
|
if settlement.Action != "release" {
|
|
t.Fatalf("queued cancellation settlement=%+v, want release", settlement)
|
|
}
|
|
if err := db.ProcessBillingSettlement(ctx, settlement); err != nil {
|
|
t.Fatalf("process cancellation release: %v", err)
|
|
}
|
|
|
|
_, frozen, spent := readWalletReservationAccount(t, ctx, db, userID)
|
|
if !walletFloatNear(frozen, 0) || !walletFloatNear(spent, 0) {
|
|
t.Fatalf("queued cancellation did not release hold: frozen=%f spent=%f", frozen, spent)
|
|
}
|
|
var reserveCount int
|
|
var releaseCount int
|
|
if err := db.pool.QueryRow(ctx, `
|
|
SELECT count(*) FILTER (WHERE transaction_type='reserve'),
|
|
count(*) FILTER (WHERE transaction_type='release')
|
|
FROM gateway_wallet_transactions
|
|
WHERE reference_type='gateway_task' AND reference_id=$1`, task.ID).Scan(&reserveCount, &releaseCount); err != nil {
|
|
t.Fatalf("count queued cancellation transactions: %v", err)
|
|
}
|
|
if reserveCount != 1 || releaseCount != 1 {
|
|
t.Fatalf("queued cancellation transactions reserve/release=%d/%d, want 1/1", reserveCount, releaseCount)
|
|
}
|
|
}
|
|
|
|
func seedWalletReservationUser(t *testing.T, ctx context.Context, db *Store) (string, string) {
|
|
t.Helper()
|
|
suffix := strconv.FormatInt(time.Now().UnixNano(), 10)
|
|
var tenantID string
|
|
if err := db.pool.QueryRow(ctx, `
|
|
INSERT INTO gateway_tenants (tenant_key, name)
|
|
VALUES ($1, $2)
|
|
RETURNING id::text`, "wallet-reservation-"+suffix, "Wallet Reservation Test "+suffix).Scan(&tenantID); err != nil {
|
|
t.Fatalf("insert test tenant: %v", err)
|
|
}
|
|
var userID string
|
|
if err := db.pool.QueryRow(ctx, `
|
|
INSERT INTO gateway_users (user_key, username, gateway_tenant_id, tenant_key, roles)
|
|
VALUES ($1, $2, $3::uuid, $4, '["basic"]'::jsonb)
|
|
RETURNING id::text`, "wallet-reservation-user-"+suffix, "wallet_reservation_"+suffix, tenantID, "wallet-reservation-"+suffix).Scan(&userID); err != nil {
|
|
t.Fatalf("insert test user: %v", err)
|
|
}
|
|
t.Cleanup(func() {
|
|
cleanupCtx := context.Background()
|
|
_, _ = db.pool.Exec(cleanupCtx, `DELETE FROM gateway_wallet_transactions WHERE gateway_user_id = $1::uuid`, userID)
|
|
_, _ = db.pool.Exec(cleanupCtx, `
|
|
DELETE FROM gateway_wallet_account_audit_guards
|
|
WHERE account_id IN (SELECT id FROM gateway_wallet_accounts WHERE gateway_user_id = $1::uuid)`, userID)
|
|
_, _ = db.pool.Exec(cleanupCtx, `DELETE FROM gateway_users WHERE id = $1::uuid`, userID)
|
|
_, _ = db.pool.Exec(cleanupCtx, `DELETE FROM gateway_tenants WHERE id = $1::uuid`, tenantID)
|
|
})
|
|
return tenantID, userID
|
|
}
|
|
|
|
func newWalletReservationTestUUID(t *testing.T, ctx context.Context, db *Store) string {
|
|
t.Helper()
|
|
var id string
|
|
if err := db.pool.QueryRow(ctx, `SELECT gen_random_uuid()::text`).Scan(&id); err != nil {
|
|
t.Fatalf("generate uuid: %v", err)
|
|
}
|
|
return id
|
|
}
|
|
|
|
func readWalletReservationAccount(t *testing.T, ctx context.Context, db *Store, userID string) (float64, float64, float64) {
|
|
t.Helper()
|
|
var balance float64
|
|
var frozen float64
|
|
var spent float64
|
|
if err := db.pool.QueryRow(ctx, `
|
|
SELECT balance::float8, frozen_balance::float8, total_spent::float8
|
|
FROM gateway_wallet_accounts
|
|
WHERE gateway_user_id = $1::uuid
|
|
AND currency = 'resource'`, userID).Scan(&balance, &frozen, &spent); err != nil {
|
|
t.Fatalf("read wallet account: %v", err)
|
|
}
|
|
return balance, frozen, spent
|
|
}
|
|
|
|
func walletFloatNear(a float64, b float64) bool {
|
|
delta := a - b
|
|
if delta < 0 {
|
|
delta = -delta
|
|
}
|
|
return delta < 0.000001
|
|
}
|