diff --git a/backend/internal/config/config.go b/backend/internal/config/config.go index 2fefe820c9..748b751463 100644 --- a/backend/internal/config/config.go +++ b/backend/internal/config/config.go @@ -643,6 +643,10 @@ type ProxyProbeConfig struct { type BillingConfig struct { CircuitBreaker CircuitBreakerConfig `mapstructure:"circuit_breaker"` + // MinimumBalanceReserve is the conservative preflight floor for balance billing. + // Requests in balance mode are rejected when the cached balance is below this + // amount, even if it is still positive. Set to 0 to keep the legacy balance > 0 gate. + MinimumBalanceReserve float64 `mapstructure:"minimum_balance_reserve"` // UserPlatformQuotaCacheTTLSeconds 用户 × 平台 quota 缓存 TTL(秒),默认 86400=1天,覆盖典型 daily 窗口。 // 消费点: // - billing_cache_service.cacheWriteWorker 异步累加 @@ -1615,6 +1619,7 @@ func setDefaults() { viper.SetDefault("billing.circuit_breaker.failure_threshold", 5) viper.SetDefault("billing.circuit_breaker.reset_timeout_seconds", 30) viper.SetDefault("billing.circuit_breaker.half_open_requests", 3) + viper.SetDefault("billing.minimum_balance_reserve", 0.000001) viper.SetDefault("billing.user_platform_quota_cache_ttl_seconds", 86400) viper.SetDefault("billing.user_platform_quota_sentinel_ttl_seconds", 3600) @@ -2277,6 +2282,9 @@ func (c *Config) Validate() error { return fmt.Errorf("billing.circuit_breaker.half_open_requests must be positive") } } + if c.Billing.MinimumBalanceReserve < 0 { + return fmt.Errorf("billing.minimum_balance_reserve must be non-negative") + } if c.Database.MaxOpenConns <= 0 { return fmt.Errorf("database.max_open_conns must be positive") } diff --git a/backend/internal/config/config_test.go b/backend/internal/config/config_test.go index ecdfaac4d1..ec0df0e3bd 100644 --- a/backend/internal/config/config_test.go +++ b/backend/internal/config/config_test.go @@ -1157,6 +1157,11 @@ func TestValidateConfigErrors(t *testing.T) { mutate: func(c *Config) { c.Billing.CircuitBreaker.HalfOpenRequests = 0 }, wantErr: "billing.circuit_breaker.half_open_requests", }, + { + name: "billing minimum balance reserve", + mutate: func(c *Config) { c.Billing.MinimumBalanceReserve = -0.01 }, + wantErr: "billing.minimum_balance_reserve", + }, { name: "database max open conns", mutate: func(c *Config) { c.Database.MaxOpenConns = 0 }, diff --git a/backend/internal/repository/usage_billing_repo.go b/backend/internal/repository/usage_billing_repo.go index 62f48b58f4..91ac536eee 100644 --- a/backend/internal/repository/usage_billing_repo.go +++ b/backend/internal/repository/usage_billing_repo.go @@ -113,11 +113,12 @@ func (r *usageBillingRepository) applyUsageBillingEffects(ctx context.Context, t } if cmd.BalanceCost > 0 { - newBalance, err := deductUsageBillingBalance(ctx, tx, cmd.UserID, cmd.BalanceCost) + newBalance, sufficient, err := deductUsageBillingBalance(ctx, tx, cmd.UserID, cmd.BalanceCost) if err != nil { return err } result.NewBalance = &newBalance + result.BalanceOverdrafted = !sufficient } if cmd.APIKeyQuotaCost > 0 { @@ -173,9 +174,23 @@ func incrementUsageBillingSubscription(ctx context.Context, tx *sql.Tx, subscrip return service.ErrSubscriptionNotFound } -func deductUsageBillingBalance(ctx context.Context, tx *sql.Tx, userID int64, amount float64) (float64, error) { +func deductUsageBillingBalance(ctx context.Context, tx *sql.Tx, userID int64, amount float64) (float64, bool, error) { var newBalance float64 err := tx.QueryRowContext(ctx, ` + UPDATE users + SET balance = balance - $1, + updated_at = NOW() + WHERE id = $2 AND deleted_at IS NULL AND balance >= $1 + RETURNING balance + `, amount, userID).Scan(&newBalance) + if err == nil { + return newBalance, true, nil + } + if !errors.Is(err, sql.ErrNoRows) { + return 0, false, err + } + + err = tx.QueryRowContext(ctx, ` UPDATE users SET balance = balance - $1, updated_at = NOW() @@ -183,12 +198,12 @@ func deductUsageBillingBalance(ctx context.Context, tx *sql.Tx, userID int64, am RETURNING balance `, amount, userID).Scan(&newBalance) if errors.Is(err, sql.ErrNoRows) { - return 0, service.ErrUserNotFound + return 0, false, service.ErrUserNotFound } if err != nil { - return 0, err + return 0, false, err } - return newBalance, nil + return newBalance, false, nil } func incrementUsageBillingAPIKeyQuota(ctx context.Context, tx *sql.Tx, apiKeyID int64, amount float64) (bool, error) { diff --git a/backend/internal/repository/usage_billing_repo_unit_test.go b/backend/internal/repository/usage_billing_repo_unit_test.go new file mode 100644 index 0000000000..8ed5530a8f --- /dev/null +++ b/backend/internal/repository/usage_billing_repo_unit_test.go @@ -0,0 +1,119 @@ +//go:build unit + +package repository + +import ( + "context" + "database/sql" + "testing" + + "github.com/DATA-DOG/go-sqlmock" + "github.com/stretchr/testify/require" + + "github.com/Wei-Shaw/sub2api/internal/service" +) + +const ( + conditionalBalanceDeductSQL = `(?s)UPDATE users\s+SET balance = balance - \$1,\s+updated_at = NOW\(\)\s+WHERE id = \$2 AND deleted_at IS NULL AND balance >= \$1\s+RETURNING balance` + overdraftBalanceDeductSQL = `(?s)UPDATE users\s+SET balance = balance - \$1,\s+updated_at = NOW\(\)\s+WHERE id = \$2 AND deleted_at IS NULL\s+RETURNING balance` +) + +func TestDeductUsageBillingBalance_UsesSufficientBalanceGuard(t *testing.T) { + ctx := context.Background() + db, mock, err := sqlmock.New() + require.NoError(t, err) + defer func() { _ = db.Close() }() + + mock.ExpectBegin() + tx, err := db.BeginTx(ctx, nil) + require.NoError(t, err) + mock.ExpectQuery(conditionalBalanceDeductSQL). + WithArgs(2.5, int64(42)). + WillReturnRows(sqlmock.NewRows([]string{"balance"}).AddRow(7.5)) + mock.ExpectCommit() + + newBalance, sufficient, err := deductUsageBillingBalance(ctx, tx, 42, 2.5) + require.NoError(t, err) + require.True(t, sufficient) + require.InDelta(t, 7.5, newBalance, 0.000001) + require.NoError(t, tx.Commit()) + require.NoError(t, mock.ExpectationsWereMet()) +} + +func TestDeductUsageBillingBalance_RecordsOverdraftWhenGuardMisses(t *testing.T) { + ctx := context.Background() + db, mock, err := sqlmock.New() + require.NoError(t, err) + defer func() { _ = db.Close() }() + + mock.ExpectBegin() + tx, err := db.BeginTx(ctx, nil) + require.NoError(t, err) + mock.ExpectQuery(conditionalBalanceDeductSQL). + WithArgs(10.0, int64(42)). + WillReturnError(sql.ErrNoRows) + mock.ExpectQuery(overdraftBalanceDeductSQL). + WithArgs(10.0, int64(42)). + WillReturnRows(sqlmock.NewRows([]string{"balance"}).AddRow(-5.0)) + mock.ExpectCommit() + + newBalance, sufficient, err := deductUsageBillingBalance(ctx, tx, 42, 10) + require.NoError(t, err) + require.False(t, sufficient) + require.InDelta(t, -5.0, newBalance, 0.000001) + require.NoError(t, tx.Commit()) + require.NoError(t, mock.ExpectationsWereMet()) +} + +func TestApplyUsageBillingEffects_FlagsBalanceOverdraft(t *testing.T) { + ctx := context.Background() + db, mock, err := sqlmock.New() + require.NoError(t, err) + defer func() { _ = db.Close() }() + + mock.ExpectBegin() + tx, err := db.BeginTx(ctx, nil) + require.NoError(t, err) + mock.ExpectQuery(conditionalBalanceDeductSQL). + WithArgs(10.0, int64(42)). + WillReturnError(sql.ErrNoRows) + mock.ExpectQuery(overdraftBalanceDeductSQL). + WithArgs(10.0, int64(42)). + WillReturnRows(sqlmock.NewRows([]string{"balance"}).AddRow(-5.0)) + mock.ExpectCommit() + + result := &service.UsageBillingApplyResult{Applied: true} + err = (&usageBillingRepository{}).applyUsageBillingEffects(ctx, tx, &service.UsageBillingCommand{ + UserID: 42, + BalanceCost: 10, + }, result) + require.NoError(t, err) + require.NotNil(t, result.NewBalance) + require.InDelta(t, -5.0, *result.NewBalance, 0.000001) + require.True(t, result.BalanceOverdrafted) + require.NoError(t, tx.Commit()) + require.NoError(t, mock.ExpectationsWereMet()) +} + +func TestDeductUsageBillingBalance_ReturnsUserNotFoundWhenNoUserUpdated(t *testing.T) { + ctx := context.Background() + db, mock, err := sqlmock.New() + require.NoError(t, err) + defer func() { _ = db.Close() }() + + mock.ExpectBegin() + tx, err := db.BeginTx(ctx, nil) + require.NoError(t, err) + mock.ExpectQuery(conditionalBalanceDeductSQL). + WithArgs(10.0, int64(42)). + WillReturnError(sql.ErrNoRows) + mock.ExpectQuery(overdraftBalanceDeductSQL). + WithArgs(10.0, int64(42)). + WillReturnError(sql.ErrNoRows) + mock.ExpectRollback() + + _, _, err = deductUsageBillingBalance(ctx, tx, 42, 10) + require.ErrorIs(t, err, service.ErrUserNotFound) + require.NoError(t, tx.Rollback()) + require.NoError(t, mock.ExpectationsWereMet()) +} diff --git a/backend/internal/repository/user_repo.go b/backend/internal/repository/user_repo.go index 98122c0bea..3ac8dcfbf8 100644 --- a/backend/internal/repository/user_repo.go +++ b/backend/internal/repository/user_repo.go @@ -757,6 +757,17 @@ func (r *userRepository) UpdateBalance(ctx context.Context, id int64, amount flo func (r *userRepository) DeductBalance(ctx context.Context, id int64, amount float64) error { client := clientFromContext(ctx, r.client) n, err := client.User.Update(). + Where(dbuser.IDEQ(id), dbuser.BalanceGTE(amount)). + AddBalance(-amount). + Save(ctx) + if err != nil { + return err + } + if n > 0 { + return nil + } + + n, err = client.User.Update(). Where(dbuser.IDEQ(id)). AddBalance(-amount). Save(ctx) diff --git a/backend/internal/service/billing_cache_service.go b/backend/internal/service/billing_cache_service.go index b734fab13e..bfdb4c20bc 100644 --- a/backend/internal/service/billing_cache_service.go +++ b/backend/internal/service/billing_cache_service.go @@ -833,6 +833,21 @@ func (s *BillingCacheService) checkRPM(ctx context.Context, user *User, group *G return nil } +func (s *BillingCacheService) minimumBalanceReserve() float64 { + if s == nil || s.cfg == nil || s.cfg.Billing.MinimumBalanceReserve <= 0 { + return 0 + } + return s.cfg.Billing.MinimumBalanceReserve +} + +func (s *BillingCacheService) balanceBelowEligibilityThreshold(balance float64) bool { + if balance <= 0 { + return true + } + minimumReserve := s.minimumBalanceReserve() + return minimumReserve > 0 && balance < minimumReserve +} + // checkBalanceEligibility 检查余额模式资格 func (s *BillingCacheService) checkBalanceEligibility(ctx context.Context, userID int64) error { balance, err := s.GetUserBalance(ctx, userID) @@ -847,7 +862,7 @@ func (s *BillingCacheService) checkBalanceEligibility(ctx context.Context, userI s.circuitBreaker.OnSuccess() } - if balance <= 0 { + if s.balanceBelowEligibilityThreshold(balance) { return ErrInsufficientBalance } diff --git a/backend/internal/service/billing_cache_service_balance_test.go b/backend/internal/service/billing_cache_service_balance_test.go new file mode 100644 index 0000000000..06deb7010c --- /dev/null +++ b/backend/internal/service/billing_cache_service_balance_test.go @@ -0,0 +1,128 @@ +//go:build unit + +package service + +import ( + "context" + "errors" + "sync/atomic" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/stretchr/testify/require" +) + +type balanceEligibilityCacheStub struct { + billingCacheWorkerStub + + balance float64 + cacheMissAfterInvalidate bool + invalidated atomic.Bool + deductCalls atomic.Int64 + invalidateCalls atomic.Int64 +} + +func (s *balanceEligibilityCacheStub) GetUserBalance(context.Context, int64) (float64, error) { + if s.cacheMissAfterInvalidate && s.invalidated.Load() { + return 0, errors.New("cache miss") + } + return s.balance, nil +} + +func (s *balanceEligibilityCacheStub) DeductUserBalance(context.Context, int64, float64) error { + s.deductCalls.Add(1) + return nil +} + +func (s *balanceEligibilityCacheStub) InvalidateUserBalance(context.Context, int64) error { + s.invalidateCalls.Add(1) + s.invalidated.Store(true) + return nil +} + +func TestCheckBillingEligibility_RejectsBalanceBelowMinimumReserve(t *testing.T) { + cache := &balanceEligibilityCacheStub{balance: 0.005} + cfg := &config.Config{} + cfg.Billing.MinimumBalanceReserve = 0.01 + svc := NewBillingCacheService(cache, nil, nil, nil, nil, nil, cfg, nil) + t.Cleanup(svc.Stop) + + err := svc.CheckBillingEligibility(context.Background(), &User{ID: 1}, nil, nil, nil, "") + require.ErrorIs(t, err, ErrInsufficientBalance) +} + +func TestCheckBillingEligibility_AllowsBalanceAtMinimumReserve(t *testing.T) { + cache := &balanceEligibilityCacheStub{balance: 0.01} + cfg := &config.Config{} + cfg.Billing.MinimumBalanceReserve = 0.01 + svc := NewBillingCacheService(cache, nil, nil, nil, nil, nil, cfg, nil) + t.Cleanup(svc.Stop) + + err := svc.CheckBillingEligibility(context.Background(), &User{ID: 1}, nil, nil, nil, "") + require.NoError(t, err) +} + +func TestSyncBalanceCacheAfterDeduction_InvalidatesExhaustedBalance(t *testing.T) { + cache := &balanceEligibilityCacheStub{ + balance: 0.50, + cacheMissAfterInvalidate: true, + } + userRepo := &balanceLoadUserRepoStub{balance: -0.25} + cfg := &config.Config{} + cfg.Billing.MinimumBalanceReserve = 0.01 + svc := NewBillingCacheService(cache, userRepo, nil, nil, nil, nil, cfg, nil) + t.Cleanup(svc.Stop) + + newBalance := -0.25 + syncBalanceCacheAfterDeduction(context.Background(), &postUsageBillingParams{ + Cost: &CostBreakdown{ActualCost: 0.75}, + User: &User{ID: 1}, + }, &billingDeps{billingCacheService: svc}, &UsageBillingApplyResult{ + NewBalance: &newBalance, + BalanceOverdrafted: true, + }) + + require.Equal(t, int64(1), cache.invalidateCalls.Load()) + require.Equal(t, int64(0), cache.deductCalls.Load()) + + err := svc.CheckBillingEligibility(context.Background(), &User{ID: 1}, nil, nil, nil, "") + require.ErrorIs(t, err, ErrInsufficientBalance) + require.Equal(t, int64(1), userRepo.calls.Load()) +} + +func TestSyncBalanceCacheAfterDeduction_InvalidatesWhenBalanceFallsBelowReserve(t *testing.T) { + cache := &balanceEligibilityCacheStub{balance: 0.50} + cfg := &config.Config{} + cfg.Billing.MinimumBalanceReserve = 0.01 + svc := NewBillingCacheService(cache, nil, nil, nil, nil, nil, cfg, nil) + t.Cleanup(svc.Stop) + + newBalance := 0.005 + syncBalanceCacheAfterDeduction(context.Background(), &postUsageBillingParams{ + Cost: &CostBreakdown{ActualCost: 0.495}, + User: &User{ID: 1}, + }, &billingDeps{billingCacheService: svc}, &UsageBillingApplyResult{NewBalance: &newBalance}) + + require.Equal(t, int64(1), cache.invalidateCalls.Load()) + require.Equal(t, int64(0), cache.deductCalls.Load()) +} + +func TestSyncBalanceCacheAfterDeduction_QueuesDeductWhenBalanceStillEligible(t *testing.T) { + cache := &balanceEligibilityCacheStub{balance: 1} + cfg := &config.Config{} + cfg.Billing.MinimumBalanceReserve = 0.01 + svc := NewBillingCacheService(cache, nil, nil, nil, nil, nil, cfg, nil) + t.Cleanup(svc.Stop) + + newBalance := 0.75 + syncBalanceCacheAfterDeduction(context.Background(), &postUsageBillingParams{ + Cost: &CostBreakdown{ActualCost: 0.25}, + User: &User{ID: 1}, + }, &billingDeps{billingCacheService: svc}, &UsageBillingApplyResult{NewBalance: &newBalance}) + + require.Equal(t, int64(0), cache.invalidateCalls.Load()) + require.Eventually(t, func() bool { + return cache.deductCalls.Load() == 1 + }, 2*time.Second, 10*time.Millisecond) +} diff --git a/backend/internal/service/gateway_service.go b/backend/internal/service/gateway_service.go index cfbc5dec48..d9f3ef93b6 100644 --- a/backend/internal/service/gateway_service.go +++ b/backend/internal/service/gateway_service.go @@ -8913,6 +8913,10 @@ func postUsageBilling(ctx context.Context, p *postUsageBillingParams, deps *bill if cost.ActualCost > 0 { if err := deps.userRepo.DeductBalance(billingCtx, p.User.ID, cost.ActualCost); err != nil { slog.Error("deduct balance failed", "user_id", p.User.ID, "error", err) + } else if deps.billingCacheService != nil { + if err := deps.billingCacheService.InvalidateUserBalance(billingCtx, p.User.ID); err != nil { + slog.Warn("invalidate balance cache after legacy deduction failed", "user_id", p.User.ID, "error", err) + } } } } @@ -9093,7 +9097,7 @@ func finalizePostUsageBilling(ctx context.Context, p *postUsageBillingParams, de deps.billingCacheService.QueueUpdateSubscriptionUsage(p.User.ID, *p.APIKey.GroupID, p.Cost.ActualCost) } } else if p.Cost.ActualCost > 0 && p.User != nil { - deps.billingCacheService.QueueDeductBalance(p.User.ID, p.Cost.ActualCost) + syncBalanceCacheAfterDeduction(ctx, p, deps, result) } if p.Cost.ActualCost > 0 && p.APIKey != nil && p.APIKey.HasRateLimits() { @@ -9142,6 +9146,24 @@ func finalizePostUsageBilling(ctx context.Context, p *postUsageBillingParams, de go notifyAccountQuota(p, deps, result) } +func syncBalanceCacheAfterDeduction(ctx context.Context, p *postUsageBillingParams, deps *billingDeps, result *UsageBillingApplyResult) { + if p == nil || p.Cost == nil || p.User == nil || deps == nil || deps.billingCacheService == nil { + return + } + if result != nil && result.NewBalance != nil && deps.billingCacheService.balanceBelowEligibilityThreshold(*result.NewBalance) { + if err := deps.billingCacheService.InvalidateUserBalance(ctx, p.User.ID); err != nil { + slog.Warn("invalidate balance cache after exhausted deduction failed", + "user_id", p.User.ID, + "new_balance", *result.NewBalance, + "balance_overdrafted", result.BalanceOverdrafted, + "error", err, + ) + } + return + } + deps.billingCacheService.QueueDeductBalance(p.User.ID, p.Cost.ActualCost) +} + // notifyBalanceLow sends balance low notification after deduction. // When result.NewBalance is available (from DB transaction RETURNING), it is used directly // to reconstruct oldBalance, avoiding stale Redis reads and concurrent-deduction races. diff --git a/backend/internal/service/usage_billing.go b/backend/internal/service/usage_billing.go index 30495624b5..accc7cb2cb 100644 --- a/backend/internal/service/usage_billing.go +++ b/backend/internal/service/usage_billing.go @@ -115,6 +115,7 @@ type UsageBillingApplyResult struct { Applied bool APIKeyQuotaExhausted bool NewBalance *float64 // post-deduction balance (nil = no balance deduction) + BalanceOverdrafted bool // true when the sufficient-balance guard missed and debt was still recorded QuotaState *AccountQuotaState // post-increment quota state (nil = no quota increment) } diff --git a/deploy/config.example.yaml b/deploy/config.example.yaml index 2cb65d8fec..67cff357e7 100644 --- a/deploy/config.example.yaml +++ b/deploy/config.example.yaml @@ -1027,6 +1027,10 @@ billing: # Number of requests to allow in half-open state # 半开状态允许通过的请求数 half_open_requests: 3 + # Conservative minimum balance required before forwarding balance-billed requests. + # Set to 0 to only require balance > 0. + # 余额计费请求转发前要求的保守最小余额;设为 0 则仅要求余额 > 0。 + minimum_balance_reserve: 0.000001 # Cache TTL (seconds) for per-user × per-platform quota records # 用户 × 平台 quota 缓存 TTL(秒),默认 86400=1天,覆盖典型 daily 窗口 user_platform_quota_cache_ttl_seconds: 86400