From 43f88c99aa3efea40eb76fb55b48ae88a9f5dc13 Mon Sep 17 00:00:00 2001 From: William Drake Date: Mon, 28 Sep 2026 14:20:49 -0400 Subject: [PATCH] fix(billing): reserve in-flight balance to prevent concurrent overdraft Balance-mode admission only checked cached balance > threshold, so N concurrent expensive requests could all pass and overdraw the user. Add a Redis in-flight reservation: at admission estimate a conservative per-request cost (input estimate + output price x max_tokens, times rate multiplier) and atomically require cached_balance - inflight >= estimate (first in-flight request always admitted). Reservations are released via defer on every path and expire by TTL. Fails open on estimation / Redis errors. Configurable via billing.inflight_reservation (enabled by default). Co-Authored-By: Claude Opus 5.5 (1M context) --- backend/internal/config/config.go | 31 +++ backend/internal/handler/gateway_handler.go | 13 ++ .../gateway_handler_chat_completions.go | 13 ++ .../handler/gateway_handler_responses.go | 13 ++ .../handler/gateway_inflight_reservation.go | 49 ++++ .../gateway_inflight_reservation_test.go | 46 ++++ .../internal/handler/gemini_v1beta_handler.go | 13 ++ .../handler/openai_chat_completions.go | 13 ++ .../handler/openai_gateway_handler.go | 13 ++ .../repository/billing_inflight_cache.go | 102 +++++++++ .../repository/billing_inflight_cache_test.go | 211 ++++++++++++++++++ .../service/billing_inflight_reservation.go | 198 ++++++++++++++++ .../billing_inflight_reservation_test.go | 60 +++++ deploy/config.example.yaml | 22 ++ 14 files changed, 797 insertions(+) create mode 100644 backend/internal/handler/gateway_inflight_reservation.go create mode 100644 backend/internal/handler/gateway_inflight_reservation_test.go create mode 100644 backend/internal/repository/billing_inflight_cache.go create mode 100644 backend/internal/repository/billing_inflight_cache_test.go create mode 100644 backend/internal/service/billing_inflight_reservation.go create mode 100644 backend/internal/service/billing_inflight_reservation_test.go diff --git a/backend/internal/config/config.go b/backend/internal/config/config.go index 11b86e6b22..f7cfcead2e 100644 --- a/backend/internal/config/config.go +++ b/backend/internal/config/config.go @@ -921,6 +921,26 @@ type BillingConfig struct { // UserPlatformQuotaSentinelTTLSeconds sentinel(无 limit 占位)entry 的 TTL, // 显著短于 quota cache 默认 86400s 以控 Redis 内存;默认 3600=1h。 UserPlatformQuotaSentinelTTLSeconds int `mapstructure:"user_platform_quota_sentinel_ttl_seconds"` + // InflightReservation 余额模式在途请求预留（Redis），防止并发请求在预检时看到同一份余额而集体透支。 + InflightReservation InflightReservationConfig `mapstructure:"inflight_reservation"` +} + +// InflightReservationConfig 余额模式在途预留配置。 +// 准入时按 输入估算 + 输出单价 × max_tokens 估算单请求费用，在 Redis 中原子地 +// 校验 缓存余额 - 在途预留合计 >= 估算 后登记预留，请求结束（任意路径）释放。 +// 估算失败或 Redis 不可用时 fail-open，退回旧的仅余额 > 阈值检查。 +type InflightReservationConfig struct { + Enabled bool `mapstructure:"enabled"` + // TTLSeconds 单条预留的最长存活时间；进程崩溃等泄漏的预留到期自动失效。 + TTLSeconds int `mapstructure:"ttl_seconds"` + // DefaultMaxTokens 请求未携带 max_tokens 时用于估算的输出 token 数。 + DefaultMaxTokens int `mapstructure:"default_max_tokens"` + // MaxOutputTokens 估算输出 token 的上限（max_tokens 超出时截断）。 + MaxOutputTokens int `mapstructure:"max_output_tokens"` + // MaxInputTokens 输入 token 估算（请求体字节数 / 4）的上限。 + MaxInputTokens int `mapstructure:"max_input_tokens"` + // MaxReservationUSD 单请求预留金额上限；0 表示不设上限。 + MaxReservationUSD float64 `mapstructure:"max_reservation_usd"` } type CircuitBreakerConfig struct { @@ -2093,6 +2113,12 @@ func setDefaults() { 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) + viper.SetDefault("billing.inflight_reservation.enabled", true) + viper.SetDefault("billing.inflight_reservation.ttl_seconds", 900) + viper.SetDefault("billing.inflight_reservation.default_max_tokens", 8192) + viper.SetDefault("billing.inflight_reservation.max_output_tokens", 128000) + viper.SetDefault("billing.inflight_reservation.max_input_tokens", 200000) + viper.SetDefault("billing.inflight_reservation.max_reservation_usd", 0) // Turnstile viper.SetDefault("turnstile.required", false) @@ -3062,6 +3088,11 @@ func (c *Config) Validate() error { if c.Billing.MinimumBalanceReserve < 0 { return fmt.Errorf("billing.minimum_balance_reserve must be non-negative") } + if c.Billing.InflightReservation.TTLSeconds < 0 || c.Billing.InflightReservation.DefaultMaxTokens < 0 || + c.Billing.InflightReservation.MaxOutputTokens < 0 || c.Billing.InflightReservation.MaxInputTokens < 0 || + c.Billing.InflightReservation.MaxReservationUSD < 0 { + return fmt.Errorf("billing.inflight_reservation values must be non-negative") + } if c.Database.MaxOpenConns <= 0 { return fmt.Errorf("database.max_open_conns must be positive") } diff --git a/backend/internal/handler/gateway_handler.go b/backend/internal/handler/gateway_handler.go index 2f52c0bfe4..acf4e21acb 100644 --- a/backend/internal/handler/gateway_handler.go +++ b/backend/internal/handler/gateway_handler.go @@ -263,6 +263,19 @@ func (h *GatewayHandler) Messages(c *gin.Context) { return } + // 余额模式在途预留：防止并发请求在预检时看到同一份余额而集体透支。 + inflightRelease, err := reserveInflightBalance(c.Request.Context(), h.billingCacheService, h.gatewayService, apiKey, subscription, reqModel, body) + if err != nil { + reqLog.Info("gateway.inflight_reservation_rejected", zap.Error(err)) + status, code, message, retryAfter := billingErrorDetails(err) + if retryAfter > 0 { + c.Header("Retry-After", strconv.Itoa(retryAfter)) + } + h.handleStreamingAwareError(c, status, code, message, streamStarted) + return + } + defer inflightRelease() + // 设置请求所属分组 ID（用于渠道级功能判断，如 WebSearch 模拟） parsedReq.GroupID = apiKey.GroupID diff --git a/backend/internal/handler/gateway_handler_chat_completions.go b/backend/internal/handler/gateway_handler_chat_completions.go index 9eec2bb3d1..b844d27d1f 100644 --- a/backend/internal/handler/gateway_handler_chat_completions.go +++ b/backend/internal/handler/gateway_handler_chat_completions.go @@ -144,6 +144,19 @@ func (h *GatewayHandler) ChatCompletions(c *gin.Context) { return } + // 余额模式在途预留：防止并发请求在预检时看到同一份余额而集体透支。 + inflightRelease, err := reserveInflightBalance(c.Request.Context(), h.billingCacheService, h.gatewayService, apiKey, subscription, reqModel, body) + if err != nil { + reqLog.Info("gateway.cc.inflight_reservation_rejected", zap.Error(err)) + status, code, message, retryAfter := billingErrorDetails(err) + if retryAfter > 0 { + c.Header("Retry-After", strconv.Itoa(retryAfter)) + } + h.chatCompletionsErrorResponse(c, status, code, message) + return + } + defer inflightRelease() + // Parse request for session hash bodyRef := service.NewRequestBodyRef(body) parsedReq, _ := service.ParseGatewayRequest(bodyRef, "chat_completions") diff --git a/backend/internal/handler/gateway_handler_responses.go b/backend/internal/handler/gateway_handler_responses.go index 8cf07f1e9d..a10ad8d4ec 100644 --- a/backend/internal/handler/gateway_handler_responses.go +++ b/backend/internal/handler/gateway_handler_responses.go @@ -153,6 +153,19 @@ func (h *GatewayHandler) Responses(c *gin.Context) { return } + // 余额模式在途预留：防止并发请求在预检时看到同一份余额而集体透支。 + inflightRelease, err := reserveInflightBalance(requestCtx, h.billingCacheService, h.gatewayService, apiKey, subscription, reqModel, body) + if err != nil { + reqLog.Info("gateway.responses.inflight_reservation_rejected", zap.Error(err)) + status, code, message, retryAfter := billingErrorDetails(err) + if retryAfter > 0 { + c.Header("Retry-After", strconv.Itoa(retryAfter)) + } + h.responsesErrorResponse(c, status, code, message) + return + } + defer inflightRelease() + // Parse request for session hash bodyRef := service.NewRequestBodyRef(body) parsedReq, _ := service.ParseGatewayRequest(bodyRef, "responses") diff --git a/backend/internal/handler/gateway_inflight_reservation.go b/backend/internal/handler/gateway_inflight_reservation.go new file mode 100644 index 0000000000..08ebaafb9f --- /dev/null +++ b/backend/internal/handler/gateway_inflight_reservation.go @@ -0,0 +1,49 @@ +package handler + +import ( + "context" + + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/tidwall/gjson" +) + +// inflightReservationEstimator 估算单请求在途预留金额（USD）。 +type inflightReservationEstimator interface { + EstimateInflightReservation(ctx context.Context, apiKey *service.APIKey, model string, bodyBytes, maxTokens int) float64 +} + +// requestMaxOutputTokens 从请求体中提取输出 token 上限（兼容 Anthropic / OpenAI Chat / Responses / Gemini）。 +func requestMaxOutputTokens(body []byte) int { + for _, path := range []string{"max_tokens", "max_completion_tokens", "max_output_tokens", "generationConfig.maxOutputTokens", "generation_config.max_output_tokens"} { + if v := gjson.GetBytes(body, path); v.Exists() && v.Type == gjson.Number && v.Int() > 0 { + return int(v.Int()) + } + } + return 0 +} + +// reserveInflightBalance 在 CheckBillingEligibility 之后为余额模式请求登记在途预留。 +// 返回的 release 必须 defer 调用（覆盖成功/错误/客户端断开/failover 等所有路径）； +// 开关关闭、订阅模式、估算失败或 Redis 故障时返回 no-op（fail-open）。 +func reserveInflightBalance( + ctx context.Context, + billing *service.BillingCacheService, + estimator inflightReservationEstimator, + apiKey *service.APIKey, + subscription *service.UserSubscription, + model string, + body []byte, +) (func(), error) { + noop := func() {} + if billing == nil || estimator == nil || apiKey == nil || apiKey.User == nil || !billing.InflightReservationEnabled() { + return noop, nil + } + if apiKey.Group != nil && apiKey.Group.IsSubscriptionType() && subscription != nil { + return noop, nil + } + estimate := estimator.EstimateInflightReservation(ctx, apiKey, model, len(body), requestMaxOutputTokens(body)) + if estimate <= 0 { + return noop, nil + } + return billing.ReserveInflightBalance(ctx, apiKey.User, apiKey.Group, subscription, estimate) +} diff --git a/backend/internal/handler/gateway_inflight_reservation_test.go b/backend/internal/handler/gateway_inflight_reservation_test.go new file mode 100644 index 0000000000..a76f5eb6cb --- /dev/null +++ b/backend/internal/handler/gateway_inflight_reservation_test.go @@ -0,0 +1,46 @@ +package handler + +import ( + "context" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/stretchr/testify/require" +) + +func TestRequestMaxOutputTokens(t *testing.T) { + require.Equal(t, 1024, requestMaxOutputTokens([]byte(`{"max_tokens":1024}`))) + require.Equal(t, 2048, requestMaxOutputTokens([]byte(`{"max_completion_tokens":2048}`))) + require.Equal(t, 4096, requestMaxOutputTokens([]byte(`{"max_output_tokens":4096}`))) + require.Equal(t, 512, requestMaxOutputTokens([]byte(`{"generationConfig":{"maxOutputTokens":512}}`))) + require.Equal(t, 0, requestMaxOutputTokens([]byte(`{"max_tokens":"x"}`))) + require.Equal(t, 0, requestMaxOutputTokens([]byte(`{}`))) +} + +type countingEstimator struct{ calls int } + +func (e *countingEstimator) EstimateInflightReservation(context.Context, *service.APIKey, string, int, int) float64 { + e.calls++ + return 1 +} + +func TestReserveInflightBalance_SkipsWhenDisabledOrSubscription(t *testing.T) { + cfg := &config.Config{} + billing := service.NewBillingCacheService(nil, nil, nil, nil, nil, nil, cfg, nil) + t.Cleanup(billing.Stop) + est := &countingEstimator{} + apiKey := &service.APIKey{User: &service.User{ID: 1}} + + release, err := reserveInflightBalance(context.Background(), billing, est, apiKey, nil, "m", []byte(`{}`)) + require.NoError(t, err) + release() + require.Equal(t, 0, est.calls, "disabled switch must not even estimate") + + cfg.Billing.InflightReservation.Enabled = true + apiKey.Group = &service.Group{SubscriptionType: service.SubscriptionTypeSubscription} + release, err = reserveInflightBalance(context.Background(), billing, est, apiKey, &service.UserSubscription{}, "m", []byte(`{}`)) + require.NoError(t, err) + release() + require.Equal(t, 0, est.calls, "subscription mode must be unaffected") +} diff --git a/backend/internal/handler/gemini_v1beta_handler.go b/backend/internal/handler/gemini_v1beta_handler.go index 140e294794..0fc9e5eba5 100644 --- a/backend/internal/handler/gemini_v1beta_handler.go +++ b/backend/internal/handler/gemini_v1beta_handler.go @@ -390,6 +390,19 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) { return } + // 余额模式在途预留：防止并发请求在预检时看到同一份余额而集体透支。 + inflightRelease, err := reserveInflightBalance(c.Request.Context(), h.billingCacheService, h.gatewayService, apiKey, subscription, modelName, body) + if err != nil { + reqLog.Info("gemini.inflight_reservation_rejected", zap.Error(err)) + status, _, message, retryAfter := billingErrorDetails(err) + if retryAfter > 0 { + c.Header("Retry-After", strconv.Itoa(retryAfter)) + } + googleError(c, status, message) + return + } + defer inflightRelease() + // 3) select account (sticky session based on request body) // 优先使用 Gemini CLI 的会话标识（privileged-user-id + tmp 目录哈希） sessionHash := extractGeminiCLISessionHash(c, body) diff --git a/backend/internal/handler/openai_chat_completions.go b/backend/internal/handler/openai_chat_completions.go index 147f4460a8..afecc1b06a 100644 --- a/backend/internal/handler/openai_chat_completions.go +++ b/backend/internal/handler/openai_chat_completions.go @@ -146,6 +146,19 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) { return } + // 余额模式在途预留：防止并发请求在预检时看到同一份余额而集体透支。 + inflightRelease, err := reserveInflightBalance(c.Request.Context(), h.billingCacheService, h.gatewayService, apiKey, subscription, reqModel, body) + if err != nil { + reqLog.Info("openai_chat_completions.inflight_reservation_rejected", zap.Error(err)) + status, code, message, retryAfter := billingErrorDetails(err) + if retryAfter > 0 { + c.Header("Retry-After", strconv.Itoa(retryAfter)) + } + h.handleStreamingAwareError(c, status, code, message, streamStarted) + return + } + defer inflightRelease() + sessionHash := h.gatewayService.GenerateSessionHash(c, body) promptCacheKey := h.gatewayService.ExtractSessionID(c, body) diff --git a/backend/internal/handler/openai_gateway_handler.go b/backend/internal/handler/openai_gateway_handler.go index bf4563f2f2..5278219eaa 100644 --- a/backend/internal/handler/openai_gateway_handler.go +++ b/backend/internal/handler/openai_gateway_handler.go @@ -606,6 +606,19 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) { return } + // 余额模式在途预留：防止并发请求在预检时看到同一份余额而集体透支。 + inflightRelease, err := reserveInflightBalance(c.Request.Context(), h.billingCacheService, h.gatewayService, apiKey, subscription, reqModel, body) + if err != nil { + reqLog.Info("openai.inflight_reservation_rejected", zap.Error(err)) + status, code, message, retryAfter := billingErrorDetails(err) + if retryAfter > 0 { + c.Header("Retry-After", strconv.Itoa(retryAfter)) + } + h.handleStreamingAwareError(c, status, code, message, streamStarted) + return + } + defer inflightRelease() + // Generate session hash (header first; fallback to prompt_cache_key) sessionHash := h.gatewayService.GenerateSessionHash(c, sessionHashBody) if h.rejectIfCyberSessionBlocked(c, apiKey, sessionHashBody, reqModel, cyberBlockFormatResponses) { diff --git a/backend/internal/repository/billing_inflight_cache.go b/backend/internal/repository/billing_inflight_cache.go new file mode 100644 index 0000000000..47692fcf74 --- /dev/null +++ b/backend/internal/repository/billing_inflight_cache.go @@ -0,0 +1,102 @@ +package repository + +import ( + "context" + "fmt" + "strconv" + "time" + + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/redis/go-redis/v9" +) + +// 余额在途预留： +// - billing:inflight:{uid} ZSET member=requestID score=过期时间(ms) +// - billing:inflight_amt:{uid} HASH field=requestID value=预留金额(USD) +// +// 两个 key 使用同一 hash tag，保证 Redis Cluster 下 Lua 可同时访问。 +const ( + billingInflightKeyPrefix = "billing:inflight:" + billingInflightAmtKeyPrefix = "billing:inflight_amt:" +) + +func billingInflightKeys(userID int64) (string, string) { + return fmt.Sprintf("%s{%d}", billingInflightKeyPrefix, userID), + fmt.Sprintf("%s{%d}", billingInflightAmtKeyPrefix, userID) +} + +var ( + // KEYS: [1]=zset [2]=hash + // ARGV: [1]=now_ms [2]=expire_at_ms [3]=member [4]=amount [5]=balance [6]=key_ttl_ms + // 返回 {allowed, inflight_sum(string), inflight_count} + // 规则：先清理已过期成员；若无在途预留则直接放行（与旧行为一致，外层已校验余额 > 阈值）； + // 否则要求 balance - sum(在途) >= amount。放行时登记预留。 + reserveInflightBalanceScript = redis.NewScript(` + local now = tonumber(ARGV[1]) + local expired = redis.call('ZRANGEBYSCORE', KEYS[1], '-inf', now) + if #expired > 0 then + redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', now) + redis.call('HDEL', KEYS[2], unpack(expired)) + end + local vals = redis.call('HVALS', KEYS[2]) + local sum = 0 + for i = 1, #vals do + sum = sum + (tonumber(vals[i]) or 0) + end + local count = redis.call('ZCARD', KEYS[1]) + local amount = tonumber(ARGV[4]) + local balance = tonumber(ARGV[5]) + if count > 0 and (balance - sum) < amount then + return {0, tostring(sum), count} + end + redis.call('ZADD', KEYS[1], tonumber(ARGV[2]), ARGV[3]) + redis.call('HSET', KEYS[2], ARGV[3], ARGV[4]) + redis.call('PEXPIRE', KEYS[1], ARGV[6]) + redis.call('PEXPIRE', KEYS[2], ARGV[6]) + return {1, tostring(sum), count} + `) + + releaseInflightBalanceScript = redis.NewScript(` + redis.call('ZREM', KEYS[1], ARGV[1]) + redis.call('HDEL', KEYS[2], ARGV[1]) + return 1 + `) +) + +// ReserveInflightBalance 实现 service.InflightBalanceReservationCache。 +func (c *billingCache) ReserveInflightBalance(ctx context.Context, userID int64, requestID string, amount, balance float64, ttl time.Duration) (bool, float64, error) { + zkey, hkey := billingInflightKeys(userID) + now := time.Now().UnixMilli() + ttlMs := ttl.Milliseconds() + if ttlMs <= 0 { + ttlMs = 1 + } + res, err := reserveInflightBalanceScript.Run(ctx, c.rdb, []string{zkey, hkey}, + now, + now+ttlMs, + requestID, + strconv.FormatFloat(amount, 'f', -1, 64), + strconv.FormatFloat(balance, 'f', -1, 64), + ttlMs, + ).Slice() + if err != nil { + return false, 0, err + } + if len(res) < 2 { + return false, 0, fmt.Errorf("unexpected inflight reservation reply: %v", res) + } + allowed, _ := res[0].(int64) + var sum float64 + if s, ok := res[1].(string); ok { + sum, _ = strconv.ParseFloat(s, 64) + } + return allowed == 1, sum, nil +} + +// ReleaseInflightBalance 实现 service.InflightBalanceReservationCache。 +func (c *billingCache) ReleaseInflightBalance(ctx context.Context, userID int64, requestID string) error { + zkey, hkey := billingInflightKeys(userID) + return releaseInflightBalanceScript.Run(ctx, c.rdb, []string{zkey, hkey}, requestID).Err() +} + +var _ service.InflightBalanceReservationCache = (*billingCache)(nil) diff --git a/backend/internal/repository/billing_inflight_cache_test.go b/backend/internal/repository/billing_inflight_cache_test.go new file mode 100644 index 0000000000..6540c83bdf --- /dev/null +++ b/backend/internal/repository/billing_inflight_cache_test.go @@ -0,0 +1,211 @@ +package repository + +import ( + "context" + "errors" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/alicebob/miniredis/v2" + "github.com/redis/go-redis/v9" + "github.com/stretchr/testify/require" +) + +func newInflightTestEnv(t *testing.T, enabled bool, ttlSeconds int) (*miniredis.Miniredis, *billingCache, *service.BillingCacheService) { + t.Helper() + mr := miniredis.RunT(t) + rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()}) + t.Cleanup(func() { _ = rdb.Close() }) + cache := &billingCache{rdb: rdb} + cfg := &config.Config{} + cfg.Billing.InflightReservation = config.InflightReservationConfig{ + Enabled: enabled, + TTLSeconds: ttlSeconds, + DefaultMaxTokens: 8192, + } + svc := service.NewBillingCacheService(cache, nil, nil, nil, nil, nil, cfg, nil) + t.Cleanup(svc.Stop) + return mr, cache, svc +} + +func inflightCount(t *testing.T, cache *billingCache, userID int64) int64 { + t.Helper() + zkey, _ := billingInflightKeys(userID) + n, err := cache.rdb.ZCard(context.Background(), zkey).Result() + require.NoError(t, err) + return n +} + +func TestInflightReservation_ConcurrentAdmitsOnlyWhatBalanceCovers(t *testing.T) { + _, cache, svc := newInflightTestEnv(t, true, 60) + ctx := context.Background() + user := &service.User{ID: 42} + require.NoError(t, cache.SetUserBalance(ctx, user.ID, 1.0)) + + const n = 20 + var admitted, rejected atomic.Int32 + releases := make(chan func(), n) + var wg sync.WaitGroup + start := make(chan struct{}) + for i := 0; i < n; i++ { + wg.Add(1) + go func() { + defer wg.Done() + <-start + release, err := svc.ReserveInflightBalance(ctx, user, nil, nil, 0.3) + if err != nil { + if !errors.Is(err, service.ErrInsufficientBalance) { + t.Errorf("unexpected error: %v", err) + } + rejected.Add(1) + return + } + admitted.Add(1) + releases <- release + }() + } + close(start) + wg.Wait() + close(releases) + + // balance 1.0, estimate 0.3: first always admitted, then while 1.0 - inflight >= 0.3 → 3 total. + require.Equal(t, int32(3), admitted.Load()) + require.Equal(t, int32(n-3), rejected.Load()) + require.Equal(t, int64(3), inflightCount(t, cache, user.ID)) + + for release := range releases { + release() + release() // idempotent + } + require.Equal(t, int64(0), inflightCount(t, cache, user.ID)) + + // After release, capacity is available again. + release, err := svc.ReserveInflightBalance(ctx, user, nil, nil, 0.3) + require.NoError(t, err) + release() +} + +func TestInflightReservation_FirstRequestAlwaysAdmitted(t *testing.T) { + _, cache, svc := newInflightTestEnv(t, true, 60) + ctx := context.Background() + user := &service.User{ID: 7} + require.NoError(t, cache.SetUserBalance(ctx, user.ID, 0.01)) + + release, err := svc.ReserveInflightBalance(ctx, user, nil, nil, 5) + require.NoError(t, err, "a lone request keeps legacy behavior even if estimate exceeds balance") + _, err = svc.ReserveInflightBalance(ctx, user, nil, nil, 0.001) + require.ErrorIs(t, err, service.ErrInsufficientBalance) + release() + release2, err := svc.ReserveInflightBalance(ctx, user, nil, nil, 0.001) + require.NoError(t, err) + release2() +} + +func TestInflightReservation_TTLExpiryFreesLeakedReservation(t *testing.T) { + _, cache, svc := newInflightTestEnv(t, true, 1) + ctx := context.Background() + user := &service.User{ID: 9} + require.NoError(t, cache.SetUserBalance(ctx, user.ID, 1.0)) + + _, err := svc.ReserveInflightBalance(ctx, user, nil, nil, 0.8) // leaked: never released + require.NoError(t, err) + _, err = svc.ReserveInflightBalance(ctx, user, nil, nil, 0.8) + require.ErrorIs(t, err, service.ErrInsufficientBalance) + + time.Sleep(1100 * time.Millisecond) + release, err := svc.ReserveInflightBalance(ctx, user, nil, nil, 0.8) + require.NoError(t, err) + require.Equal(t, int64(1), inflightCount(t, cache, user.ID)) + release() +} + +func TestInflightReservation_DisabledKeepsLegacyBehavior(t *testing.T) { + _, cache, svc := newInflightTestEnv(t, false, 60) + ctx := context.Background() + user := &service.User{ID: 11} + require.NoError(t, cache.SetUserBalance(ctx, user.ID, 0.01)) + require.False(t, svc.InflightReservationEnabled()) + for i := 0; i < 5; i++ { + _, err := svc.ReserveInflightBalance(ctx, user, nil, nil, 10) + require.NoError(t, err) + } + require.Equal(t, int64(0), inflightCount(t, cache, user.ID)) +} + +func TestInflightReservation_SubscriptionUnaffected(t *testing.T) { + _, cache, svc := newInflightTestEnv(t, true, 60) + ctx := context.Background() + user := &service.User{ID: 12} + require.NoError(t, cache.SetUserBalance(ctx, user.ID, 0)) + group := &service.Group{ID: 1, SubscriptionType: service.SubscriptionTypeSubscription} + sub := &service.UserSubscription{ID: 1} + for i := 0; i < 5; i++ { + _, err := svc.ReserveInflightBalance(ctx, user, group, sub, 10) + require.NoError(t, err) + } + require.Equal(t, int64(0), inflightCount(t, cache, user.ID)) +} + +func TestInflightReservation_ZeroEstimateSkips(t *testing.T) { + _, cache, svc := newInflightTestEnv(t, true, 60) + ctx := context.Background() + user := &service.User{ID: 13} + require.NoError(t, cache.SetUserBalance(ctx, user.ID, 1)) + _, err := svc.ReserveInflightBalance(ctx, user, nil, nil, 0) + require.NoError(t, err) + require.Equal(t, int64(0), inflightCount(t, cache, user.ID)) +} + +// downReserveCache 余额读取正常，但预留走一个已关闭的 Redis。 +type downReserveCache struct { + service.BillingCache + down *billingCache +} + +func (d *downReserveCache) ReserveInflightBalance(ctx context.Context, userID int64, requestID string, amount, balance float64, ttl time.Duration) (bool, float64, error) { + return d.down.ReserveInflightBalance(ctx, userID, requestID, amount, balance, ttl) +} + +func (d *downReserveCache) ReleaseInflightBalance(ctx context.Context, userID int64, requestID string) error { + return d.down.ReleaseInflightBalance(ctx, userID, requestID) +} + +func TestInflightReservation_RedisDownFailsOpen(t *testing.T) { + mr := miniredis.RunT(t) + rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()}) + t.Cleanup(func() { _ = rdb.Close() }) + up := &billingCache{rdb: rdb} + require.NoError(t, up.SetUserBalance(context.Background(), 5, 0.01)) + + downMR := miniredis.RunT(t) + downRDB := redis.NewClient(&redis.Options{Addr: downMR.Addr(), MaxRetries: -1, DialTimeout: 100 * time.Millisecond}) + t.Cleanup(func() { _ = downRDB.Close() }) + downMR.Close() + + cfg := &config.Config{} + cfg.Billing.InflightReservation = config.InflightReservationConfig{Enabled: true, TTLSeconds: 60} + svc := service.NewBillingCacheService(&downReserveCache{BillingCache: up, down: &billingCache{rdb: downRDB}}, nil, nil, nil, nil, nil, cfg, nil) + t.Cleanup(svc.Stop) + + user := &service.User{ID: 5} + for i := 0; i < 3; i++ { + release, err := svc.ReserveInflightBalance(context.Background(), user, nil, nil, 100) + require.NoError(t, err) + release() + } +} + +func TestInflightReservation_ScriptErrorSurfaces(t *testing.T) { + mr := miniredis.RunT(t) + rdb := redis.NewClient(&redis.Options{Addr: mr.Addr(), MaxRetries: -1}) + cache := &billingCache{rdb: rdb} + mr.Close() + _, _, err := cache.ReserveInflightBalance(context.Background(), 1, "r", 1, 1, time.Second) + require.Error(t, err) + require.False(t, errors.Is(err, context.Canceled)) + _ = rdb.Close() +} diff --git a/backend/internal/service/billing_inflight_reservation.go b/backend/internal/service/billing_inflight_reservation.go new file mode 100644 index 0000000000..86b7e69a86 --- /dev/null +++ b/backend/internal/service/billing_inflight_reservation.go @@ -0,0 +1,198 @@ +package service + +import ( + "context" + "math" + "sync" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/pkg/logger" + "github.com/google/uuid" +) + +// InflightBalanceReservationCache 余额在途预留的缓存能力（可选）。 +// BillingCache 的 Redis 实现同时实现此接口；未实现时在途预留自动关闭（fail-open）。 +type InflightBalanceReservationCache interface { + // ReserveInflightBalance 原子地：清理过期预留；若用户已有在途预留且 + // balance - sum(在途) < amount 则拒绝；否则登记 requestID 的预留（ttl 后自动失效）。 + // 返回是否放行以及登记前的在途合计。 + ReserveInflightBalance(ctx context.Context, userID int64, requestID string, amount, balance float64, ttl time.Duration) (bool, float64, error) + // ReleaseInflightBalance 释放 requestID 的预留（幂等）。 + ReleaseInflightBalance(ctx context.Context, userID int64, requestID string) error +} + +const ( + defaultInflightReservationTTL = 15 * time.Minute + defaultInflightDefaultMaxTokens = 8192 + inflightReservationReleaseTimeout = 2 * time.Second + inflightReservationReserveTimeout = 2 * time.Second + inflightInputBytesPerTokenEstimate = 4 +) + +func noopRelease() {} + +func (s *BillingCacheService) inflightReservationConfig() (config.InflightReservationConfig, bool) { + if s == nil || s.cfg == nil { + return config.InflightReservationConfig{}, false + } + cfg := s.cfg.Billing.InflightReservation + if !cfg.Enabled || s.cfg.RunMode == config.RunModeSimple { + return cfg, false + } + return cfg, true +} + +// InflightReservationEnabled 是否启用余额在途预留。 +func (s *BillingCacheService) InflightReservationEnabled() bool { + _, ok := s.inflightReservationConfig() + return ok +} + +// ReserveInflightBalance 在余额模式下为本次请求登记在途预留。 +// +// 必须在 CheckBillingEligibility 通过后调用。返回的 release 必须在请求结束时调用 +// （所有路径，建议 defer），多次调用安全。 +// +// 以下情况直接放行且不登记预留（fail-open，保持旧行为）： +// 开关关闭 / 简易模式 / 订阅模式 / estimate <= 0 / 缓存不支持 / 余额读取失败 / Redis 执行失败。 +// 仅当 Redis 明确判定 缓存余额 - 在途合计 < estimate（且已有在途请求）时返回 ErrInsufficientBalance。 +func (s *BillingCacheService) ReserveInflightBalance(ctx context.Context, user *User, group *Group, subscription *UserSubscription, estimate float64) (func(), error) { + cfg, ok := s.inflightReservationConfig() + if !ok || user == nil { + return noopRelease, nil + } + if group != nil && group.IsSubscriptionType() && subscription != nil { + return noopRelease, nil + } + if cfg.MaxReservationUSD > 0 && estimate > cfg.MaxReservationUSD { + estimate = cfg.MaxReservationUSD + } + if estimate <= 0 || math.IsNaN(estimate) || math.IsInf(estimate, 0) { + return noopRelease, nil + } + rc, ok := s.cache.(InflightBalanceReservationCache) + if !ok || rc == nil { + return noopRelease, nil + } + + balance, err := s.GetUserBalance(ctx, user.ID) + if err != nil { + logger.LegacyPrintf("service.billing_cache", "Warning: inflight reservation balance read failed for user %d (fail-open): %v", user.ID, err) + return noopRelease, nil + } + + ttl := defaultInflightReservationTTL + if cfg.TTLSeconds > 0 { + ttl = time.Duration(cfg.TTLSeconds) * time.Second + } + requestID := uuid.NewString() + reserveCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), inflightReservationReserveTimeout) + allowed, inflight, err := rc.ReserveInflightBalance(reserveCtx, user.ID, requestID, estimate, balance, ttl) + cancel() + if err != nil { + logger.LegacyPrintf("service.billing_cache", "Warning: inflight reservation failed for user %d (fail-open): %v", user.ID, err) + return noopRelease, nil + } + if !allowed { + logger.LegacyPrintf("service.billing_cache", "inflight reservation rejected: user=%d balance=%.6f inflight=%.6f estimate=%.6f", user.ID, balance, inflight, estimate) + return noopRelease, ErrInsufficientBalance + } + + var once sync.Once + userID := user.ID + return func() { + once.Do(func() { + relCtx, relCancel := context.WithTimeout(context.Background(), inflightReservationReleaseTimeout) + defer relCancel() + if err := rc.ReleaseInflightBalance(relCtx, userID, requestID); err != nil { + logger.LegacyPrintf("service.billing_cache", "Warning: inflight reservation release failed for user %d (expires by ttl): %v", userID, err) + } + }) + }, nil +} + +// EstimateInflightReservationCost 估算单请求的保守费用（USD，已乘倍率）。 +// +// input_tokens = min(bodyBytes / 4, max_input_tokens) +// output_tokens = min(max_tokens 或 default_max_tokens, max_output_tokens) +// cost = (input_tokens × 输入单价 + output_tokens × 输出单价) × rateMultiplier +// +// 无法取得定价时返回 (0, false)，调用方应 fail-open。 +func EstimateInflightReservationCost(billing *BillingService, cfg config.InflightReservationConfig, model string, bodyBytes, maxTokens int, rateMultiplier float64) (float64, bool) { + if billing == nil || model == "" || rateMultiplier <= 0 { + return 0, false + } + pricing, err := billing.GetModelPricing(model) + if err != nil || pricing == nil { + return 0, false + } + inputTokens := 0 + if bodyBytes > 0 { + inputTokens = bodyBytes / inflightInputBytesPerTokenEstimate + } + if cfg.MaxInputTokens > 0 && inputTokens > cfg.MaxInputTokens { + inputTokens = cfg.MaxInputTokens + } + outputTokens := maxTokens + if outputTokens <= 0 { + outputTokens = cfg.DefaultMaxTokens + if outputTokens <= 0 { + outputTokens = defaultInflightDefaultMaxTokens + } + } + if cfg.MaxOutputTokens > 0 && outputTokens > cfg.MaxOutputTokens { + outputTokens = cfg.MaxOutputTokens + } + cost := (float64(inputTokens)*pricing.InputPricePerToken + float64(outputTokens)*pricing.OutputPricePerToken) * rateMultiplier + if cost <= 0 || math.IsNaN(cost) || math.IsInf(cost, 0) { + return 0, false + } + return cost, true +} + +func inflightReservationCfg(cfg *config.Config) config.InflightReservationConfig { + if cfg == nil { + return config.InflightReservationConfig{} + } + return cfg.Billing.InflightReservation +} + +func groupRateMultiplierOrOne(group *Group) float64 { + if group == nil { + return 1 + } + return group.RateMultiplier +} + +// EstimateInflightReservation 使用网关计费倍率估算在途预留金额；0 表示不预留。 +func (s *GatewayService) EstimateInflightReservation(ctx context.Context, apiKey *APIKey, model string, bodyBytes, maxTokens int) float64 { + if s == nil || apiKey == nil || apiKey.User == nil { + return 0 + } + rate := groupRateMultiplierOrOne(apiKey.Group) + if apiKey.GroupID != nil && apiKey.Group != nil { + rate = s.getUserGroupRateMultiplier(ctx, apiKey.User.ID, *apiKey.GroupID, rate) + } + cost, ok := EstimateInflightReservationCost(s.billingService, inflightReservationCfg(s.cfg), model, bodyBytes, maxTokens, rate) + if !ok { + return 0 + } + return cost +} + +// EstimateInflightReservation 使用 OpenAI 网关计费倍率估算在途预留金额；0 表示不预留。 +func (s *OpenAIGatewayService) EstimateInflightReservation(ctx context.Context, apiKey *APIKey, model string, bodyBytes, maxTokens int) float64 { + if s == nil || apiKey == nil || apiKey.User == nil { + return 0 + } + rate := groupRateMultiplierOrOne(apiKey.Group) + if apiKey.GroupID != nil && apiKey.Group != nil && s.userGroupRateResolver != nil { + rate = s.userGroupRateResolver.Resolve(ctx, apiKey.User.ID, *apiKey.GroupID, rate) + } + cost, ok := EstimateInflightReservationCost(s.billingService, inflightReservationCfg(s.cfg), model, bodyBytes, maxTokens, rate) + if !ok { + return 0 + } + return cost +} diff --git a/backend/internal/service/billing_inflight_reservation_test.go b/backend/internal/service/billing_inflight_reservation_test.go new file mode 100644 index 0000000000..f1dff41ec8 --- /dev/null +++ b/backend/internal/service/billing_inflight_reservation_test.go @@ -0,0 +1,60 @@ +//go:build unit + +package service + +import ( + "context" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/stretchr/testify/require" +) + +func TestEstimateInflightReservationCost(t *testing.T) { + billing := NewBillingService(&config.Config{}, nil) + pricing, err := billing.GetModelPricing("claude-sonnet-4-5") + require.NoError(t, err) + + cfg := config.InflightReservationConfig{DefaultMaxTokens: 8192, MaxOutputTokens: 64000, MaxInputTokens: 1000} + + // explicit max_tokens; body 400 bytes → 100 input tokens + cost, ok := EstimateInflightReservationCost(billing, cfg, "claude-sonnet-4-5", 400, 1000, 1) + require.True(t, ok) + require.InDelta(t, 100*pricing.InputPricePerToken+1000*pricing.OutputPricePerToken, cost, 1e-12) + + // missing max_tokens → default; rate multiplier applied + cost, ok = EstimateInflightReservationCost(billing, cfg, "claude-sonnet-4-5", 0, 0, 2) + require.True(t, ok) + require.InDelta(t, 2*8192*pricing.OutputPricePerToken, cost, 1e-12) + + // caps: input and output clamped + cost, ok = EstimateInflightReservationCost(billing, cfg, "claude-sonnet-4-5", 1_000_000, 1_000_000, 1) + require.True(t, ok) + require.InDelta(t, 1000*pricing.InputPricePerToken+64000*pricing.OutputPricePerToken, cost, 1e-12) + + // free group / missing billing → fail open + _, ok = EstimateInflightReservationCost(billing, cfg, "claude-sonnet-4-5", 10, 10, 0) + require.False(t, ok) + _, ok = EstimateInflightReservationCost(nil, cfg, "claude-sonnet-4-5", 10, 10, 1) + require.False(t, ok) + _, ok = EstimateInflightReservationCost(billing, cfg, "", 10, 10, 1) + require.False(t, ok) +} + +func TestReserveInflightBalance_NoReservationCacheFailsOpen(t *testing.T) { + cfg := &config.Config{} + cfg.Billing.InflightReservation.Enabled = true + svc := NewBillingCacheService(nil, nil, nil, nil, nil, nil, cfg, nil) + t.Cleanup(svc.Stop) + release, err := svc.ReserveInflightBalance(context.Background(), &User{ID: 1}, nil, nil, 100) + require.NoError(t, err) + release() +} + +func TestReserveInflightBalance_SimpleModeDisabled(t *testing.T) { + cfg := &config.Config{RunMode: config.RunModeSimple} + cfg.Billing.InflightReservation.Enabled = true + svc := NewBillingCacheService(nil, nil, nil, nil, nil, nil, cfg, nil) + t.Cleanup(svc.Stop) + require.False(t, svc.InflightReservationEnabled()) +} diff --git a/deploy/config.example.yaml b/deploy/config.example.yaml index e1e1c6df7e..de77ce6ad8 100644 --- a/deploy/config.example.yaml +++ b/deploy/config.example.yaml @@ -1126,6 +1126,28 @@ billing: # Cache TTL (seconds) for per-user × per-platform quota records # 用户 × 平台 quota 缓存 TTL（秒），默认 86400=1天，覆盖典型 daily 窗口 user_platform_quota_cache_ttl_seconds: 86400 + # In-flight balance reservation (Redis): at admission, estimate per-request cost + # (input estimate + output price x max_tokens) and atomically require + # cached_balance - in-flight reservations >= estimate. Released when the request ends. + # The first in-flight request of a user is always admitted (same as before). + # Fails open when the estimate cannot be computed or Redis is unavailable. + # 余额模式在途预留：准入时估算单请求费用，要求 缓存余额 - 在途预留合计 >= 估算； + # 用户无在途请求时始终放行（与旧行为一致）；估算失败/Redis 不可用时 fail-open。 + inflight_reservation: + enabled: true + # Max lifetime of a single reservation (self-heals leaked reservations) + # 单条预留最长存活时间（秒），泄漏的预留到期自动释放 + ttl_seconds: 900 + # Output tokens assumed when the request has no max_tokens + # 请求未携带 max_tokens 时假定的输出 token 数 + default_max_tokens: 8192 + # Cap for estimated output / input tokens + # 估算输出 / 输入 token 上限 + max_output_tokens: 128000 + max_input_tokens: 200000 + # Cap for a single reservation in USD (0 = no cap) + # 单请求预留金额上限（USD，0 为不设上限） + max_reservation_usd: 0 # ============================================================================= # Turnstile Configuration