fix(billing): 防止余额计费持续透支

This commit is contained in:
wucm667
2026-06-22 10:30:27 +08:00
parent e5f38a6f60
commit 9f5b57fc96
10 changed files with 335 additions and 7 deletions
+8
View File
@@ -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")
}
+5
View File
@@ -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 },
@@ -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) {
@@ -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())
}
+11
View File
@@ -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)
@@ -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
}
@@ -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)
}
+23 -1
View File
@@ -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.
@@ -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)
}
+4
View File
@@ -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