fix(gemini): 显式缓存修改串行化、失败回滚,确定金额费用按 Key / 订阅预留

同一缓存的并发 PATCH 可能以与上游提交不同的顺序写回本地到期时间;
确定金额的存储费只预留用户余额,并发续期可越过 API Key 总额度与订阅额度;
上游已修改而本地写回失败时没有可恢复的状态。

- 新增迁移 245:缓存记录增加修改意图(mutation_token / mutation_lease_until)。
  PATCH 在调用上游前登记意图并取得 2 分钟租约,条件为没有进行中的修改且到期时间
  未变;上游修改、本地写回与计费均由持有者完成,并发的修改返回 409 ABORTED
- 上游结果未知、本地写回失败或进程退出时修改意图保留且不计费;租约过期后由下一次
  修改或后台每分钟的恢复流程把上游有效期回滚到修改前的到期时间(已到期则删除上游
  缓存与本地记录)。上游明确拒绝或请求未发出时直接清除意图
- 修改 / 删除上游缓存的单次调用超时改为 30 秒;请求发出前的失败包装为
  ErrGeminiCachedContentRequestNotSent
- CheckChargeAffordable 与 ReserveInflightCharge 合并为 ReserveFixedCharge:在途预留
  新增订阅、API Key 两个维度(Redis billing:inflight_scope:*),先登记本笔预留再读取
  库中权威用量(Key quota_used 与限速窗口、订阅日 / 周 / 月用量)叠加同维度在途合计
  校验;预留随计费任务交接,落库后释放
- GetRateLimitData 一并读取 quota_used
This commit is contained in:
feeeei
2026-10-12 17:14:56 +08:00
parent daebb7a80c
commit 26b31a83e4
20 changed files with 1259 additions and 172 deletions
+7
View File
@@ -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()
+9 -2
View File
@@ -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()
+1
View File
@@ -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,
@@ -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) {
@@ -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 {
@@ -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 {
+2 -2
View File
@@ -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()
@@ -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() {
@@ -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)
)
@@ -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))
}
@@ -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 {
@@ -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)
@@ -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
@@ -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
}
@@ -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) {
+229 -35
View File
@@ -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 {
@@ -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))
})
}
@@ -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 {
+8 -1
View File
@@ -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)),
@@ -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;