diff --git a/backend/cmd/server/wire.go b/backend/cmd/server/wire.go index 06674251af..6834ad78df 100644 --- a/backend/cmd/server/wire.go +++ b/backend/cmd/server/wire.go @@ -104,6 +104,7 @@ func provideCleanup( subscriptionExpiry *service.SubscriptionExpiryService, usageCleanup *service.UsageCleanupService, idempotencyCleanup *service.IdempotencyCleanupService, + geminiCachedContent *service.GeminiCachedContentService, batchImageCleanup *service.BatchImageCleanupService, batchImageWorker *service.BatchImageWorkerRuntime, pricing *service.PricingService, @@ -244,6 +245,12 @@ func provideCleanup( } return nil }}, + {"GeminiCachedContentService", func() error { + if geminiCachedContent != nil { + geminiCachedContent.Stop() + } + return nil + }}, {"BatchImageCleanupService", func() error { if batchImageCleanup != nil { batchImageCleanup.Stop() diff --git a/backend/cmd/server/wire_gen.go b/backend/cmd/server/wire_gen.go index 2c2f0cfbc9..53cef32052 100644 --- a/backend/cmd/server/wire_gen.go +++ b/backend/cmd/server/wire_gen.go @@ -297,7 +297,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) { legacyEngine := securityaudit.NewLegacyModerationAdapter(contentModerationService) coordinator := securityaudit.NewCoordinator(legacyEngine, promptService) geminiCachedContentRepository := repository.NewGeminiCachedContentRepository(db) - geminiCachedContentService := service.NewGeminiCachedContentService(geminiCachedContentRepository, geminiMessagesCompatService, configConfig) + geminiCachedContentService := service.ProvideGeminiCachedContentService(geminiCachedContentRepository, geminiMessagesCompatService, configConfig) gatewayHandler := handler.ProvideGatewayHandler(gatewayService, openAIGatewayService, geminiMessagesCompatService, antigravityGatewayService, userService, concurrencyService, billingCacheService, usageService, apiKeyService, usageRecordWorkerPool, errorPassthroughService, contentModerationService, userMessageQueueService, configConfig, settingService, coordinator, geminiCachedContentService) openAIGatewayHandler := handler.ProvideOpenAIGatewayHandler(openAIGatewayService, pluginManager, concurrencyService, billingCacheService, apiKeyService, usageRecordWorkerPool, errorPassthroughService, contentModerationService, opsService, grokQuotaService, configConfig, coordinator, compositeRouteResolver) handlerSettingHandler := handler.ProvideSettingHandler(settingService, buildInfo, notificationEmailService) @@ -355,7 +355,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) { channelMonitorRunner := service.ProvideChannelMonitorRunner(channelMonitorService, settingService, channelMonitorQuotaFetcher) channelMonitorV2Aggregator := service.ProvideChannelMonitorV2Aggregator(channelMonitorV2Repository, db, settingService) userPlatformQuotaUsageFlusher := service.ProvideUserPlatformQuotaUsageFlusher(configConfig, billingCache, serviceUserPlatformQuotaRepository, timingWheelService) - v := provideCleanup(client, redisClient, opsMetricsCollector, opsAggregationService, opsAlertEvaluatorService, opsCleanupService, opsScheduledReportService, opsSystemLogSink, opsService, opsIngressRejectAggregator, apiKeyService, authCacheInvalidationWorker, schedulerSnapshotService, tokenRefreshService, accountExpiryService, cnProviderBalanceCheckService, openAICodexVersionSyncService, claudeCodeVersionSyncService, proxyExpiryService, subscriptionExpiryService, usageCleanupService, idempotencyCleanupService, batchImageCleanupService, batchImageWorkerRuntime, pricingService, emailQueueService, billingCacheService, usageRecordWorkerPool, subscriptionService, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, grokOAuthService, openAIGatewayService, scheduledTestRunnerService, backupService, paymentOrderExpiryService, channelMonitorRunner, channelMonitorV2Aggregator, userPlatformQuotaUsageFlusher, upstreamBillingProbeService, ollamaCloudUsageService, openCodeGoUsageService, auditLogService, openAIQuotaAutoResetService, promptService, pluginManager) + v := provideCleanup(client, redisClient, opsMetricsCollector, opsAggregationService, opsAlertEvaluatorService, opsCleanupService, opsScheduledReportService, opsSystemLogSink, opsService, opsIngressRejectAggregator, apiKeyService, authCacheInvalidationWorker, schedulerSnapshotService, tokenRefreshService, accountExpiryService, cnProviderBalanceCheckService, openAICodexVersionSyncService, claudeCodeVersionSyncService, proxyExpiryService, subscriptionExpiryService, usageCleanupService, idempotencyCleanupService, geminiCachedContentService, batchImageCleanupService, batchImageWorkerRuntime, pricingService, emailQueueService, billingCacheService, usageRecordWorkerPool, subscriptionService, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, grokOAuthService, openAIGatewayService, scheduledTestRunnerService, backupService, paymentOrderExpiryService, channelMonitorRunner, channelMonitorV2Aggregator, userPlatformQuotaUsageFlusher, upstreamBillingProbeService, ollamaCloudUsageService, openCodeGoUsageService, auditLogService, openAIQuotaAutoResetService, promptService, pluginManager) application := &Application{ Server: httpServer, PromptAudit: promptService, @@ -415,6 +415,7 @@ func provideCleanup( subscriptionExpiry *service.SubscriptionExpiryService, usageCleanup *service.UsageCleanupService, idempotencyCleanup *service.IdempotencyCleanupService, + geminiCachedContent *service.GeminiCachedContentService, batchImageCleanup *service.BatchImageCleanupService, batchImageWorker *service.BatchImageWorkerRuntime, pricing *service.PricingService, @@ -554,6 +555,12 @@ func provideCleanup( } return nil }}, + {"GeminiCachedContentService", func() error { + if geminiCachedContent != nil { + geminiCachedContent.Stop() + } + return nil + }}, {"BatchImageCleanupService", func() error { if batchImageCleanup != nil { batchImageCleanup.Stop() diff --git a/backend/cmd/server/wire_gen_test.go b/backend/cmd/server/wire_gen_test.go index d14ca19db6..bb56c6735f 100644 --- a/backend/cmd/server/wire_gen_test.go +++ b/backend/cmd/server/wire_gen_test.go @@ -74,6 +74,7 @@ func TestProvideCleanup_WithMinimalDependencies_NoPanic(t *testing.T) { subscriptionExpirySvc, &service.UsageCleanupService{}, idempotencyCleanupSvc, + service.NewGeminiCachedContentService(nil, nil, cfg), &service.BatchImageCleanupService{}, nil, // batchImageWorker pricingSvc, diff --git a/backend/internal/handler/gemini_cached_contents_billing_test.go b/backend/internal/handler/gemini_cached_contents_billing_test.go index a1fd414519..fdf0faea93 100644 --- a/backend/internal/handler/gemini_cached_contents_billing_test.go +++ b/backend/internal/handler/gemini_cached_contents_billing_test.go @@ -4,6 +4,7 @@ package handler import ( "context" + "encoding/json" "errors" "fmt" "io" @@ -11,6 +12,7 @@ import ( "net/http/httptest" "os" "path/filepath" + "strconv" "strings" "sync" "testing" @@ -32,13 +34,20 @@ const ( ) type cachedBillingRepo struct { - mu sync.Mutex - records map[string]*service.GeminiCachedContent - failUpdate bool + mu sync.Mutex + records map[string]*service.GeminiCachedContent + // mutations 按记录 id 保存修改意图。 + mutations map[int64]cachedBillingMutation + failCommit bool nextID int64 deleted []int64 } +type cachedBillingMutation struct { + token string + leaseUntil time.Time +} + func (r *cachedBillingRepo) Create(_ context.Context, record *service.GeminiCachedContent) error { r.mu.Lock() defer r.mu.Unlock() @@ -64,20 +73,83 @@ func (r *cachedBillingRepo) ListForOwner(context.Context, int64, int64, time.Tim return nil, nil } -func (r *cachedBillingRepo) UpdateExpireTime(_ context.Context, id int64, expireTime time.Time) (time.Time, error) { - r.mu.Lock() - defer r.mu.Unlock() - if r.failUpdate { - return time.Time{}, errors.New("db unavailable") +// liveLocked 返回未删除的记录；调用方须持有 mu。 +func (r *cachedBillingRepo) liveLocked(id int64) *service.GeminiCachedContent { + for _, deleted := range r.deleted { + if deleted == id { + return nil + } } for _, record := range r.records { if record.ID == id { - previous := record.ExpireTime - record.ExpireTime = expireTime - return previous, nil + return record } } - return time.Time{}, service.ErrGeminiCachedContentNotFound + return nil +} + +func (r *cachedBillingRepo) BeginMutation(_ context.Context, id int64, baseline time.Time, token string, lease time.Duration) (bool, error) { + r.mu.Lock() + defer r.mu.Unlock() + record := r.liveLocked(id) + if record == nil { + return false, service.ErrGeminiCachedContentNotFound + } + if _, pending := r.mutations[id]; pending || !record.ExpireTime.Equal(baseline) { + return false, nil + } + r.mutations[id] = cachedBillingMutation{token: token, leaseUntil: time.Now().Add(lease)} + return true, nil +} + +func (r *cachedBillingRepo) TakeOverStaleMutation(_ context.Context, id int64, token string, lease time.Duration) (*service.GeminiCachedContent, error) { + r.mu.Lock() + defer r.mu.Unlock() + record := r.liveLocked(id) + mutation, pending := r.mutations[id] + if record == nil || !pending || mutation.leaseUntil.After(time.Now()) { + return nil, nil + } + r.mutations[id] = cachedBillingMutation{token: token, leaseUntil: time.Now().Add(lease)} + cp := *record + return &cp, nil +} + +func (r *cachedBillingRepo) CommitMutation(_ context.Context, id int64, token string, expireTime time.Time) error { + r.mu.Lock() + defer r.mu.Unlock() + if r.failCommit { + return errors.New("db unavailable") + } + record := r.liveLocked(id) + mutation, pending := r.mutations[id] + if record == nil || !pending || mutation.token != token { + return service.ErrGeminiCachedContentMutationConflict + } + record.ExpireTime = expireTime + delete(r.mutations, id) + return nil +} + +func (r *cachedBillingRepo) ReleaseMutation(_ context.Context, id int64, token string) error { + r.mu.Lock() + defer r.mu.Unlock() + if mutation, pending := r.mutations[id]; pending && mutation.token == token { + delete(r.mutations, id) + } + return nil +} + +func (r *cachedBillingRepo) ListStaleMutations(_ context.Context, limit int) ([]int64, error) { + r.mu.Lock() + defer r.mu.Unlock() + var ids []int64 + for id, mutation := range r.mutations { + if len(ids) < limit && mutation.leaseUntil.Before(time.Now()) && r.liveLocked(id) != nil { + ids = append(ids, id) + } + } + return ids, nil } func (r *cachedBillingRepo) SoftDelete(_ context.Context, id int64) error { @@ -87,6 +159,23 @@ func (r *cachedBillingRepo) SoftDelete(_ context.Context, id int64) error { return nil } +// pendingMutation 返回记录是否有未清除的修改意图。 +func (r *cachedBillingRepo) pendingMutation(id int64) bool { + r.mu.Lock() + defer r.mu.Unlock() + _, pending := r.mutations[id] + return pending +} + +// expireLease 让记录上的修改意图租约立即过期。 +func (r *cachedBillingRepo) expireLease(id int64) { + r.mu.Lock() + defer r.mu.Unlock() + mutation := r.mutations[id] + mutation.leaseUntil = time.Now().Add(-time.Second) + r.mutations[id] = mutation +} + type cachedBillingAccountRepo struct { service.AccountRepository accounts map[int64]*service.Account @@ -106,9 +195,12 @@ type cachedBillingUpstream struct { mu sync.Mutex createBody string calls []string + // patchTTLs 按顺序记录 PATCH 请求体中的 ttl。 + patchTTLs []string beforePost func(n int) - // failures 按请求方法返回上游错误响应。 - failures map[string]cachedBillingUpstreamFailure + // failures 按请求方法返回上游错误响应；transportErrors 按请求方法返回传输层错误（上游结果未知）。 + failures map[string]cachedBillingUpstreamFailure + transportErrors map[string]error } type cachedBillingUpstreamFailure struct { @@ -117,11 +209,24 @@ type cachedBillingUpstreamFailure struct { } func (u *cachedBillingUpstream) Do(req *http.Request, _ string, _ int64, _ int) (*http.Response, error) { + var patch struct { + TTL string `json:"ttl"` + } + if req.Method == http.MethodPatch { + _ = json.NewDecoder(req.Body).Decode(&patch) + } u.mu.Lock() u.calls = append(u.calls, req.Method+" "+req.URL.Path) n := len(u.calls) + if req.Method == http.MethodPatch { + u.patchTTLs = append(u.patchTTLs, patch.TTL) + } failure, failed := u.failures[req.Method] + transportErr := u.transportErrors[req.Method] u.mu.Unlock() + if transportErr != nil { + return nil, transportErr + } if failed { return &http.Response{StatusCode: failure.status, Header: http.Header{"Content-Type": []string{"application/json"}}, Body: io.NopCloser(strings.NewReader(failure.body))}, nil @@ -134,7 +239,8 @@ func (u *cachedBillingUpstream) Do(req *http.Request, _ string, _ int64, _ int) if u.beforePost != nil { u.beforePost(n) } - body = fmt.Sprintf(`{"name":"cachedContents/up1","expireTime":%q}`, time.Now().Add(24*time.Hour).UTC().Format(time.RFC3339Nano)) + seconds, _ := strconv.ParseInt(strings.TrimSuffix(patch.TTL, "s"), 10, 64) + body = fmt.Sprintf(`{"name":"cachedContents/up1","expireTime":%q}`, time.Now().Add(time.Duration(seconds)*time.Second).UTC().Format(time.RFC3339Nano)) } return &http.Response{StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{"application/json"}}, Body: io.NopCloser(strings.NewReader(body))}, nil @@ -240,7 +346,7 @@ func newCachedBillingEnv(t *testing.T, balance float64, accountIDs ...int64) *ca billing, nil, &service.BillingCacheService{}, nil, nil, &service.DeferredService{}, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil) upstream := &cachedBillingUpstream{} compat := service.NewGeminiMessagesCompatService(&cachedBillingAccountRepo{accounts: accounts}, nil, nil, nil, nil, nil, upstream, nil, cfg) - repo := &cachedBillingRepo{records: map[string]*service.GeminiCachedContent{}} + repo := &cachedBillingRepo{records: map[string]*service.GeminiCachedContent{}, mutations: map[int64]cachedBillingMutation{}} h := &GatewayHandler{ gatewayService: gw, billingCacheService: bcs, @@ -377,37 +483,258 @@ func TestGeminiCachedContentsCreate_UnusableUpstreamResponseCleansUpWithoutRetry require.Empty(t, env.usage.totals()) } -func TestGeminiCachedContentsPatch_ConcurrentExtensionsChargeOnlyActualProgress(t *testing.T) { - env := newCachedBillingEnv(t, 1000, 900) - original := time.Now().Add(time.Hour) - env.seed(original) - env.upstream.beforePost = func(n int) { - if n == 1 { - // 第一个 PATCH 的上游调用返回前，第二个 PATCH 完整执行（读到的仍是旧到期时间）。 - require.Equal(t, http.StatusOK, env.patch().Code) - } - } - require.Equal(t, http.StatusOK, env.patch().Code) - - totals := env.waitForUsage(t, 1) - final := env.repo.records["abc"].ExpireTime - fair := cachedBillingTokens * cachedBillingStoragePrice * final.Sub(original).Hours() - sum := 0.0 - for _, v := range totals { - sum += v - } - require.InDelta(t, fair, sum, cachedBillingTokens*cachedBillingStoragePrice*(time.Minute.Hours()), "两次延长合计只收实际推进的时长") +func (e *cachedBillingEnv) patchTTL(id, ttl string) *httptest.ResponseRecorder { + return e.serve(http.MethodPatch, "/v1beta/cachedContents/"+id+"?updateMask=ttl", fmt.Sprintf(`{"ttl":%q}`, ttl), + gin.Params{{Key: "cacheID", Value: id}}, e.h.GeminiCachedContentsPatch) } -func TestGeminiCachedContentsPatch_PersistFailureDoesNotCharge(t *testing.T) { +// requireTTLNear 断言发往上游的 ttl（"Ns"）与 want 相差不超过 5 秒。 +func requireTTLNear(t *testing.T, ttl string, want time.Duration) { + t.Helper() + seconds, err := strconv.ParseInt(strings.TrimSuffix(ttl, "s"), 10, 64) + require.NoError(t, err, ttl) + require.InDelta(t, want.Seconds(), float64(seconds), 5, ttl) +} + +// 同一缓存的修改串行执行：第一个 PATCH 的上游调用返回前，第二个 PATCH 直接得到 409 且不调用上游， +// 本地到期时间与上游最后生效的修改一致，只按第一个 PATCH 实际推进的时长收费。 +func TestGeminiCachedContentsPatch_ConcurrentUpdateConflictsAndKeepsUpstreamInSync(t *testing.T) { env := newCachedBillingEnv(t, 1000, 900) env.seed(time.Now().Add(time.Hour)) - env.repo.failUpdate = true + env.upstream.beforePost = func(n int) { + if n == 1 { + rec := env.patchTTL("abc", "86400s") + require.Equal(t, http.StatusConflict, rec.Code, rec.Body.String()) + require.Contains(t, rec.Body.String(), `"ABORTED"`) + } + } + + rec := env.patchTTL("abc", "7200s") + require.Equal(t, http.StatusOK, rec.Code, rec.Body.String()) + require.Equal(t, 1, env.upstream.methodCalls(http.MethodPatch), "冲突的修改不调用上游") + require.WithinDuration(t, time.Now().Add(2*time.Hour), env.repo.records["abc"].ExpireTime, 5*time.Second, "本地到期时间与上游一致") + require.False(t, env.repo.pendingMutation(100)) + totals := env.waitForUsage(t, 1) + require.Len(t, totals, 1) + require.InDelta(t, cachedBillingTokens*cachedBillingStoragePrice, totals[0], cachedBillingTokens*cachedBillingStoragePrice*time.Minute.Hours(), "只收推进的 1 小时") +} + +// 上游已修改但本地写回失败：不计费并保留修改意图；租约内的修改返回 409， +// 租约过期后后台恢复把上游有效期回滚到修改前的到期时间并清除意图。 +func TestGeminiCachedContentsPatch_CommitFailureRollsBackUpstreamWithoutCharge(t *testing.T) { + env := newCachedBillingEnv(t, 1000, 900) + baseline := time.Now().Add(time.Hour) + env.seed(baseline) + env.repo.failCommit = true + + require.Equal(t, http.StatusInternalServerError, env.patch().Code) + require.Equal(t, 1, env.upstream.methodCalls(http.MethodPatch)) + require.True(t, env.repo.pendingMutation(100), "上游可能已修改：保留修改意图") + require.Empty(t, env.usage.totals(), "未提交的修改不收费") + + env.repo.failCommit = false + require.Equal(t, http.StatusConflict, env.patch().Code, "租约内的修改被拒绝") + require.Equal(t, 1, env.upstream.methodCalls(http.MethodPatch)) + + env.repo.expireLease(100) + recovered, err := env.h.geminiCachedContentService.RecoverStaleMutations(context.Background()) + require.NoError(t, err) + require.Equal(t, 1, recovered) + require.Equal(t, 2, env.upstream.methodCalls(http.MethodPatch)) + requireTTLNear(t, env.upstream.patchTTLs[1], time.Until(baseline)) + require.False(t, env.repo.pendingMutation(100)) + require.True(t, env.repo.records["abc"].ExpireTime.Equal(baseline), "本地到期时间保持修改前的值") + require.Empty(t, env.usage.totals()) +} + +// 租约过期的遗留修改意图在下一次修改前先回滚，随后的修改按修改前的到期时间计费。 +func TestGeminiCachedContentsPatch_StaleMutationRolledBackBeforeNextUpdate(t *testing.T) { + env := newCachedBillingEnv(t, 1000, 900) + baseline := time.Now().Add(time.Hour) + env.seed(baseline) + env.repo.failCommit = true + require.Equal(t, http.StatusInternalServerError, env.patch().Code) + env.repo.failCommit = false + env.repo.expireLease(100) rec := env.patch() - require.Equal(t, http.StatusInternalServerError, rec.Code) - require.Equal(t, 1, env.upstream.methodCalls(http.MethodPatch)) - require.Empty(t, env.usage.totals(), "落库失败不收费") + require.Equal(t, http.StatusOK, rec.Code, rec.Body.String()) + require.Equal(t, 3, env.upstream.methodCalls(http.MethodPatch), "失败的修改、回滚、本次修改") + requireTTLNear(t, env.upstream.patchTTLs[1], time.Until(baseline)) + require.Equal(t, "86400s", env.upstream.patchTTLs[2]) + totals := env.waitForUsage(t, 1) + require.Len(t, totals, 1) + require.InDelta(t, cachedBillingTokens*cachedBillingStoragePrice*23, totals[0], 0.01) +} + +// 上游结果未知（传输层错误）时保留修改意图等待回滚；上游明确拒绝时清除修改意图。两种情况都不收费。 +func TestGeminiCachedContentsPatch_MutationKeptOnlyWhenUpstreamOutcomeUnknown(t *testing.T) { + t.Run("transport error", func(t *testing.T) { + env := newCachedBillingEnv(t, 1000, 900) + env.seed(time.Now().Add(time.Hour)) + env.upstream.transportErrors = map[string]error{http.MethodPatch: errors.New("connection reset")} + + require.Equal(t, http.StatusBadGateway, env.patch().Code) + require.True(t, env.repo.pendingMutation(100)) + require.Empty(t, env.usage.totals()) + }) + t.Run("upstream rejected", func(t *testing.T) { + env := newCachedBillingEnv(t, 1000, 900) + env.seed(time.Now().Add(time.Hour)) + env.upstream.failures = map[string]cachedBillingUpstreamFailure{http.MethodPatch: {status: http.StatusBadRequest, + body: `{"error":{"code":400,"message":"invalid ttl","status":"INVALID_ARGUMENT"}}`}} + + require.Equal(t, http.StatusBadRequest, env.patch().Code) + require.False(t, env.repo.pendingMutation(100)) + require.Empty(t, env.usage.totals()) + }) +} + +// 修改前的到期时间已过：恢复时删除上游缓存与本地记录。 +func TestGeminiCachedContentMutationRecovery_DeletesUpstreamWhenBaselineExpired(t *testing.T) { + env := newCachedBillingEnv(t, 1000, 900) + env.seed(time.Now().Add(-time.Minute)) + env.repo.mutations[100] = cachedBillingMutation{token: "stale", leaseUntil: time.Now().Add(-time.Second)} + + recovered, err := env.h.geminiCachedContentService.RecoverStaleMutations(context.Background()) + require.NoError(t, err) + require.Equal(t, 1, recovered) + require.Equal(t, 1, env.upstream.methodCalls(http.MethodDelete)) + require.Zero(t, env.upstream.methodCalls(http.MethodPatch)) + require.Equal(t, []int64{100}, env.repo.deleted) +} + +// cachedBillingReservationCache 是启用在途预留时的计费缓存桩：余额与订阅用量取固定值，各维度预留按 key 记账。 +type cachedBillingReservationCache struct { + service.BillingCache + mu sync.Mutex + balance float64 + sub *service.SubscriptionCacheData + reserved map[string]map[string]float64 +} + +func (c *cachedBillingReservationCache) GetUserBalance(context.Context, int64) (float64, error) { + return c.balance, nil +} + +func (c *cachedBillingReservationCache) GetSubscriptionCache(context.Context, int64, int64) (*service.SubscriptionCacheData, error) { + return c.sub, nil +} + +func (c *cachedBillingReservationCache) reserveLocked(key, requestID string, amount float64) float64 { + if c.reserved == nil { + c.reserved = map[string]map[string]float64{} + } + if c.reserved[key] == nil { + c.reserved[key] = map[string]float64{} + } + others := 0.0 + for _, v := range c.reserved[key] { + others += v + } + c.reserved[key][requestID] = amount + return others +} + +func (c *cachedBillingReservationCache) ReserveInflightBalance(_ context.Context, userID int64, requestID string, amount, balance float64, _ time.Duration) (bool, float64, error) { + c.mu.Lock() + defer c.mu.Unlock() + key := fmt.Sprintf("user:%d", userID) + others := c.reserveLocked(key, requestID, amount) + if len(c.reserved[key]) > 1 && balance-others < amount { + delete(c.reserved[key], requestID) + return false, others, nil + } + return true, others, nil +} + +func (c *cachedBillingReservationCache) ReleaseInflightBalance(_ context.Context, userID int64, requestID string) error { + c.mu.Lock() + defer c.mu.Unlock() + delete(c.reserved[fmt.Sprintf("user:%d", userID)], requestID) + return nil +} + +func (c *cachedBillingReservationCache) ReserveInflightScope(_ context.Context, scope, requestID string, amount float64, _ time.Duration) (float64, error) { + c.mu.Lock() + defer c.mu.Unlock() + return c.reserveLocked(scope, requestID, amount), nil +} + +func (c *cachedBillingReservationCache) ReleaseInflightScope(_ context.Context, scope, requestID string) error { + c.mu.Lock() + defer c.mu.Unlock() + delete(c.reserved[scope], requestID) + return nil +} + +func enableCachedBillingReservations(t *testing.T, env *cachedBillingEnv, cache *cachedBillingReservationCache) { + t.Helper() + cfg := &config.Config{} + cfg.Billing.InflightReservation.Enabled = true + bcs := service.NewBillingCacheService(cache, nil, nil, nil, nil, nil, cfg, nil) + t.Cleanup(bcs.Stop) + env.h.billingCacheService = bcs +} + +// 确定金额的续期费按 API Key 与余额维度预留：第一笔续期（约 $103.5）尚未入账时， +// 第二个缓存的续期超出 Key 剩余额度或余额，在调用上游前被拒绝。 +func TestGeminiCachedContentsPatch_UnsettledChargeHoldsKeyQuotaAndBalance(t *testing.T) { + for _, tc := range []struct { + name string + balance, quota float64 + }{ + {name: "api key quota", balance: 1000, quota: 110}, + {name: "balance", balance: 110}, + } { + t.Run(tc.name, func(t *testing.T) { + env := newCachedBillingEnv(t, tc.balance, 900) + env.seed(time.Now().Add(time.Hour)) + second := *env.repo.records["abc"] + second.ID, second.PublicID, second.UpstreamName = 101, "def", "cachedContents/up2" + env.repo.records["def"] = &second + env.apiKey.Quota = tc.quota + enableCachedBillingReservations(t, env, &cachedBillingReservationCache{balance: tc.balance}) + env.upstream.beforePost = func(n int) { + if n == 1 { + rec := env.patchTTL("def", "86400s") + require.NotEqual(t, http.StatusOK, rec.Code, rec.Body.String()) + } + } + + require.Equal(t, http.StatusOK, env.patch().Code) + require.Equal(t, 1, env.upstream.methodCalls(http.MethodPatch), "第二笔续期没有到达上游") + }) + } +} + +// 订阅模式的确定金额费用按订阅维度预留：第一笔未入账时第二笔超出日剩余额度被拒绝，释放后放行。 +func TestReserveGeminiCachedContentCharge_UnsettledChargeHoldsSubscriptionHeadroom(t *testing.T) { + env := newCachedBillingEnv(t, 1000, 900) + limit := 110.0 + env.group.SubscriptionType = service.SubscriptionTypeSubscription + env.group.DailyLimitUSD = &limit + enableCachedBillingReservations(t, env, &cachedBillingReservationCache{balance: 1000, + sub: &service.SubscriptionCacheData{Status: service.SubscriptionStatusActive, ExpiresAt: time.Now().Add(time.Hour)}}) + sub := &service.UserSubscription{ID: 1, UserID: env.apiKey.UserID, GroupID: env.group.ID, Status: service.SubscriptionStatusActive, + ExpiresAt: time.Now().Add(time.Hour)} + reserve := func() (func(), error) { + var done func() + var err error + env.serve(http.MethodPatch, "/", `{}`, nil, func(c *gin.Context) { + done, err = env.h.reserveGeminiCachedContentCharge(c, env.apiKey, sub, 103.5) + }) + return done, err + } + + firstDone, err := reserve() + require.NoError(t, err) + _, err = reserve() + require.ErrorIs(t, err, service.ErrDailyLimitExceeded) + firstDone() + secondDone, err := reserve() + require.NoError(t, err) + secondDone() } func TestGeminiCachedContentsPatch_UnaffordableExtensionRejectedBeforeUpstream(t *testing.T) { diff --git a/backend/internal/handler/gemini_cached_contents_handler.go b/backend/internal/handler/gemini_cached_contents_handler.go index ac99eaf1d5..2071c16129 100644 --- a/backend/internal/handler/gemini_cached_contents_handler.go +++ b/backend/internal/handler/gemini_cached_contents_handler.go @@ -395,6 +395,8 @@ func (h *GatewayHandler) GeminiCachedContentsGet(c *gin.Context) { // GeminiCachedContentsPatch PATCH /v1beta/cachedContents/:cacheID // 只允许修改有效期；延长部分在本次请求内按增量时长收取存储费，缩短不退费。 +// 同一缓存的修改串行执行：上游修改、本地写回与计费都在修改意图的租约内完成，并发的修改返回 409。 +// 上游结果未知或本地写回失败的修改不计费，租约过期后由恢复流程把上游有效期回滚到修改前的到期时间。 func (h *GatewayHandler) GeminiCachedContentsPatch(c *gin.Context) { apiKey, _, reqLog, ok := h.geminiCachedContentsPreamble(c, "handler.gemini_v1beta.cached_contents.patch") if !ok { @@ -455,10 +457,33 @@ func (h *GatewayHandler) GeminiCachedContentsPatch(c *gin.Context) { } setOpsSelectedAccount(c, account.ID, account.Platform) + mutation, err := svc.BeginMutation(c.Request.Context(), record) + if err != nil { + switch { + case errors.Is(err, service.ErrGeminiCachedContentNotFound): + googleError(c, http.StatusForbidden, service.GeminiCachedContentNotFoundMessage) + case errors.Is(err, service.ErrGeminiCachedContentMutationConflict): + reqLog.Info("gemini.cached_content.mutation_conflict", zap.Int64("account_id", account.ID), zap.Error(err)) + writeGeminiCachedContentMutationConflict(c) + default: + reqLog.Error("gemini.cached_content.begin_mutation_failed", zap.Int64("account_id", account.ID), zap.Error(err)) + googleError(c, http.StatusInternalServerError, "Failed to update cached content") + } + return + } upstream, err := svc.PatchUpstream(c.Request.Context(), account, record, ttl) if err != nil { var upErr *service.GeminiCachedContentUpstreamError - if errors.As(err, &upErr) { + upstreamRejected := errors.As(err, &upErr) + if upstreamRejected || errors.Is(err, service.ErrGeminiCachedContentRequestNotSent) { + // 上游有效期未变：清除修改意图。 + persistCtx, cancelPersist := geminiCachedContentPersistContext(c) + if releaseErr := svc.ReleaseMutation(persistCtx, mutation); releaseErr != nil { + reqLog.Warn("gemini.cached_content.release_mutation_failed", zap.Error(releaseErr)) + } + cancelPersist() + } + if upstreamRejected { if upErr.NotFound() { if forgetErr := svc.Forget(c.Request.Context(), record); forgetErr != nil { reqLog.Warn("gemini.cached_content.forget_failed", zap.Error(forgetErr)) @@ -473,19 +498,19 @@ func (h *GatewayHandler) GeminiCachedContentsPatch(c *gin.Context) { } expireTime := service.ClampGeminiCachedContentExpire(upstream.ExpireTime, svc.Now().Add(ttl)) persistCtx, cancelPersist := geminiCachedContentPersistContext(c) - previousExpire, err := svc.UpdateExpireTime(persistCtx, record, expireTime) + err = svc.CommitMutation(persistCtx, mutation, expireTime) cancelPersist() if err != nil { - reqLog.Error("gemini.cached_content.update_expire_failed", zap.Int64("account_id", account.ID), zap.Error(err)) - if errors.Is(err, service.ErrGeminiCachedContentNotFound) { - googleError(c, http.StatusForbidden, service.GeminiCachedContentNotFoundMessage) + reqLog.Error("gemini.cached_content.commit_mutation_failed", zap.Int64("account_id", account.ID), zap.Error(err)) + if errors.Is(err, service.ErrGeminiCachedContentMutationConflict) { + writeGeminiCachedContentMutationConflict(c) return } googleError(c, http.StatusInternalServerError, "Failed to persist cached content expiration") return } - // 按库中实际推进的时长收费：并发延长各自只为自己推进的部分付费。 - if extended := expireTime.Sub(previousExpire); extended > 0 { + // 提交前库中的到期时间一直是修改前的值：按本次实际推进的时长收费。 + if extended := expireTime.Sub(mutation.Baseline()); extended > 0 { h.submitGeminiCachedContentUsage(c, apiKey, account, subscription, pricingAt, service.GeminiCachedContentPatchUsageRequestID(record.PublicID), record.Model, record.UpstreamModelID(), 0, service.GeminiCachedContentStorageTokenHours(record.TotalTokenCount, extended), record.ChannelUsage, body) @@ -706,13 +731,11 @@ func (h *GatewayHandler) cleanupGeminiCachedContentUpstream(reqLog *zap.Logger, } // reserveGeminiCachedContentCharge 校验余额 / 订阅额度 / API Key 额度足以支付确定金额 amount（倍率后）， -// 并按该金额登记在途预留（不受预留上限封顶）。返回的 done 须在 handler 结束时调用。 +// 并按该金额在用户、订阅、API Key 维度登记在途预留（不受预留上限封顶，见 ReserveFixedCharge）。 +// 返回的 done 须在 handler 结束时调用；预留随计费任务交接，落库后释放。 func (h *GatewayHandler) reserveGeminiCachedContentCharge(c *gin.Context, apiKey *service.APIKey, subscription *service.UserSubscription, amount float64) (func(), error) { ctx := c.Request.Context() - if err := h.billingCacheService.CheckChargeAffordable(ctx, apiKey.User, apiKey, apiKey.Group, subscription, amount); err != nil { - return inflightNoop, err - } - res, err := h.billingCacheService.ReserveInflightCharge(ctx, apiKey.User, apiKey.Group, subscription, amount) + res, err := h.billingCacheService.ReserveFixedCharge(ctx, apiKey.User, apiKey, apiKey.Group, subscription, amount) if err != nil { return inflightNoop, err } @@ -723,6 +746,17 @@ func (h *GatewayHandler) reserveGeminiCachedContentCharge(c *gin.Context, apiKey return res.HandlerDone, nil } +// writeGeminiCachedContentMutationConflict 回写同一缓存的修改冲突（Google 并发冲突语义 ABORTED），客户端可稍后重试。 +func writeGeminiCachedContentMutationConflict(c *gin.Context) { + c.JSON(http.StatusConflict, gin.H{ + "error": gin.H{ + "code": http.StatusConflict, + "message": "Another update to this cached content is in progress, please retry later", + "status": "ABORTED", + }, + }) +} + func writeGeminiCachedContentBillingError(c *gin.Context, err error) { status, code, message, retryAfter := billingErrorDetails(err) if retryAfter > 0 { diff --git a/backend/internal/handler/gemini_cached_contents_handler_test.go b/backend/internal/handler/gemini_cached_contents_handler_test.go index 26bc5e7183..f9a1d2b493 100644 --- a/backend/internal/handler/gemini_cached_contents_handler_test.go +++ b/backend/internal/handler/gemini_cached_contents_handler_test.go @@ -52,8 +52,24 @@ func (r *geminiCachedContentHandlerRepoStub) ListForOwner(_ context.Context, api return out, nil } -func (r *geminiCachedContentHandlerRepoStub) UpdateExpireTime(context.Context, int64, time.Time) (time.Time, error) { - return time.Time{}, nil +func (r *geminiCachedContentHandlerRepoStub) BeginMutation(context.Context, int64, time.Time, string, time.Duration) (bool, error) { + return true, nil +} + +func (r *geminiCachedContentHandlerRepoStub) TakeOverStaleMutation(context.Context, int64, string, time.Duration) (*service.GeminiCachedContent, error) { + return nil, nil +} + +func (r *geminiCachedContentHandlerRepoStub) CommitMutation(context.Context, int64, string, time.Time) error { + return nil +} + +func (r *geminiCachedContentHandlerRepoStub) ReleaseMutation(context.Context, int64, string) error { + return nil +} + +func (r *geminiCachedContentHandlerRepoStub) ListStaleMutations(context.Context, int) ([]int64, error) { + return nil, nil } func (r *geminiCachedContentHandlerRepoStub) SoftDelete(context.Context, int64) error { diff --git a/backend/internal/repository/api_key_repo.go b/backend/internal/repository/api_key_repo.go index dd831283b7..1fa3049ca9 100644 --- a/backend/internal/repository/api_key_repo.go +++ b/backend/internal/repository/api_key_repo.go @@ -857,7 +857,7 @@ func (r *apiKeyRepository) ResetRateLimitWindows(ctx context.Context, id int64) // GetRateLimitData returns the current rate limit usage and window start times for an API key. func (r *apiKeyRepository) GetRateLimitData(ctx context.Context, id int64) (result *service.APIKeyRateLimitData, err error) { rows, err := r.sql.QueryContext(ctx, ` - SELECT usage_5h, usage_1d, usage_7d, window_5h_start, window_1d_start, window_7d_start + SELECT quota_used, usage_5h, usage_1d, usage_7d, window_5h_start, window_1d_start, window_7d_start FROM api_keys WHERE id = $1 AND deleted_at IS NULL`, id) @@ -873,7 +873,7 @@ func (r *apiKeyRepository) GetRateLimitData(ctx context.Context, id int64) (resu return nil, service.ErrAPIKeyNotFound } data := &service.APIKeyRateLimitData{} - if err := rows.Scan(&data.Usage5h, &data.Usage1d, &data.Usage7d, &data.Window5hStart, &data.Window1dStart, &data.Window7dStart); err != nil { + if err := rows.Scan(&data.QuotaUsed, &data.Usage5h, &data.Usage1d, &data.Usage7d, &data.Window5hStart, &data.Window1dStart, &data.Window7dStart); err != nil { return nil, err } return data, rows.Err() diff --git a/backend/internal/repository/api_key_repo_integration_test.go b/backend/internal/repository/api_key_repo_integration_test.go index ada0a201b3..9d41f269b6 100644 --- a/backend/internal/repository/api_key_repo_integration_test.go +++ b/backend/internal/repository/api_key_repo_integration_test.go @@ -464,6 +464,10 @@ func (s *APIKeyRepoSuite) TestIncrementQuotaUsed_Basic() { newQuota, err = s.repo.IncrementQuotaUsed(s.ctx, key.ID, 2.5) s.Require().NoError(err, "IncrementQuotaUsed second") s.Require().Equal(4.0, newQuota, "第二次递增后应为 4.0") + + data, err := s.repo.GetRateLimitData(s.ctx, key.ID) + s.Require().NoError(err, "GetRateLimitData") + s.Require().Equal(4.0, data.QuotaUsed, "用量读取带回总额度已用量") } func (s *APIKeyRepoSuite) TestIncrementQuotaUsed_NotFound() { diff --git a/backend/internal/repository/billing_inflight_cache.go b/backend/internal/repository/billing_inflight_cache.go index 898cf3b52a..8f83717ce6 100644 --- a/backend/internal/repository/billing_inflight_cache.go +++ b/backend/internal/repository/billing_inflight_cache.go @@ -25,6 +25,20 @@ func billingInflightKeys(userID int64) (string, string) { fmt.Sprintf("%s{%d}", billingInflightAmtKeyPrefix, userID) } +// 订阅 / API Key 维度在途预留（确定金额费用）： +// - billing:inflight_scope:{scope} ZSET member=requestID score=过期时间(ms) +// - billing:inflight_scope_amt:{scope} HASH field=requestID value=预留金额(USD) +// +// 过期清理与释放语义与余额维度相同，不续期。 +const ( + billingInflightScopeKeyPrefix = "billing:inflight_scope:" + billingInflightScopeAmtKeyPrefix = "billing:inflight_scope_amt:" +) + +func billingInflightScopeKeys(scope string) (string, string) { + return billingInflightScopeKeyPrefix + "{" + scope + "}", billingInflightScopeAmtKeyPrefix + "{" + scope + "}" +} + var ( // KEYS: [1]=zset [2]=hash // ARGV: [1]=now_ms [2]=expire_at_ms [3]=member [4]=amount [5]=balance [6]=key_ttl_ms @@ -85,6 +99,28 @@ var ( redis.call('HDEL', KEYS[2], ARGV[1]) return 1 `) + + // KEYS: [1]=zset [2]=hash + // ARGV: [1]=now_ms [2]=expire_at_ms [3]=member [4]=amount [5]=key_ttl_ms + // 先清理已过期成员，再登记本次预留；返回登记前其他在途预留的合计（string），是否放行由调用方按权威用量判定。 + reserveInflightScopeScript = 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 + 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[5]) + redis.call('PEXPIRE', KEYS[2], ARGV[5]) + return tostring(sum) + `) ) // ReserveInflightBalance 实现 service.InflightBalanceReservationCache。 @@ -138,7 +174,35 @@ func (c *billingCache) RenewInflightBalance(ctx context.Context, userID int64, r return n == 1, nil } +// ReserveInflightScope 实现 service.InflightScopeReservationCache。 +func (c *billingCache) ReserveInflightScope(ctx context.Context, scope, requestID string, amount float64, ttl time.Duration) (float64, error) { + zkey, hkey := billingInflightScopeKeys(scope) + now := time.Now().UnixMilli() + ttlMs := ttl.Milliseconds() + if ttlMs <= 0 { + ttlMs = 1 + } + res, err := reserveInflightScopeScript.Run(ctx, c.rdb, []string{zkey, hkey}, + now, + now+ttlMs, + requestID, + strconv.FormatFloat(amount, 'f', -1, 64), + ttlMs, + ).Text() + if err != nil { + return 0, err + } + return strconv.ParseFloat(res, 64) +} + +// ReleaseInflightScope 实现 service.InflightScopeReservationCache。 +func (c *billingCache) ReleaseInflightScope(ctx context.Context, scope, requestID string) error { + zkey, hkey := billingInflightScopeKeys(scope) + return releaseInflightBalanceScript.Run(ctx, c.rdb, []string{zkey, hkey}, requestID).Err() +} + var ( _ service.InflightBalanceReservationCache = (*billingCache)(nil) _ service.InflightBalanceReservationRenewer = (*billingCache)(nil) + _ service.InflightScopeReservationCache = (*billingCache)(nil) ) diff --git a/backend/internal/repository/billing_inflight_cache_test.go b/backend/internal/repository/billing_inflight_cache_test.go index ddab7c4de0..2fd642fe38 100644 --- a/backend/internal/repository/billing_inflight_cache_test.go +++ b/backend/internal/repository/billing_inflight_cache_test.go @@ -285,3 +285,72 @@ func TestInflightReservation_RenewDoesNotResurrectExpiredMember(t *testing.T) { require.NoError(t, err) require.False(t, exists) } + +func TestInflightScopeReservation_ReturnsOthersAndReleases(t *testing.T) { + mr, cache, _ := newInflightTestEnv(t, true, 60) + ctx := context.Background() + + others, err := cache.ReserveInflightScope(ctx, "apikey:7", "r1", 3, time.Minute) + require.NoError(t, err) + require.Zero(t, others) + others, err = cache.ReserveInflightScope(ctx, "apikey:7", "r2", 2, time.Minute) + require.NoError(t, err) + require.InDelta(t, 3, others, 1e-9, "返回登记前其他在途预留的合计") + others, err = cache.ReserveInflightScope(ctx, "sub:7", "r3", 5, time.Minute) + require.NoError(t, err) + require.Zero(t, others, "不同维度互不影响") + + require.NoError(t, cache.ReleaseInflightScope(ctx, "apikey:7", "r1")) + require.NoError(t, cache.ReleaseInflightScope(ctx, "apikey:7", "r1")) + others, err = cache.ReserveInflightScope(ctx, "apikey:7", "r4", 1, time.Minute) + require.NoError(t, err) + require.InDelta(t, 2, others, 1e-9) + + mr.FastForward(2 * time.Minute) + others, err = cache.ReserveInflightScope(ctx, "apikey:7", "r5", 1, time.Minute) + require.NoError(t, err) + require.Zero(t, others, "过期预留不再计入") +} + +func TestReserveFixedCharge_ConcurrentChargesAdmitOnlyWhatAPIKeyQuotaCovers(t *testing.T) { + _, cache, svc := newInflightTestEnv(t, true, 60) + ctx := context.Background() + user := &service.User{ID: 42} + apiKey := &service.APIKey{ID: 9, UserID: user.ID, Quota: 1.0} + require.NoError(t, cache.SetUserBalance(ctx, user.ID, 100)) + + const n = 20 + var admitted atomic.Int32 + reservations := make(chan *service.InflightReservation, n) + var wg sync.WaitGroup + start := make(chan struct{}) + for i := 0; i < n; i++ { + wg.Add(1) + go func() { + defer wg.Done() + <-start + res, err := svc.ReserveFixedCharge(ctx, user, apiKey, nil, nil, 0.3) + if err != nil { + if !errors.Is(err, service.ErrAPIKeyQuotaExhausted) { + t.Errorf("unexpected error: %v", err) + } + return + } + admitted.Add(1) + reservations <- res + }() + } + close(start) + wg.Wait() + close(reservations) + + require.Equal(t, int32(3), admitted.Load(), "Key 额度 1.0、每笔 0.3：同时在途的不超过 3 笔") + for res := range reservations { + res.Release() + } + zkey, _ := billingInflightScopeKeys("apikey:9") + remaining, err := cache.rdb.ZCard(ctx, zkey).Result() + require.NoError(t, err) + require.Zero(t, remaining) + require.Equal(t, int64(0), inflightCount(t, cache, user.ID)) +} diff --git a/backend/internal/repository/gemini_cached_content_repo.go b/backend/internal/repository/gemini_cached_content_repo.go index ddb05bbdcf..42775448d7 100644 --- a/backend/internal/repository/gemini_cached_content_repo.go +++ b/backend/internal/repository/gemini_cached_content_repo.go @@ -4,6 +4,7 @@ import ( "context" "database/sql" "encoding/json" + "errors" "time" "github.com/Wei-Shaw/sub2api/internal/service" @@ -78,20 +79,88 @@ func (r *geminiCachedContentRepository) ListForOwner(ctx context.Context, apiKey return out, rows.Err() } -func (r *geminiCachedContentRepository) UpdateExpireTime(ctx context.Context, id int64, expireTime time.Time) (time.Time, error) { - var previous time.Time - err := r.sql.QueryRowContext(ctx, `WITH prev AS ( - SELECT id, expire_time FROM gemini_cached_contents WHERE id = $1 AND deleted_at IS NULL FOR UPDATE -) -UPDATE gemini_cached_contents AS g - SET expire_time = $2, updated_at = NOW() - FROM prev - WHERE g.id = prev.id - RETURNING prev.expire_time`, id, expireTime).Scan(&previous) +func (r *geminiCachedContentRepository) BeginMutation(ctx context.Context, id int64, baseline time.Time, token string, lease time.Duration) (bool, error) { + result, err := r.sql.ExecContext(ctx, `UPDATE gemini_cached_contents + SET mutation_token = $2, mutation_lease_until = NOW() + make_interval(secs => $3::double precision) + WHERE id = $1 AND deleted_at IS NULL AND mutation_token IS NULL AND expire_time = $4`, + id, token, lease.Seconds(), baseline) if err != nil { - return time.Time{}, translatePersistenceError(err, service.ErrGeminiCachedContentNotFound, nil) + return false, err } - return previous, nil + affected, err := result.RowsAffected() + if err != nil { + return false, err + } + if affected > 0 { + return true, nil + } + var exists bool + if err := r.sql.QueryRowContext(ctx, `SELECT EXISTS ( + SELECT 1 FROM gemini_cached_contents WHERE id = $1 AND deleted_at IS NULL)`, id).Scan(&exists); err != nil { + return false, err + } + if !exists { + return false, service.ErrGeminiCachedContentNotFound + } + return false, nil +} + +func (r *geminiCachedContentRepository) TakeOverStaleMutation(ctx context.Context, id int64, token string, lease time.Duration) (*service.GeminiCachedContent, error) { + record, err := scanGeminiCachedContent(r.sql.QueryRowContext(ctx, `UPDATE gemini_cached_contents + SET mutation_token = $2, mutation_lease_until = NOW() + make_interval(secs => $3::double precision) + WHERE id = $1 AND deleted_at IS NULL AND mutation_token IS NOT NULL AND mutation_lease_until < NOW() + RETURNING `+geminiCachedContentColumns, id, token, lease.Seconds())) + if errors.Is(err, sql.ErrNoRows) { + return nil, nil + } + if err != nil { + return nil, err + } + return record, nil +} + +func (r *geminiCachedContentRepository) CommitMutation(ctx context.Context, id int64, token string, expireTime time.Time) error { + result, err := r.sql.ExecContext(ctx, `UPDATE gemini_cached_contents + SET expire_time = $3, updated_at = NOW(), mutation_token = NULL, mutation_lease_until = NULL + WHERE id = $1 AND deleted_at IS NULL AND mutation_token = $2`, id, token, expireTime) + if err != nil { + return err + } + affected, err := result.RowsAffected() + if err != nil { + return err + } + if affected == 0 { + return service.ErrGeminiCachedContentMutationConflict + } + return nil +} + +func (r *geminiCachedContentRepository) ReleaseMutation(ctx context.Context, id int64, token string) error { + _, err := r.sql.ExecContext(ctx, `UPDATE gemini_cached_contents + SET mutation_token = NULL, mutation_lease_until = NULL + WHERE id = $1 AND mutation_token = $2`, id, token) + return err +} + +func (r *geminiCachedContentRepository) ListStaleMutations(ctx context.Context, limit int) ([]int64, error) { + rows, err := r.sql.QueryContext(ctx, `SELECT id FROM gemini_cached_contents + WHERE mutation_token IS NOT NULL AND deleted_at IS NULL AND mutation_lease_until < NOW() + ORDER BY mutation_lease_until + LIMIT $1`, limit) + if err != nil { + return nil, err + } + defer func() { _ = rows.Close() }() + var ids []int64 + for rows.Next() { + var id int64 + if err := rows.Scan(&id); err != nil { + return nil, err + } + ids = append(ids, id) + } + return ids, rows.Err() } func (r *geminiCachedContentRepository) SoftDelete(ctx context.Context, id int64) error { diff --git a/backend/internal/repository/gemini_cached_content_repo_integration_test.go b/backend/internal/repository/gemini_cached_content_repo_integration_test.go index fc537a45f3..89399552c2 100644 --- a/backend/internal/repository/gemini_cached_content_repo_integration_test.go +++ b/backend/internal/repository/gemini_cached_content_repo_integration_test.go @@ -61,22 +61,58 @@ func TestGeminiCachedContentRepository_Lifecycle(t *testing.T) { require.Len(t, page, 1) require.Equal(t, first.PublicID, page[0].PublicID) - newExpire := now.Add(3 * time.Hour) - previous, err := repo.UpdateExpireTime(ctx, first.ID, newExpire) + lease := 2 * time.Minute + acquired, err := repo.BeginMutation(ctx, first.ID, first.ExpireTime, "t1", lease) require.NoError(t, err) - require.True(t, previous.Equal(first.ExpireTime), "返回更新前的到期时间") + require.True(t, acquired) + acquired, err = repo.BeginMutation(ctx, first.ID, first.ExpireTime, "t2", lease) + require.NoError(t, err) + require.False(t, acquired, "有进行中的修改时不能再登记") + taken, err := repo.TakeOverStaleMutation(ctx, first.ID, "t2", lease) + require.NoError(t, err) + require.Nil(t, taken, "租约未过期不能接管") + require.ErrorIs(t, repo.CommitMutation(ctx, first.ID, "t2", now), service.ErrGeminiCachedContentMutationConflict) + + newExpire := now.Add(3 * time.Hour) + require.NoError(t, repo.CommitMutation(ctx, first.ID, "t1", newExpire)) got, err = repo.GetForOwner(ctx, apiKeyID, first.PublicID) require.NoError(t, err) require.True(t, got.ExpireTime.Equal(newExpire)) - previous, err = repo.UpdateExpireTime(ctx, first.ID, newExpire.Add(time.Hour)) + require.ErrorIs(t, repo.CommitMutation(ctx, first.ID, "t1", now), service.ErrGeminiCachedContentMutationConflict, "提交后修改意图已清除") + + acquired, err = repo.BeginMutation(ctx, first.ID, first.ExpireTime, "t3", lease) require.NoError(t, err) - require.True(t, previous.Equal(newExpire), "连续更新各自拿到自己的前值") - newExpire = newExpire.Add(time.Hour) + require.False(t, acquired, "到期时间已被其他修改改变") + acquired, err = repo.BeginMutation(ctx, first.ID, newExpire, "t3", lease) + require.NoError(t, err) + require.True(t, acquired) + stale, err := repo.ListStaleMutations(ctx, 10) + require.NoError(t, err) + require.NotContains(t, stale, first.ID) + + _, err = tx.ExecContext(ctx, `UPDATE gemini_cached_contents SET mutation_lease_until = NOW() - INTERVAL '1 second' WHERE id = $1`, first.ID) + require.NoError(t, err) + stale, err = repo.ListStaleMutations(ctx, 10) + require.NoError(t, err) + require.Contains(t, stale, first.ID) + taken, err = repo.TakeOverStaleMutation(ctx, first.ID, "t4", lease) + require.NoError(t, err) + require.NotNil(t, taken) + require.True(t, taken.ExpireTime.Equal(newExpire), "接管返回修改前的到期时间") + require.ErrorIs(t, repo.CommitMutation(ctx, first.ID, "t3", now), service.ErrGeminiCachedContentMutationConflict, "被接管后原持有者不能提交") + require.NoError(t, repo.ReleaseMutation(ctx, first.ID, "t3")) + acquired, err = repo.BeginMutation(ctx, first.ID, newExpire, "t5", lease) + require.NoError(t, err) + require.False(t, acquired, "原持有者释放不影响接管者") + require.NoError(t, repo.ReleaseMutation(ctx, first.ID, "t4")) + stale, err = repo.ListStaleMutations(ctx, 10) + require.NoError(t, err) + require.NotContains(t, stale, first.ID) require.NoError(t, repo.SoftDelete(ctx, first.ID)) _, err = repo.GetForOwner(ctx, apiKeyID, first.PublicID) require.ErrorIs(t, err, service.ErrGeminiCachedContentNotFound) - _, err = repo.UpdateExpireTime(ctx, first.ID, newExpire) + _, err = repo.BeginMutation(ctx, first.ID, newExpire, "t6", lease) require.ErrorIs(t, err, service.ErrGeminiCachedContentNotFound) list, err = repo.ListForOwner(ctx, apiKeyID, 42, now, 0, 10) require.NoError(t, err) diff --git a/backend/internal/service/api_key_service.go b/backend/internal/service/api_key_service.go index b00b03d43f..e20a2697eb 100644 --- a/backend/internal/service/api_key_service.go +++ b/backend/internal/service/api_key_service.go @@ -130,6 +130,8 @@ type apiKeyAllByUserIDLister interface { // APIKeyRateLimitData holds rate limit usage and window state for an API key. type APIKeyRateLimitData struct { + // QuotaUsed 是总额度已用量（quota_used）。 + QuotaUsed float64 Usage5h float64 Usage1d float64 Usage7d float64 diff --git a/backend/internal/service/billing_cache_service.go b/backend/internal/service/billing_cache_service.go index 70f75a7ae8..ea22efe2eb 100644 --- a/backend/internal/service/billing_cache_service.go +++ b/backend/internal/service/billing_cache_service.go @@ -790,48 +790,106 @@ func (s *BillingCacheService) CheckBillingEligibility(ctx context.Context, user return nil } -// CheckChargeAffordable 校验扣费前即可确定金额的费用（倍率后 USD，如显式缓存存储费）不会透支： -// 余额模式要求缓存余额覆盖 amount，订阅模式要求日 / 周 / 月剩余额度覆盖 amount， -// API Key 总额度与限速窗口同样叠加 amount 检查。须在 CheckBillingEligibility 通过后调用；简易模式不校验。 -func (s *BillingCacheService) CheckChargeAffordable(ctx context.Context, user *User, apiKey *APIKey, group *Group, subscription *UserSubscription, amount float64) error { +// ReserveFixedCharge 为扣费前即可确定金额的费用（倍率后 USD，如显式缓存存储费）做准入并登记在途预留。 +// +// 余额模式要求缓存余额覆盖 amount 及用户的其他在途预留；订阅模式要求日 / 周 / 月剩余额度覆盖 amount； +// API Key 总额度与限速窗口同样叠加 amount。订阅与 API Key 的用量取库中权威值，并叠加同一订阅 / Key 上其他 +// 在途预留：先登记本次预留再读用量，并发的两笔费用中后读用量的一方必然看到另一方（预留或已落库的扣费）。 +// 返回的句柄可能为 nil（未登记任何预留）；调用方须在 handler 结束时调用 HandlerDone，预留在计费任务落库后释放。 +// 须在 CheckBillingEligibility 通过后调用；简易模式不校验。 +func (s *BillingCacheService) ReserveFixedCharge(ctx context.Context, user *User, apiKey *APIKey, group *Group, subscription *UserSubscription, amount float64) (*InflightReservation, error) { if user == nil || amount <= 0 || math.IsNaN(amount) || math.IsInf(amount, 0) { - return nil + return nil, nil } if s.cfg != nil && s.cfg.RunMode == config.RunModeSimple { - return nil + return nil, nil } - if group != nil && group.IsSubscriptionType() && subscription != nil { - subData, err := s.GetSubscriptionStatus(ctx, user.ID, group.ID) - if err != nil { - return ErrBillingServiceUnavailable.WithCause(err) - } - if group.HasDailyLimit() && subData.DailyUsage+amount > *group.DailyLimitUSD { - return ErrDailyLimitExceeded - } - if group.HasWeeklyLimit() && subData.WeeklyUsage+amount > *group.WeeklyLimitUSD { - return ErrWeeklyLimitExceeded - } - if group.HasMonthlyLimit() && subData.MonthlyUsage+amount > *group.MonthlyLimitUSD { - return ErrMonthlyLimitExceeded - } - } else { + subscriptionBill := group != nil && group.IsSubscriptionType() && subscription != nil + if !subscriptionBill { balance, err := s.GetUserBalance(ctx, user.ID) if err != nil { - return ErrBillingServiceUnavailable.WithCause(err) + return nil, ErrBillingServiceUnavailable.WithCause(err) } if balance < amount { - return ErrInsufficientBalance + return nil, ErrInsufficientBalance } } - if apiKey != nil { - if apiKey.Quota > 0 && apiKey.QuotaUsed+amount > apiKey.Quota { + res, err := s.reserveInflight(ctx, user, group, subscription, amount, false, false) + if err != nil { + return nil, err + } + var subscriptionPending, apiKeyPending float64 + if subscriptionBill && (group.HasDailyLimit() || group.HasWeeklyLimit() || group.HasMonthlyLimit()) { + res, subscriptionPending = s.reserveInflightScope(ctx, res, user.ID, inflightSubscriptionScope(subscription.ID), amount) + } + if apiKey != nil && (apiKey.Quota > 0 || apiKey.HasRateLimits()) { + res, apiKeyPending = s.reserveInflightScope(ctx, res, user.ID, inflightAPIKeyScope(apiKey.ID), amount) + } + if subscriptionBill { + err = s.checkFixedChargeSubscription(ctx, user.ID, group, amount+subscriptionPending) + } + if err == nil && apiKey != nil { + err = s.checkFixedChargeAPIKey(ctx, apiKey, amount+apiKeyPending) + } + if err != nil { + res.Release() + return nil, err + } + return res, nil +} + +// checkFixedChargeSubscription 按库中订阅用量校验日 / 周 / 月剩余额度覆盖 pending。 +func (s *BillingCacheService) checkFixedChargeSubscription(ctx context.Context, userID int64, group *Group, pending float64) error { + read := s.GetSubscriptionStatus + if s.subRepo != nil { + read = s.getSubscriptionFromDB + } + subData, err := read(ctx, userID, group.ID) + if err != nil { + return ErrBillingServiceUnavailable.WithCause(err) + } + if group.HasDailyLimit() && subData.DailyUsage+pending > *group.DailyLimitUSD { + return ErrDailyLimitExceeded + } + if group.HasWeeklyLimit() && subData.WeeklyUsage+pending > *group.WeeklyLimitUSD { + return ErrWeeklyLimitExceeded + } + if group.HasMonthlyLimit() && subData.MonthlyUsage+pending > *group.MonthlyLimitUSD { + return ErrMonthlyLimitExceeded + } + return nil +} + +// checkFixedChargeAPIKey 按库中 Key 用量校验总额度与限速窗口覆盖 pending；未接入仓储时使用请求上下文与限速缓存中的用量。 +func (s *BillingCacheService) checkFixedChargeAPIKey(ctx context.Context, apiKey *APIKey, pending float64) error { + if apiKey.Quota <= 0 && !apiKey.HasRateLimits() { + return nil + } + if s.apiKeyRateLimitLoader == nil { + if apiKey.Quota > 0 && apiKey.QuotaUsed+pending > apiKey.Quota { return ErrAPIKeyQuotaExhausted } if apiKey.HasRateLimits() { - if err := s.checkAPIKeyRateLimitsWithPending(ctx, apiKey, amount); err != nil { - return err - } + return s.checkAPIKeyRateLimitsWithPending(ctx, apiKey, pending) } + return nil + } + data, err := s.apiKeyRateLimitLoader.GetRateLimitData(ctx, apiKey.ID) + if err != nil { + return ErrBillingServiceUnavailable.WithCause(err) + } + exceeds := func(usage, limit float64) bool { + return limit > 0 && usage+pending > limit + } + switch { + case exceeds(data.QuotaUsed, apiKey.Quota): + return ErrAPIKeyQuotaExhausted + case exceeds(data.EffectiveUsage5h(), apiKey.RateLimit5h): + return ErrAPIKeyRateLimit5hExceeded + case exceeds(data.EffectiveUsage1d(), apiKey.RateLimit1d): + return ErrAPIKeyRateLimit1dExceeded + case exceeds(data.EffectiveUsage7d(), apiKey.RateLimit7d): + return ErrAPIKeyRateLimit7dExceeded } return nil } diff --git a/backend/internal/service/billing_inflight_reservation.go b/backend/internal/service/billing_inflight_reservation.go index 714143d554..e63c9db339 100644 --- a/backend/internal/service/billing_inflight_reservation.go +++ b/backend/internal/service/billing_inflight_reservation.go @@ -3,6 +3,7 @@ package service import ( "context" "math" + "strconv" "strings" "sync" "sync/atomic" @@ -25,6 +26,16 @@ type InflightBalanceReservationCache interface { ReleaseInflightBalance(ctx context.Context, userID int64, requestID string) error } +// InflightScopeReservationCache 订阅 / API Key 维度在途预留的缓存能力（可选），用于扣费前即可确定金额的费用； +// 未实现时这两个维度不登记预留（fail-open），准入仍按库中用量校验。 +type InflightScopeReservationCache interface { + // ReserveInflightScope 原子地清理 scope 下的过期预留并登记 requestID 的预留（ttl 后自动失效）， + // 返回登记前该 scope 下其他在途预留的合计。 + ReserveInflightScope(ctx context.Context, scope, requestID string, amount float64, ttl time.Duration) (float64, error) + // ReleaseInflightScope 释放 scope 下 requestID 的预留（幂等）。 + ReleaseInflightScope(ctx context.Context, scope, requestID string) error +} + // InflightBalanceReservationRenewer 可选：续期仍存活的预留（流式长请求期间防止 TTL 过期）。 type InflightBalanceReservationRenewer interface { // RenewInflightBalance 若 requestID 仍存在则把其过期时间推迟到 now+ttl；返回是否仍存在。 @@ -54,11 +65,15 @@ const ( // 不再续期，所以最长持有时间受 TTL 约束（计费任务卡死时预留自动过期）。 // 所有方法对 nil 接收者安全。 type InflightReservation struct { + // cache 为 nil 表示未登记余额维度的预留。 cache InflightBalanceReservationCache userID int64 requestID string amount float64 ttl time.Duration + // scopes 是同一金额在订阅 / API Key 维度登记的预留（与余额维度共用 requestID），不续期。 + scopeCache InflightScopeReservationCache + scopes []string refs atomic.Int64 releaseOnce sync.Once @@ -116,8 +131,15 @@ func (r *InflightReservation) Release() { r.refs.Store(0) relCtx, relCancel := context.WithTimeout(context.Background(), inflightReservationReleaseTimeout) defer relCancel() - if err := r.cache.ReleaseInflightBalance(relCtx, r.userID, r.requestID); err != nil { - logger.LegacyPrintf("service.billing_cache", "Warning: inflight reservation release failed for user %d (expires by ttl): %v", r.userID, err) + if r.cache != nil { + if err := r.cache.ReleaseInflightBalance(relCtx, r.userID, r.requestID); err != nil { + logger.LegacyPrintf("service.billing_cache", "Warning: inflight reservation release failed for user %d (expires by ttl): %v", r.userID, err) + } + } + for _, scope := range r.scopes { + if err := r.scopeCache.ReleaseInflightScope(relCtx, scope, r.requestID); err != nil { + logger.LegacyPrintf("service.billing_cache", "Warning: inflight reservation release failed for %s (expires by ttl): %v", scope, err) + } } }) } @@ -235,10 +257,44 @@ func (s *BillingCacheService) ReserveInflight(ctx context.Context, user *User, g return s.reserveInflight(ctx, user, group, subscription, estimate, true, true) } -// ReserveInflightCharge 为扣费前即可确定金额的费用登记在途预留，金额不受 max_reservation_usd 封顶， -// 其余语义与 ReserveInflight 相同。 -func (s *BillingCacheService) ReserveInflightCharge(ctx context.Context, user *User, group *Group, subscription *UserSubscription, amount float64) (*InflightReservation, error) { - return s.reserveInflight(ctx, user, group, subscription, amount, true, false) +// inflightSubscriptionScope / inflightAPIKeyScope 是订阅 / API Key 维度在途预留的 scope 名。 +func inflightSubscriptionScope(subscriptionID int64) string { + return "sub:" + strconv.FormatInt(subscriptionID, 10) +} + +func inflightAPIKeyScope(apiKeyID int64) string { + return "apikey:" + strconv.FormatInt(apiKeyID, 10) +} + +// reserveInflightScope 在 scope 维度为 amount 登记在途预留并挂到 res 上（res 为 nil 时新建句柄）， +// 返回句柄与登记前该 scope 下其他在途预留的合计。开关关闭 / 缓存不支持 / Redis 执行失败时不登记（fail-open），合计按 0。 +func (s *BillingCacheService) reserveInflightScope(ctx context.Context, res *InflightReservation, userID int64, scope string, amount float64) (*InflightReservation, float64) { + cfg, ok := s.inflightReservationConfig() + if !ok { + return res, 0 + } + sc, ok := s.cache.(InflightScopeReservationCache) + if !ok || sc == nil { + return res, 0 + } + if res == nil { + ttl := defaultInflightReservationTTL + if cfg.TTLSeconds > 0 { + ttl = time.Duration(cfg.TTLSeconds) * time.Second + } + res = &InflightReservation{userID: userID, requestID: uuid.NewString(), amount: amount, ttl: ttl} + res.refs.Store(1) + } + reserveCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), inflightReservationReserveTimeout) + pending, err := sc.ReserveInflightScope(reserveCtx, scope, res.requestID, amount, res.ttl) + cancel() + if err != nil { + logger.LegacyPrintf("service.billing_cache", "Warning: inflight reservation failed for %s (fail-open): %v", scope, err) + return res, 0 + } + res.scopeCache = sc + res.scopes = append(res.scopes, scope) + return res, pending } func (s *BillingCacheService) reserveInflight(ctx context.Context, user *User, group *Group, subscription *UserSubscription, estimate float64, renew, capped bool) (*InflightReservation, error) { diff --git a/backend/internal/service/gemini_cached_content.go b/backend/internal/service/gemini_cached_content.go index 92157dffbb..36b3f79a71 100644 --- a/backend/internal/service/gemini_cached_content.go +++ b/backend/internal/service/gemini_cached_content.go @@ -16,12 +16,14 @@ import ( "slices" "strconv" "strings" + "sync" "time" "unicode/utf8" "github.com/Wei-Shaw/sub2api/internal/config" infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" "github.com/Wei-Shaw/sub2api/internal/pkg/geminicli" + "github.com/Wei-Shaw/sub2api/internal/pkg/logger" "github.com/tidwall/gjson" "github.com/tidwall/sjson" ) @@ -42,6 +44,17 @@ const ( geminiCachedContentUpstreamTimeout = 5 * time.Minute geminiCachedContentOfficialHost = "generativelanguage.googleapis.com" + // geminiCachedContentMutationUpstreamTimeout 是修改 / 删除上游缓存的单次调用超时。 + geminiCachedContentMutationUpstreamTimeout = 30 * time.Second + // geminiCachedContentMutationLease 是修改意图的租约时长，须覆盖持有者在租约内的全部操作 + // （一次上游修改与一次落库，或恢复流程的一次上游回滚与一次落库），租约过期后持有者不会再改动上游。 + geminiCachedContentMutationLease = 2 * time.Minute + // geminiCachedContentMutationRecoveryInterval / geminiCachedContentMutationRecoveryBatch 约束后台回滚恢复的节奏； + // geminiCachedContentMutationRecoveryTimeout 是恢复单条记录的落库超时，与一次上游回滚合计不超过租约。 + geminiCachedContentMutationRecoveryInterval = time.Minute + geminiCachedContentMutationRecoveryBatch = 50 + geminiCachedContentMutationRecoveryTimeout = time.Minute + // GeminiCachedContentListDefaultPageSize / GeminiCachedContentListMaxPageSize 约束 list 分页。 GeminiCachedContentListDefaultPageSize = 50 GeminiCachedContentListMaxPageSize = 1000 @@ -50,6 +63,10 @@ const ( var ( ErrGeminiCachedContentNotFound = infraerrors.New(http.StatusNotFound, "GEMINI_CACHED_CONTENT_NOT_FOUND", "gemini cached content not found") ErrGeminiCachedContentExists = infraerrors.New(http.StatusConflict, "GEMINI_CACHED_CONTENT_EXISTS", "gemini cached content already exists") + // ErrGeminiCachedContentMutationConflict 表示同一缓存有进行中的修改、到期时间已被其他修改改变，或修改意图已不属于调用方。 + ErrGeminiCachedContentMutationConflict = infraerrors.New(http.StatusConflict, "GEMINI_CACHED_CONTENT_MUTATION_CONFLICT", "another update to this cached content is in progress") + // ErrGeminiCachedContentRequestNotSent 表示上游缓存请求在发出前失败（地址、凭据等），上游状态未被改变。 + ErrGeminiCachedContentRequestNotSent = errors.New("gemini cached content request was not sent") geminiCachedContentPublicIDPattern = regexp.MustCompile(`^[a-z0-9]{1,64}$`) geminiCachedContentAIStudioName = regexp.MustCompile(`^cachedContents/[A-Za-z0-9_-]+$`) @@ -100,8 +117,18 @@ type GeminiCachedContentRepository interface { GetForOwner(ctx context.Context, apiKeyID int64, publicID string) (*GeminiCachedContent, error) // ListForOwner 按 id 倒序返回分组内 expire_time 晚于 activeAt 的记录；beforeID>0 时只返回 id 更小的记录。 ListForOwner(ctx context.Context, apiKeyID, groupID int64, activeAt time.Time, beforeID int64, limit int) ([]*GeminiCachedContent, error) - // UpdateExpireTime 原子地写入新的到期时间并返回写入前的值。 - UpdateExpireTime(ctx context.Context, id int64, expireTime time.Time) (time.Time, error) + // BeginMutation 在记录没有修改意图且 expire_time 仍为 baseline 时登记修改意图（token 与自数据库当前时间起的租约）， + // 返回是否登记成功；记录不存在或已删除时返回 ErrGeminiCachedContentNotFound。 + BeginMutation(ctx context.Context, id int64, baseline time.Time, token string, lease time.Duration) (bool, error) + // TakeOverStaleMutation 把记录上租约已过期的修改意图换成 token 与新租约并返回记录；没有可接管的意图时返回 nil。 + TakeOverStaleMutation(ctx context.Context, id int64, token string, lease time.Duration) (*GeminiCachedContent, error) + // CommitMutation 由 token 的持有者写入新的到期时间并清除修改意图；token 已失效或记录已删除时返回 + // ErrGeminiCachedContentMutationConflict。 + CommitMutation(ctx context.Context, id int64, token string, expireTime time.Time) error + // ReleaseMutation 由 token 的持有者清除修改意图，不改变到期时间。 + ReleaseMutation(ctx context.Context, id int64, token string) error + // ListStaleMutations 返回修改意图租约已过期的未删除记录 id。 + ListStaleMutations(ctx context.Context, limit int) ([]int64, error) SoftDelete(ctx context.Context, id int64) error } @@ -594,10 +621,48 @@ type GeminiCachedContentService struct { compat *GeminiMessagesCompatService cfg *config.Config now func() time.Time + + startOnce sync.Once + stopOnce sync.Once + stopCh chan struct{} } func NewGeminiCachedContentService(repo GeminiCachedContentRepository, compat *GeminiMessagesCompatService, cfg *config.Config) *GeminiCachedContentService { - return &GeminiCachedContentService{repo: repo, compat: compat, cfg: cfg, now: time.Now} + return &GeminiCachedContentService{repo: repo, compat: compat, cfg: cfg, now: time.Now, stopCh: make(chan struct{})} +} + +// Start 启动后台回滚恢复：定期处理租约已过期仍未清除的修改意图。 +func (s *GeminiCachedContentService) Start() { + if s == nil || s.repo == nil { + return + } + s.startOnce.Do(func() { go s.runMutationRecovery() }) +} + +// Stop 停止后台回滚恢复。 +func (s *GeminiCachedContentService) Stop() { + if s == nil { + return + } + s.stopOnce.Do(func() { close(s.stopCh) }) +} + +func (s *GeminiCachedContentService) runMutationRecovery() { + ticker := time.NewTicker(geminiCachedContentMutationRecoveryInterval) + defer ticker.Stop() + for { + select { + case <-s.stopCh: + return + case <-ticker.C: + recovered, err := s.RecoverStaleMutations(context.Background()) + if err != nil { + logger.LegacyPrintf("service.gemini_cached_content", "Warning: gemini cached content mutation recovery failed (recovered=%d): %v", recovered, err) + } else if recovered > 0 { + logger.LegacyPrintf("service.gemini_cached_content", "gemini cached content mutation recovery rolled back %d stale mutations", recovered) + } + } + } } // Now 返回服务时钟。 @@ -662,15 +727,135 @@ func (s *GeminiCachedContentService) List(ctx context.Context, apiKeyID, groupID return s.repo.ListForOwner(ctx, apiKeyID, groupID, s.now(), beforeID, limit) } -// UpdateExpireTime 更新到期时间并返回更新前库中的到期时间（并发更新按提交顺序各自拿到自己的前值）。 -func (s *GeminiCachedContentService) UpdateExpireTime(ctx context.Context, record *GeminiCachedContent, expireTime time.Time) (time.Time, error) { - previous, err := s.repo.UpdateExpireTime(ctx, record.ID, expireTime) - if err != nil { - return time.Time{}, err +// GeminiCachedContentMutation 是一次已登记修改意图、尚未提交的有效期修改。持有期间同一缓存的其他修改被拒绝； +// 未提交的修改（上游结果未知、落库失败、进程退出）在租约过期后由恢复流程把上游有效期回滚到 Baseline，不计费。 +type GeminiCachedContentMutation struct { + record *GeminiCachedContent + baseline time.Time + token string +} + +// Baseline 返回修改前的到期时间：提交前它始终是库中的到期时间。 +func (m *GeminiCachedContentMutation) Baseline() time.Time { + return m.baseline +} + +// BeginMutation 登记修改 record 有效期的意图并取得租约。同一缓存有进行中的修改，或库中到期时间已不是 +// record.ExpireTime 时返回 ErrGeminiCachedContentMutationConflict；租约已过期的遗留意图先回滚恢复再登记。 +func (s *GeminiCachedContentService) BeginMutation(ctx context.Context, record *GeminiCachedContent) (*GeminiCachedContentMutation, error) { + for attempt := 0; ; attempt++ { + token, err := NewGeminiCachedContentPublicID() + if err != nil { + return nil, err + } + acquired, err := s.repo.BeginMutation(ctx, record.ID, record.ExpireTime, token, geminiCachedContentMutationLease) + if err != nil { + return nil, err + } + if acquired { + return &GeminiCachedContentMutation{record: record, baseline: record.ExpireTime, token: token}, nil + } + if attempt > 0 { + return nil, ErrGeminiCachedContentMutationConflict + } + recovered, err := s.recoverStaleMutation(ctx, record.ID) + if err != nil { + return nil, errors.Join(ErrGeminiCachedContentMutationConflict, err) + } + if !recovered { + return nil, ErrGeminiCachedContentMutationConflict + } } - record.ExpireTime = expireTime - record.UpdatedAt = s.now() - return previous, nil +} + +// CommitMutation 写入新的到期时间并清除修改意图。修改意图已不属于 m（缓存已删除或租约已被接管）时返回 +// ErrGeminiCachedContentMutationConflict；其他失败时意图保留，由恢复流程回滚上游。 +func (s *GeminiCachedContentService) CommitMutation(ctx context.Context, m *GeminiCachedContentMutation, expireTime time.Time) error { + if err := s.repo.CommitMutation(ctx, m.record.ID, m.token, expireTime); err != nil { + return err + } + m.record.ExpireTime = expireTime + m.record.UpdatedAt = s.now() + return nil +} + +// ReleaseMutation 放弃修改并清除修改意图，用于上游明确未改变有效期的情形。 +func (s *GeminiCachedContentService) ReleaseMutation(ctx context.Context, m *GeminiCachedContentMutation) error { + return s.repo.ReleaseMutation(ctx, m.record.ID, m.token) +} + +// RecoverStaleMutations 回滚恢复一批租约已过期的修改意图，返回完成恢复的数量。 +func (s *GeminiCachedContentService) RecoverStaleMutations(ctx context.Context) (int, error) { + listCtx, cancel := context.WithTimeout(ctx, geminiCachedContentMutationRecoveryTimeout) + ids, err := s.repo.ListStaleMutations(listCtx, geminiCachedContentMutationRecoveryBatch) + cancel() + if err != nil { + return 0, err + } + recovered := 0 + var errs []error + for _, id := range ids { + select { + case <-s.stopCh: + return recovered, errors.Join(errs...) + default: + } + itemCtx, cancel := context.WithTimeout(ctx, geminiCachedContentMutationRecoveryTimeout) + ok, err := s.recoverStaleMutation(itemCtx, id) + cancel() + if err != nil { + errs = append(errs, fmt.Errorf("cached content %d: %w", id, err)) + continue + } + if ok { + recovered++ + } + } + return recovered, errors.Join(errs...) +} + +// recoverStaleMutation 接管 id 上租约已过期的修改意图，把上游有效期回滚到库中的到期时间后清除意图； +// 库中到期时间已到时删除上游缓存与本地记录。没有可接管的意图时返回 false；回滚失败时意图保留，租约过期后重试。 +func (s *GeminiCachedContentService) recoverStaleMutation(ctx context.Context, id int64) (bool, error) { + token, err := NewGeminiCachedContentPublicID() + if err != nil { + return false, err + } + record, err := s.repo.TakeOverStaleMutation(ctx, id, token, geminiCachedContentMutationLease) + if err != nil || record == nil { + return false, err + } + forget, err := s.rollbackUpstream(ctx, record) + if err != nil { + return false, err + } + if forget { + return true, s.repo.SoftDelete(ctx, record.ID) + } + return true, s.repo.ReleaseMutation(ctx, record.ID, token) +} + +// rollbackUpstream 把上游有效期改回 record.ExpireTime，剩余不足 1 秒时删除上游缓存； +// 返回 true 表示上游缓存已不存在，本地记录应移除。持有账号已不可用时不改动上游。 +func (s *GeminiCachedContentService) rollbackUpstream(ctx context.Context, record *GeminiCachedContent) (bool, error) { + account, err := s.BoundAccount(ctx, record) + if err != nil || account == nil { + return false, err + } + remaining := record.ExpireTime.Sub(s.now()) + if remaining < time.Second { + err = s.DeleteUpstream(ctx, account, record) + } else { + _, err = s.PatchUpstream(ctx, account, record, remaining) + } + var upErr *GeminiCachedContentUpstreamError + if errors.As(err, &upErr) && upErr.NotFound() { + return true, nil + } + if err != nil { + return false, err + } + return remaining < time.Second, nil } func (s *GeminiCachedContentService) Forget(ctx context.Context, record *GeminiCachedContent) error { @@ -743,7 +928,7 @@ func (s *GeminiCachedContentService) CreateUpstream(ctx context.Context, account if err != nil { return nil, err } - respBody, err := s.doUpstream(ctx, account, http.MethodPost, collectionURL, body) + respBody, err := s.doUpstream(ctx, account, http.MethodPost, collectionURL, body, geminiCachedContentUpstreamTimeout) if err != nil { return nil, err } @@ -759,17 +944,13 @@ func (s *GeminiCachedContentService) CreateUpstream(ctx context.Context, account // PatchUpstream 把上游缓存的有效期改为自现在起 ttl。 func (s *GeminiCachedContentService) PatchUpstream(ctx context.Context, account *Account, record *GeminiCachedContent, ttl time.Duration) (*GeminiCachedContentUpstream, error) { - base, err := s.apiBase(account) - if err != nil { - return nil, err - } - resourceURL, err := geminiCachedContentResourceURL(account, base, record.UpstreamName) + resourceURL, err := s.upstreamResourceURL(account, record) if err != nil { return nil, err } seconds := int64(math.Ceil(ttl.Seconds())) body, _ := json.Marshal(map[string]any{"ttl": fmt.Sprintf("%ds", seconds)}) - respBody, err := s.doUpstream(ctx, account, http.MethodPatch, resourceURL+"?updateMask=ttl", body) + respBody, err := s.doUpstream(ctx, account, http.MethodPatch, resourceURL+"?updateMask=ttl", body, geminiCachedContentMutationUpstreamTimeout) if err != nil { return nil, err } @@ -784,25 +965,38 @@ func (s *GeminiCachedContentService) PatchUpstream(ctx context.Context, account // DeleteUpstream 删除上游缓存。 func (s *GeminiCachedContentService) DeleteUpstream(ctx context.Context, account *Account, record *GeminiCachedContent) error { - base, err := s.apiBase(account) + resourceURL, err := s.upstreamResourceURL(account, record) if err != nil { return err } - resourceURL, err := geminiCachedContentResourceURL(account, base, record.UpstreamName) - if err != nil { - return err - } - _, err = s.doUpstream(ctx, account, http.MethodDelete, resourceURL, nil) + _, err = s.doUpstream(ctx, account, http.MethodDelete, resourceURL, nil, geminiCachedContentMutationUpstreamTimeout) return err } -// doUpstream 调用上游缓存接口并返回 2xx 响应体。上游调用与客户端连接脱钩：客户端断开不会中途取消 -// 已发出的创建 / 修改，调用方据结果完成落库、计费或清理。 -func (s *GeminiCachedContentService) doUpstream(ctx context.Context, account *Account, method, target string, body []byte) ([]byte, error) { - if s.compat == nil || s.compat.httpUpstream == nil { - return nil, errors.New("gemini upstream client not configured") +// upstreamResourceURL 返回 record 对应上游缓存资源的地址；失败时请求不会发出。 +func (s *GeminiCachedContentService) upstreamResourceURL(account *Account, record *GeminiCachedContent) (string, error) { + base, err := s.apiBase(account) + if err != nil { + return "", geminiCachedContentNotSent(err) } - ctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), geminiCachedContentUpstreamTimeout) + resourceURL, err := geminiCachedContentResourceURL(account, base, record.UpstreamName) + if err != nil { + return "", geminiCachedContentNotSent(err) + } + return resourceURL, nil +} + +func geminiCachedContentNotSent(err error) error { + return fmt.Errorf("%w: %w", ErrGeminiCachedContentRequestNotSent, err) +} + +// doUpstream 调用上游缓存接口并返回 2xx 响应体，单次调用不超过 timeout。上游调用与客户端连接脱钩：客户端断开不会中途取消 +// 已发出的创建 / 修改，调用方据结果完成落库、计费或清理。请求发出前的失败包装为 ErrGeminiCachedContentRequestNotSent。 +func (s *GeminiCachedContentService) doUpstream(ctx context.Context, account *Account, method, target string, body []byte, timeout time.Duration) ([]byte, error) { + if s.compat == nil || s.compat.httpUpstream == nil { + return nil, geminiCachedContentNotSent(errors.New("gemini upstream client not configured")) + } + ctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), timeout) defer cancel() var reader io.Reader if body != nil { @@ -810,7 +1004,7 @@ func (s *GeminiCachedContentService) doUpstream(ctx context.Context, account *Ac } req, err := http.NewRequestWithContext(ctx, method, target, reader) if err != nil { - return nil, err + return nil, geminiCachedContentNotSent(err) } if body != nil { req.Header.Set("Content-Type", "application/json") @@ -819,20 +1013,20 @@ func (s *GeminiCachedContentService) doUpstream(ctx context.Context, account *Ac case AccountTypeAPIKey: apiKey := strings.TrimSpace(account.GetCredential("api_key")) if apiKey == "" { - return nil, errors.New("gemini api_key not configured") + return nil, geminiCachedContentNotSent(errors.New("gemini api_key not configured")) } req.Header.Set("x-goog-api-key", apiKey) case AccountTypeServiceAccount: if s.compat.tokenProvider == nil { - return nil, errors.New("gemini token provider not configured") + return nil, geminiCachedContentNotSent(errors.New("gemini token provider not configured")) } accessToken, err := s.compat.tokenProvider.GetAccessToken(ctx, account) if err != nil { - return nil, err + return nil, geminiCachedContentNotSent(err) } req.Header.Set("Authorization", "Bearer "+accessToken) default: - return nil, fmt.Errorf("account type %s does not support explicit context caching", account.Type) + return nil, geminiCachedContentNotSent(fmt.Errorf("account type %s does not support explicit context caching", account.Type)) } proxyURL := "" if account.ProxyID != nil && account.Proxy != nil { diff --git a/backend/internal/service/gemini_cached_content_hardening_test.go b/backend/internal/service/gemini_cached_content_hardening_test.go index b1b893187a..f93b9424a0 100644 --- a/backend/internal/service/gemini_cached_content_hardening_test.go +++ b/backend/internal/service/gemini_cached_content_hardening_test.go @@ -295,75 +295,186 @@ func TestGeminiForwardNative_BoundCacheUpstreamGatewayErrorsSkipAccountPenalty(t } } -type chargeAffordableCacheStub struct { +type fixedChargeCacheStub struct { billingCacheWorkerStub balance float64 sub *SubscriptionCacheData rate *APIKeyRateLimitCacheData + // reserved 按维度（user:{id} / scope 名）记录在途预留。 + reserved map[string]map[string]float64 } -func (s *chargeAffordableCacheStub) GetUserBalance(context.Context, int64) (float64, error) { +func (s *fixedChargeCacheStub) GetUserBalance(context.Context, int64) (float64, error) { return s.balance, nil } -func (s *chargeAffordableCacheStub) GetSubscriptionCache(context.Context, int64, int64) (*SubscriptionCacheData, error) { +func (s *fixedChargeCacheStub) GetSubscriptionCache(context.Context, int64, int64) (*SubscriptionCacheData, error) { if s.sub == nil { return nil, errors.New("cache miss") } return s.sub, nil } -func (s *chargeAffordableCacheStub) GetAPIKeyRateLimit(context.Context, int64) (*APIKeyRateLimitCacheData, error) { +func (s *fixedChargeCacheStub) GetAPIKeyRateLimit(context.Context, int64) (*APIKeyRateLimitCacheData, error) { if s.rate == nil { return nil, errors.New("cache miss") } return s.rate, nil } -func TestCheckChargeAffordable(t *testing.T) { +func (s *fixedChargeCacheStub) reserve(key, requestID string, amount float64) float64 { + if s.reserved == nil { + s.reserved = map[string]map[string]float64{} + } + if s.reserved[key] == nil { + s.reserved[key] = map[string]float64{} + } + sum := 0.0 + for _, v := range s.reserved[key] { + sum += v + } + s.reserved[key][requestID] = amount + return sum +} + +func (s *fixedChargeCacheStub) ReserveInflightBalance(_ context.Context, userID int64, requestID string, amount, balance float64, _ time.Duration) (bool, float64, error) { + key := fmt.Sprintf("user:%d", userID) + others := s.reserve(key, requestID, amount) + if len(s.reserved[key]) > 1 && balance-others < amount { + delete(s.reserved[key], requestID) + return false, others, nil + } + return true, others, nil +} + +func (s *fixedChargeCacheStub) ReleaseInflightBalance(_ context.Context, userID int64, requestID string) error { + delete(s.reserved[fmt.Sprintf("user:%d", userID)], requestID) + return nil +} + +func (s *fixedChargeCacheStub) ReserveInflightScope(_ context.Context, scope, requestID string, amount float64, _ time.Duration) (float64, error) { + return s.reserve(scope, requestID, amount), nil +} + +func (s *fixedChargeCacheStub) ReleaseInflightScope(_ context.Context, scope, requestID string) error { + delete(s.reserved[scope], requestID) + return nil +} + +type fixedChargeAPIKeyRepoStub struct { + APIKeyRepository + data *APIKeyRateLimitData +} + +func (r *fixedChargeAPIKeyRepoStub) GetRateLimitData(context.Context, int64) (*APIKeyRateLimitData, error) { + return r.data, nil +} + +type fixedChargeSubscriptionRepoStub struct { + UserSubscriptionRepository + sub *UserSubscription +} + +func (r *fixedChargeSubscriptionRepoStub) GetActiveByUserIDAndGroupID(context.Context, int64, int64) (*UserSubscription, error) { + return r.sub, nil +} + +func TestReserveFixedCharge(t *testing.T) { ctx := context.Background() user := &User{ID: 1} - newSvc := func(cache *chargeAffordableCacheStub) *BillingCacheService { + newSvc := func(cache *fixedChargeCacheStub) *BillingCacheService { svc := NewBillingCacheService(cache, nil, nil, nil, nil, nil, &config.Config{}, nil) t.Cleanup(svc.Stop) return svc } + reservingSvc := func(cache *fixedChargeCacheStub, subRepo UserSubscriptionRepository, apiKeyRepo APIKeyRepository) *BillingCacheService { + cfg := &config.Config{} + cfg.Billing.InflightReservation.Enabled = true + svc := NewBillingCacheService(cache, nil, subRepo, apiKeyRepo, nil, nil, cfg, nil) + t.Cleanup(svc.Stop) + return svc + } + reserve := func(svc *BillingCacheService, apiKey *APIKey, group *Group, sub *UserSubscription, amount float64) error { + res, err := svc.ReserveFixedCharge(ctx, user, apiKey, group, sub, amount) + res.Release() + return err + } t.Run("balance must cover the amount", func(t *testing.T) { - svc := newSvc(&chargeAffordableCacheStub{balance: 10}) - require.NoError(t, svc.CheckChargeAffordable(ctx, user, nil, nil, nil, 10)) - require.ErrorIs(t, svc.CheckChargeAffordable(ctx, user, nil, nil, nil, 10.01), ErrInsufficientBalance) - require.NoError(t, svc.CheckChargeAffordable(ctx, user, nil, nil, nil, 0)) + svc := newSvc(&fixedChargeCacheStub{balance: 10}) + require.NoError(t, reserve(svc, nil, nil, nil, 10)) + require.ErrorIs(t, reserve(svc, nil, nil, nil, 10.01), ErrInsufficientBalance) + require.NoError(t, reserve(svc, nil, nil, nil, 0)) }) t.Run("subscription windows must cover the amount", func(t *testing.T) { daily := 5.0 group := &Group{ID: 2, SubscriptionType: SubscriptionTypeSubscription, DailyLimitUSD: &daily} - svc := newSvc(&chargeAffordableCacheStub{sub: &SubscriptionCacheData{ + svc := newSvc(&fixedChargeCacheStub{sub: &SubscriptionCacheData{ Status: SubscriptionStatusActive, ExpiresAt: time.Now().Add(time.Hour), DailyUsage: 3, }}) sub := &UserSubscription{ID: 1} - require.NoError(t, svc.CheckChargeAffordable(ctx, user, nil, group, sub, 2)) - require.ErrorIs(t, svc.CheckChargeAffordable(ctx, user, nil, group, sub, 2.5), ErrDailyLimitExceeded) + require.NoError(t, reserve(svc, nil, group, sub, 2)) + require.ErrorIs(t, reserve(svc, nil, group, sub, 2.5), ErrDailyLimitExceeded) }) t.Run("api key quota and rate windows include the amount", func(t *testing.T) { - svc := newSvc(&chargeAffordableCacheStub{balance: 100, rate: &APIKeyRateLimitCacheData{ + svc := newSvc(&fixedChargeCacheStub{balance: 100, rate: &APIKeyRateLimitCacheData{ Usage1d: 8, Window1d: time.Now().Unix(), }}) quotaKey := &APIKey{ID: 5, Quota: 10, QuotaUsed: 9} - require.NoError(t, svc.CheckChargeAffordable(ctx, user, quotaKey, nil, nil, 1)) - require.ErrorIs(t, svc.CheckChargeAffordable(ctx, user, quotaKey, nil, nil, 1.5), ErrAPIKeyQuotaExhausted) + require.NoError(t, reserve(svc, quotaKey, nil, nil, 1)) + require.ErrorIs(t, reserve(svc, quotaKey, nil, nil, 1.5), ErrAPIKeyQuotaExhausted) windowKey := &APIKey{ID: 6, RateLimit1d: 10} - require.NoError(t, svc.CheckChargeAffordable(ctx, user, windowKey, nil, nil, 2)) - require.ErrorIs(t, svc.CheckChargeAffordable(ctx, user, windowKey, nil, nil, 2.5), ErrAPIKeyRateLimit1dExceeded) + require.NoError(t, reserve(svc, windowKey, nil, nil, 2)) + require.ErrorIs(t, reserve(svc, windowKey, nil, nil, 2.5), ErrAPIKeyRateLimit1dExceeded) + }) + + t.Run("usage is read from the database instead of request snapshots", func(t *testing.T) { + now := time.Now() + apiKeyRepo := &fixedChargeAPIKeyRepoStub{data: &APIKeyRateLimitData{QuotaUsed: 9.5, Usage1d: 8, Window1dStart: &now}} + svc := reservingSvc(&fixedChargeCacheStub{balance: 100}, nil, apiKeyRepo) + require.ErrorIs(t, reserve(svc, &APIKey{ID: 5, Quota: 10}, nil, nil, 1), ErrAPIKeyQuotaExhausted, "请求上下文中的 QuotaUsed 可能滞后") + require.ErrorIs(t, reserve(svc, &APIKey{ID: 6, RateLimit1d: 10}, nil, nil, 2.5), ErrAPIKeyRateLimit1dExceeded) + + daily := 5.0 + group := &Group{ID: 2, SubscriptionType: SubscriptionTypeSubscription, DailyLimitUSD: &daily} + subRepo := &fixedChargeSubscriptionRepoStub{sub: &UserSubscription{Status: SubscriptionStatusActive, DailyUsageUSD: 4}} + svc = reservingSvc(&fixedChargeCacheStub{sub: &SubscriptionCacheData{Status: SubscriptionStatusActive}}, subRepo, nil) + require.ErrorIs(t, reserve(svc, nil, group, &UserSubscription{ID: 1}, 2), ErrDailyLimitExceeded, "订阅缓存用量可能滞后") + }) + + t.Run("unsettled charges hold subscription and api key headroom", func(t *testing.T) { + daily := 110.0 + group := &Group{ID: 2, SubscriptionType: SubscriptionTypeSubscription, DailyLimitUSD: &daily} + sub := &UserSubscription{ID: 1} + cache := &fixedChargeCacheStub{balance: 1000, sub: &SubscriptionCacheData{Status: SubscriptionStatusActive}} + svc := reservingSvc(cache, nil, nil) + first, err := svc.ReserveFixedCharge(ctx, user, nil, group, sub, 103.5) + require.NoError(t, err) + _, err = svc.ReserveFixedCharge(ctx, user, nil, group, sub, 103.5) + require.ErrorIs(t, err, ErrDailyLimitExceeded) + first.Release() + second, err := svc.ReserveFixedCharge(ctx, user, nil, group, sub, 103.5) + require.NoError(t, err) + second.Release() + + apiKey := &APIKey{ID: 5, Quota: 110} + first, err = svc.ReserveFixedCharge(ctx, user, apiKey, nil, nil, 103.5) + require.NoError(t, err) + _, err = svc.ReserveFixedCharge(ctx, user, apiKey, nil, nil, 103.5) + require.ErrorIs(t, err, ErrAPIKeyQuotaExhausted) + first.Release() + for key, reservations := range cache.reserved { + require.Empty(t, reservations, "拒绝与释放都不留下预留：%s", key) + } }) t.Run("simple mode skips the check", func(t *testing.T) { cfg := &config.Config{RunMode: config.RunModeSimple} - svc := NewBillingCacheService(&chargeAffordableCacheStub{}, nil, nil, nil, nil, nil, cfg, nil) + svc := NewBillingCacheService(&fixedChargeCacheStub{}, nil, nil, nil, nil, nil, cfg, nil) t.Cleanup(svc.Stop) - require.NoError(t, svc.CheckChargeAffordable(ctx, user, nil, nil, nil, 1000)) + require.NoError(t, reserve(svc, nil, nil, nil, 1000)) }) } diff --git a/backend/internal/service/gemini_cached_content_test.go b/backend/internal/service/gemini_cached_content_test.go index 94e1027e23..e3afbdd5e6 100644 --- a/backend/internal/service/gemini_cached_content_test.go +++ b/backend/internal/service/gemini_cached_content_test.go @@ -53,15 +53,30 @@ func (r *geminiCachedContentRepoStub) ListForOwner(_ context.Context, apiKeyID, return out, nil } -func (r *geminiCachedContentRepoStub) UpdateExpireTime(_ context.Context, id int64, expireTime time.Time) (time.Time, error) { +func (r *geminiCachedContentRepoStub) BeginMutation(context.Context, int64, time.Time, string, time.Duration) (bool, error) { + return true, nil +} + +func (r *geminiCachedContentRepoStub) TakeOverStaleMutation(context.Context, int64, string, time.Duration) (*GeminiCachedContent, error) { + return nil, nil +} + +func (r *geminiCachedContentRepoStub) CommitMutation(_ context.Context, id int64, _ string, expireTime time.Time) error { for _, record := range r.records { if record.ID == id { - previous := record.ExpireTime record.ExpireTime = expireTime - return previous, nil + return nil } } - return time.Time{}, ErrGeminiCachedContentNotFound + return ErrGeminiCachedContentMutationConflict +} + +func (r *geminiCachedContentRepoStub) ReleaseMutation(context.Context, int64, string) error { + return nil +} + +func (r *geminiCachedContentRepoStub) ListStaleMutations(context.Context, int) ([]int64, error) { + return nil, nil } func (r *geminiCachedContentRepoStub) SoftDelete(_ context.Context, id int64) error { diff --git a/backend/internal/service/wire.go b/backend/internal/service/wire.go index 56fe96f52d..688d729b31 100644 --- a/backend/internal/service/wire.go +++ b/backend/internal/service/wire.go @@ -640,6 +640,13 @@ func ProvideIdempotencyCleanupService(repo IdempotencyRepository, cfg *config.Co return svc } +// ProvideGeminiCachedContentService creates GeminiCachedContentService and starts its mutation recovery loop. +func ProvideGeminiCachedContentService(repo GeminiCachedContentRepository, compat *GeminiMessagesCompatService, cfg *config.Config) *GeminiCachedContentService { + svc := NewGeminiCachedContentService(repo, compat, cfg) + svc.Start() + return svc +} + // ProvideScheduledTestService creates ScheduledTestService. func ProvideScheduledTestService( planRepo ScheduledTestPlanRepository, @@ -885,7 +892,7 @@ var ProviderSet = wire.NewSet( ProvideGrokOAuthService, wire.Bind(new(GrokOAuthTokenService), new(*GrokOAuthService)), NewGeminiOAuthService, - NewGeminiCachedContentService, + ProvideGeminiCachedContentService, NewGeminiQuotaService, NewCompositeTokenCacheInvalidator, wire.Bind(new(TokenCacheInvalidator), new(*CompositeTokenCacheInvalidator)), diff --git a/backend/migrations/245_gemini_cached_contents_mutation_lease.sql b/backend/migrations/245_gemini_cached_contents_mutation_lease.sql new file mode 100644 index 0000000000..e0970efebf --- /dev/null +++ b/backend/migrations/245_gemini_cached_contents_mutation_lease.sql @@ -0,0 +1,10 @@ +-- Gemini 显式缓存修改有效期的进行中操作（修改意图）：mutation_token 标识持有者， +-- mutation_lease_until 之前同一缓存的其他修改被拒绝；expire_time 在提交前保持修改前的值， +-- 租约过期仍未清除的意图由恢复流程把上游有效期回滚到 expire_time。 +ALTER TABLE gemini_cached_contents + ADD COLUMN IF NOT EXISTS mutation_token VARCHAR(64) NULL, + ADD COLUMN IF NOT EXISTS mutation_lease_until TIMESTAMPTZ NULL; + +CREATE INDEX IF NOT EXISTS idx_gemini_cached_contents_pending_mutation + ON gemini_cached_contents (mutation_lease_until) + WHERE mutation_token IS NOT NULL AND deleted_at IS NULL;