fix(billing): 重排任务复用钱包预留

平台额度跨窗口重排时保留任务已有的钱包预留,后续 Worker 领取复用同一幂等键,避免重复 reserve/release 事务。\n\n补齐排队任务取消后的原子结算释放,并允许任务 attempt 持久化软负载、满载原因和选择原因等有界路由快照。\n\n验证:go test ./...;go vet ./...;钱包预留复用与排队取消 PostgreSQL 集成测试。
This commit is contained in:
2026-08-03 18:09:47 +08:00
parent 13faf1d072
commit c8d04ca731
4 changed files with 147 additions and 1 deletions
@@ -36,6 +36,21 @@ func TestMinimalTaskResultDropsProviderRawCopy(t *testing.T) {
}
}
func TestMinimalTaskAttemptMetricsKeepsBoundedRoutingSnapshot(t *testing.T) {
metrics := minimalTaskAttemptMetrics(map[string]any{
"loadMetrics": map[string]any{"concurrent": map[string]any{"softRatio": 0.5}},
"fullReasons": []any{"waiting"},
"selectionReason": "full_avoided",
"requestBody": map[string]any{"secret": "must not persist"},
})
if metrics["loadMetrics"] == nil || metrics["fullReasons"] == nil || metrics["selectionReason"] != "full_avoided" {
t.Fatalf("routing snapshot was dropped: %+v", metrics)
}
if metrics["requestBody"] != nil {
t.Fatalf("unbounded attempt payload was retained: %+v", metrics)
}
}
func TestTaskEventsAreMinimalAndConsecutiveDuplicatesAreSkipped(t *testing.T) {
db := billingV2IntegrationStore(t)
ctx := context.Background()
+21 -1
View File
@@ -901,6 +901,12 @@ SET status = 'cancelled',
error = NULL,
error_code = 'task_cancelled',
error_message = NULLIF($2::text, ''),
billing_status = CASE
WHEN run_mode NOT IN ('production', 'acceptance', 'acceptance_canary') OR gateway_user_id IS NULL THEN 'not_required'
WHEN reservation_amount > 0 THEN 'pending'
ELSE 'released'
END,
billing_updated_at = now(),
remote_task_payload = '{}'::jsonb,
locked_by = NULL,
locked_at = NULL,
@@ -930,7 +936,20 @@ WHERE task_id = $1::uuid
if _, err := tx.Exec(ctx, `DELETE FROM gateway_task_admissions WHERE task_id = $1::uuid`, taskID); err != nil {
return err
}
return nil
payloadJSON, _ := json.Marshal(map[string]any{"taskId": taskID, "reason": "queued_cancelled"})
_, err = tx.Exec(ctx, `
INSERT INTO settlement_outbox (
task_id, event_type, action, amount, currency, pricing_snapshot, payload, status, next_attempt_at
)
SELECT id, 'task.billing.release', 'release', reservation_amount, billing_currency,
pricing_snapshot, $2::jsonb, 'pending', now()
FROM gateway_tasks
WHERE id = $1::uuid
AND run_mode IN ('production', 'acceptance', 'acceptance_canary')
AND gateway_user_id IS NOT NULL
AND reservation_amount > 0
ON CONFLICT (task_id, event_type) DO NOTHING`, taskID, string(payloadJSON))
return err
})
if err != nil {
return GatewayTask{}, false, err
@@ -1568,6 +1587,7 @@ func minimalTaskAttemptUsage(usage map[string]any) map[string]any {
func minimalTaskAttemptMetrics(metrics map[string]any) map[string]any {
return whitelistedTaskAttemptMap(metrics, []string{
"platformPriority", "currentPriority", "loadRatio", "loadAvoided",
"loadMetrics", "fullReasons", "selectionReason",
"cacheAffinityKey", "cacheAdjustedPriority", "cacheAffinitySamples",
"cacheAffinityScore", "cacheAffinityConfidence", "cacheAffinityHitRatio",
"cacheAffinityEMAHitRatio", "cacheAffinityLastHitRatio",
@@ -84,6 +84,28 @@ func TestReserveTaskBillingSerializesConcurrentWalletReservations(t *testing.T)
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) {
@@ -138,6 +160,85 @@ func TestWalletReservationLockSetSerializesSameWallet(t *testing.T) {
<-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)