From ccd535ef8e50cf6e5846a59278c40b7ff59d1b7d Mon Sep 17 00:00:00 2001 From: CaIon Date: Mon, 10 Aug 2026 22:21:29 +0800 Subject: [PATCH] fix: harden concurrent quota and status updates --- controller/token.go | 4 +- model/channel.go | 33 ++-- model/channel_status_test.go | 102 ++++++++++++ model/quota_reserve.go | 182 +++++++++++++++++++- model/quota_reserve_test.go | 228 ++++++++++++++++++++++++++ model/subscription.go | 16 +- model/token.go | 97 ++++------- model/token_auto_groups_cache_test.go | 19 ++- model/token_cache.go | 116 ++++++++----- model/user.go | 37 +---- model/user_cache.go | 30 ---- service/billing_session.go | 11 ++ service/funding_source.go | 12 +- service/quota.go | 18 +- 14 files changed, 702 insertions(+), 203 deletions(-) create mode 100644 model/channel_status_test.go create mode 100644 model/quota_reserve_test.go diff --git a/controller/token.go b/controller/token.go index c26d82e3..ff2aca8f 100644 --- a/controller/token.go +++ b/controller/token.go @@ -279,7 +279,7 @@ func AddToken(c *gin.Context) { common.ApiErrorI18n(c, i18n.MsgTokenQuotaNegative) return } - maxQuotaValue := int((1000000000 * common.QuotaPerUnit)) + maxQuotaValue := common.QuotaFromFloat(1000000000 * common.QuotaPerUnit) if token.RemainQuota > maxQuotaValue { common.ApiErrorI18n(c, i18n.MsgTokenQuotaExceedMax, map[string]any{"Max": maxQuotaValue}) return @@ -373,7 +373,7 @@ func UpdateToken(c *gin.Context) { common.ApiErrorI18n(c, i18n.MsgTokenQuotaNegative) return } - maxQuotaValue := int((1000000000 * common.QuotaPerUnit)) + maxQuotaValue := common.QuotaFromFloat(1000000000 * common.QuotaPerUnit) if token.RemainQuota > maxQuotaValue { common.ApiErrorI18n(c, i18n.MsgTokenQuotaExceedMax, map[string]any{"Max": maxQuotaValue}) return diff --git a/model/channel.go b/model/channel.go index 2cd7c311..0f8cdb10 100644 --- a/model/channel.go +++ b/model/channel.go @@ -346,11 +346,21 @@ func (channel *Channel) Save() error { return DB.Save(channel).Error } -func (channel *Channel) SaveWithoutKey() error { +// saveStatusState persists only the fields owned by the channel status flow. +// Keeping this allowlist here prevents a stale channel snapshot from +// overwriting credentials, accounting counters, or channel configuration. +func (channel *Channel) saveStatusState() error { if channel.Id == 0 { return errors.New("channel ID is 0") } - return DB.Omit("key").Save(channel).Error + updates := map[string]any{ + "status": channel.Status, + "other_info": channel.OtherInfo, + } + if channel.ChannelInfo.IsMultiKey { + updates["channel_info"] = channel.ChannelInfo + } + return DB.Model(&Channel{}).Where("id = ?", channel.Id).Updates(updates).Error } func GetAllChannels(startIdx int, num int, selectAll bool, idSort bool, sortOptions ...ChannelSortOptions) ([]*Channel, error) { @@ -713,19 +723,24 @@ func UpdateChannelStatus(channelId int, usingKey string, status int, reason stri if common.MemoryCacheEnabled { channelStatusLock.Lock() defer channelStatusLock.Unlock() + } + // ChannelInfo stores both multi-key status and the polling cursor. Hold the + // same per-channel lock from the first read through persistence so neither + // writer can save a stale JSON snapshot over the other. + pollingLock := GetChannelPollingLock(channelId) + pollingLock.Lock() + defer pollingLock.Unlock() + + if common.MemoryCacheEnabled { channelCache, _ := CacheGetChannel(channelId) if channelCache == nil { return false } if channelCache.ChannelInfo.IsMultiKey { - // Use per-channel lock to prevent concurrent map read/write with GetNextEnabledKey beforeStatus := channelCache.Status - pollingLock := GetChannelPollingLock(channelId) - pollingLock.Lock() // 如果是多Key模式，更新缓存中的状态 handlerMultiKeyUpdate(channelCache, usingKey, status, reason) - pollingLock.Unlock() if beforeStatus != channelCache.Status { CacheUpdateChannelStatus(channelId, channelCache.Status) } @@ -759,11 +774,7 @@ func UpdateChannelStatus(channelId int, usingKey string, status int, reason stri if channel.ChannelInfo.IsMultiKey { beforeStatus := channel.Status - // Protect map writes with the same per-channel lock used by readers - pollingLock := GetChannelPollingLock(channelId) - pollingLock.Lock() handlerMultiKeyUpdate(channel, usingKey, status, reason) - pollingLock.Unlock() if beforeStatus != channel.Status { shouldUpdateAbilities = true } @@ -775,7 +786,7 @@ func UpdateChannelStatus(channelId int, usingKey string, status int, reason stri channel.Status = status shouldUpdateAbilities = true } - err = channel.SaveWithoutKey() + err = channel.saveStatusState() if err != nil { common.SysLog(fmt.Sprintf("failed to update channel status: channel_id=%d, status=%d, error=%v", channel.Id, status, err)) return false diff --git a/model/channel_status_test.go b/model/channel_status_test.go new file mode 100644 index 00000000..e4ad86f8 --- /dev/null +++ b/model/channel_status_test.go @@ -0,0 +1,102 @@ +package model + +import ( + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func setupChannelStatusTest(t *testing.T) { + t.Helper() + truncateTables(t) + require.NoError(t, DB.Exec("DELETE FROM abilities").Error) + require.NoError(t, DB.Exec("DELETE FROM channels").Error) + + memoryCacheEnabled := common.MemoryCacheEnabled + common.MemoryCacheEnabled = false + t.Cleanup(func() { + common.MemoryCacheEnabled = memoryCacheEnabled + }) +} + +func TestUpdateChannelStatusPersistsMultiKeyState(t *testing.T) { + setupChannelStatusTest(t) + + channel := Channel{ + Name: "multi-key-status", + Key: "key-a\nkey-b", + Status: common.ChannelStatusEnabled, + ChannelInfo: ChannelInfo{ + IsMultiKey: true, + MultiKeySize: 2, + MultiKeyMode: constant.MultiKeyModePolling, + MultiKeyPollingIndex: 1, + }, + } + require.NoError(t, DB.Create(&channel).Error) + + changed := UpdateChannelStatus(channel.Id, "key-a", common.ChannelStatusAutoDisabled, "provider rejected key") + require.True(t, changed) + + var stored Channel + require.NoError(t, DB.First(&stored, channel.Id).Error) + assert.Equal(t, common.ChannelStatusEnabled, stored.Status) + assert.Equal(t, common.ChannelStatusAutoDisabled, stored.ChannelInfo.MultiKeyStatusList[0]) + assert.Equal(t, "provider rejected key", stored.ChannelInfo.MultiKeyDisabledReason[0]) + assert.NotZero(t, stored.ChannelInfo.MultiKeyDisabledTime[0]) + assert.Equal(t, 1, stored.ChannelInfo.MultiKeyPollingIndex) +} + +func TestSaveStatusStateFromSingleKeySnapshotPreservesUnownedColumns(t *testing.T) { + setupChannelStatusTest(t) + + channel := Channel{ + Name: "single-key-status", + Key: "original-key", + Status: common.ChannelStatusEnabled, + Models: "original-model", + Group: "default", + UsedQuota: 100, + ChannelInfo: ChannelInfo{}, + } + require.NoError(t, DB.Create(&channel).Error) + + stale, err := GetChannelById(channel.Id, true) + require.NoError(t, err) + + concurrentChannelInfo := ChannelInfo{ + IsMultiKey: true, + MultiKeySize: 2, + MultiKeyMode: constant.MultiKeyModePolling, + MultiKeyPollingIndex: 1, + } + require.NoError(t, DB.Model(&Channel{}).Where("id = ?", channel.Id).Updates(map[string]any{ + "key": "rotated-key", + "used_quota": gorm.Expr("used_quota + ?", 250), + "models": "concurrent-model", + "channel_info": concurrentChannelInfo, + }).Error) + + stale.Status = common.ChannelStatusManuallyDisabled + stale.SetOtherInfo(map[string]interface{}{ + "status_reason": "manual operation", + "status_time": int64(1234), + }) + require.NoError(t, stale.saveStatusState()) + + var stored Channel + require.NoError(t, DB.First(&stored, channel.Id).Error) + assert.Equal(t, common.ChannelStatusManuallyDisabled, stored.Status) + assert.Equal(t, "rotated-key", stored.Key) + assert.Equal(t, int64(350), stored.UsedQuota) + assert.Equal(t, "concurrent-model", stored.Models) + assert.Equal(t, concurrentChannelInfo, stored.ChannelInfo) + + otherInfo := stored.GetOtherInfo() + assert.Equal(t, "manual operation", otherInfo["status_reason"]) + assert.Equal(t, float64(1234), otherInfo["status_time"]) +} diff --git a/model/quota_reserve.go b/model/quota_reserve.go index 352377d1..0d0e0525 100644 --- a/model/quota_reserve.go +++ b/model/quota_reserve.go @@ -2,8 +2,11 @@ package model import ( "context" + "errors" + "fmt" "github.com/QuantumNous/new-api/common" + "gorm.io/gorm" ) type cacheQuotaResult int @@ -14,11 +17,6 @@ const ( cacheQuotaMiss ) -// 下列脚本都是守卫式的：只在完整哈希（Id 匹配且配额字段存在）上操作， -// 哈希缺失时返回 miss 而不是创建残缺哈希。脚本不修改 TTL（HINCRBY 天然保留 -// 水合时设置的 TTL），因此即使某个写库路径绕过了缓存，偏差也会在一个 TTL -// 窗口内随缓存过期而自愈。 - const userQuotaReserveScript = ` if tonumber(redis.call('HGET', KEYS[1], 'Id') or '0') ~= tonumber(ARGV[2]) or tonumber(redis.call('HGET', KEYS[1], 'CacheSchema') or '0') ~= tonumber(ARGV[3]) @@ -41,6 +39,32 @@ end redis.call('HINCRBY', KEYS[1], 'Quota', tonumber(ARGV[1])) return 1` +const tokenQuotaReserveScript = ` +if tonumber(redis.call('HGET', KEYS[1], 'Id') or '0') ~= tonumber(ARGV[2]) + or redis.call('HEXISTS', KEYS[1], 'RemainQuota') == 0 + or redis.call('HEXISTS', KEYS[1], 'UsedQuota') == 0 then + return -1 +end +local remain = tonumber(redis.call('HGET', KEYS[1], 'RemainQuota')) +if remain == nil or remain < tonumber(ARGV[1]) then + return 0 +end +redis.call('HINCRBY', KEYS[1], 'RemainQuota', -tonumber(ARGV[1])) +redis.call('HINCRBY', KEYS[1], 'UsedQuota', tonumber(ARGV[1])) +redis.call('HSET', KEYS[1], 'AccessedTime', ARGV[3]) +return 1` + +const tokenQuotaDeltaScript = ` +if tonumber(redis.call('HGET', KEYS[1], 'Id') or '0') ~= tonumber(ARGV[2]) + or redis.call('HEXISTS', KEYS[1], 'RemainQuota') == 0 + or redis.call('HEXISTS', KEYS[1], 'UsedQuota') == 0 then + return -1 +end +redis.call('HINCRBY', KEYS[1], 'RemainQuota', tonumber(ARGV[1])) +redis.call('HINCRBY', KEYS[1], 'UsedQuota', -tonumber(ARGV[1])) +redis.call('HSET', KEYS[1], 'AccessedTime', ARGV[3]) +return 1` + func quotaResultFromLua(result int, err error) (cacheQuotaResult, error) { if err != nil { return cacheQuotaMiss, err @@ -66,3 +90,151 @@ func cacheApplyUserQuotaDelta(userID int, delta int64) (cacheQuotaResult, error) []string{getUserCacheKey(userID)}, delta, userID, userCacheSchemaVersion).Int() return quotaResultFromLua(result, err) } + +func cacheTryReserveTokenQuota(id int, key string, amount int64) (cacheQuotaResult, error) { + result, err := common.RDB.Eval(context.Background(), tokenQuotaReserveScript, + []string{getTokenCacheKey(key)}, amount, id, common.GetTimestamp()).Int() + return quotaResultFromLua(result, err) +} + +func cacheApplyTokenQuotaDelta(id int, key string, delta int64) (cacheQuotaResult, error) { + result, err := common.RDB.Eval(context.Background(), tokenQuotaDeltaScript, + []string{getTokenCacheKey(key)}, delta, id, common.GetTimestamp()).Int() + return quotaResultFromLua(result, err) +} + +// persistUserQuotaDelta 把已在缓存侧预扣成功的增量落库；批量模式下入队， +// 直写模式下要求行存在（用户已删除时报错，交由调用方补偿缓存）。 +func persistUserQuotaDelta(id int, delta int) error { + if common.BatchUpdateEnabled { + addNewRecord(BatchUpdateTypeUserQuota, id, delta) + return nil + } + result := DB.Model(&User{}).Where("id = ?", id).Update("quota", gorm.Expr("quota + ?", delta)) + if result.Error != nil { + return result.Error + } + if result.RowsAffected != 1 { + return gorm.ErrRecordNotFound + } + return nil +} + +func persistTokenQuotaDelta(id int, delta int) error { + if common.BatchUpdateEnabled { + addNewRecord(BatchUpdateTypeTokenQuota, id, delta) + return nil + } + result := DB.Model(&Token{}).Where("id = ?", id).Updates( + map[string]interface{}{ + "remain_quota": gorm.Expr("remain_quota + ?", delta), + "used_quota": gorm.Expr("used_quota - ?", delta), + "accessed_time": common.GetTimestamp(), + }, + ) + if result.Error != nil { + return result.Error + } + if result.RowsAffected != 1 { + return gorm.ErrRecordNotFound + } + return nil +} + +func reserveUserQuotaDB(id int, quota int) (bool, error) { + result := DB.Model(&User{}). + Where("id = ? AND quota >= ?", id, quota). + Update("quota", gorm.Expr("quota - ?", quota)) + return result.RowsAffected == 1, result.Error +} + +func reserveTokenQuotaDB(id int, quota int) (bool, error) { + result := DB.Model(&Token{}). + Where("id = ? AND remain_quota >= ?", id, quota). + Updates(map[string]interface{}{ + "remain_quota": gorm.Expr("remain_quota - ?", quota), + "used_quota": gorm.Expr("used_quota + ?", quota), + "accessed_time": common.GetTimestamp(), + }) + return result.RowsAffected == 1, result.Error +} + +// TryReserveUserQuota atomically checks and deducts a user's wallet quota. +// 缓存命中时以缓存余额为准（避免批量模式下过期的数据库余额放大并发超扣）； +// Redis 异常或水合失败时降级为数据库条件更新，保证服务可用。 +func TryReserveUserQuota(id int, quota int) (bool, error) { + if quota < 0 { + return false, errors.New("quota 不能为负数！") + } + if quota == 0 { + return true, nil + } + if !common.RedisEnabled { + return reserveUserQuotaDB(id, quota) + } + + result, err := cacheTryReserveUserQuota(id, int64(quota)) + if err == nil && result == cacheQuotaMiss { + if _, hydrateErr := GetUserCache(id); hydrateErr == nil { + result, err = cacheTryReserveUserQuota(id, int64(quota)) + } + } + if err != nil || result == cacheQuotaMiss { + if err != nil { + common.SysLog("user quota cache reserve unavailable, falling back to database: " + err.Error()) + } + return reserveUserQuotaDB(id, quota) + } + if result == cacheQuotaInsufficient { + return false, nil + } + if err = persistUserQuotaDelta(id, -quota); err != nil { + compensated, compensateErr := cacheApplyUserQuotaDelta(id, int64(quota)) + if compensateErr != nil || compensated != cacheQuotaOK { + common.SysError(fmt.Sprintf("failed to compensate reserved user quota: result=%d error=%v", compensated, compensateErr)) + } + return false, err + } + return true, nil +} + +// TryReserveTokenQuota atomically checks and deducts a token quota. Unlimited +// tokens skip the balance check but still update remain/used accounting. +func TryReserveTokenQuota(id int, key string, quota int, unlimited bool) (bool, error) { + if quota < 0 { + return false, errors.New("quota 不能为负数！") + } + if quota == 0 { + return true, nil + } + if unlimited { + return true, DecreaseTokenQuota(id, key, quota) + } + if !common.RedisEnabled { + return reserveTokenQuotaDB(id, quota) + } + + result, err := cacheTryReserveTokenQuota(id, key, int64(quota)) + if err == nil && result == cacheQuotaMiss { + if _, hydrateErr := GetTokenByKey(key, true); hydrateErr == nil { + result, err = cacheTryReserveTokenQuota(id, key, int64(quota)) + } + } + if err != nil || result == cacheQuotaMiss { + if err != nil { + common.SysLog("token quota cache reserve unavailable, falling back to database: " + err.Error()) + } + return reserveTokenQuotaDB(id, quota) + } + if result == cacheQuotaInsufficient { + return false, nil + } + if err = persistTokenQuotaDelta(id, -quota); err != nil { + compensated, compensateErr := cacheApplyTokenQuotaDelta(id, key, int64(quota)) + if compensateErr != nil || compensated != cacheQuotaOK { + common.SysError(fmt.Sprintf("failed to compensate reserved token quota: result=%d error=%v", compensated, compensateErr)) + } + return false, err + } + return true, nil +} diff --git a/model/quota_reserve_test.go b/model/quota_reserve_test.go new file mode 100644 index 00000000..76eab83a --- /dev/null +++ b/model/quota_reserve_test.go @@ -0,0 +1,228 @@ +package model + +import ( + "testing" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func createReserveTestUser(t *testing.T, quota int) User { + t.Helper() + user := User{ + Username: "reserve-user-" + common.GetRandomString(6), + Password: "unused-password-hash", + Role: common.RoleCommonUser, + Status: common.UserStatusEnabled, + Group: "default", + AuthVersion: 1, + Quota: quota, + AffCode: "reserve-aff-" + common.GetRandomString(8), + } + require.NoError(t, DB.Create(&user).Error) + return user +} + +func createReserveTestToken(t *testing.T, remainQuota int) Token { + t.Helper() + token := Token{ + UserId: 1, + Key: "reserve-token-" + common.GetRandomString(8), + Name: "reserve-test", + Status: common.TokenStatusEnabled, + ExpiredTime: -1, + RemainQuota: remainQuota, + } + require.NoError(t, token.Insert()) + return token +} + +func getUserQuotaFromDB(t *testing.T, id int) int { + t.Helper() + var user User + require.NoError(t, DB.Select("quota").First(&user, id).Error) + return user.Quota +} + +func getTokenFromDB(t *testing.T, id int) Token { + t.Helper() + var token Token + require.NoError(t, DB.First(&token, id).Error) + return token +} + +func resetBatchUpdateTestState(t *testing.T) { + t.Helper() + oldBatchEnabled := common.BatchUpdateEnabled + common.BatchUpdateEnabled = false + for i := 0; i < BatchUpdateTypeCount; i++ { + batchUpdateLocks[i].Lock() + batchUpdateStores[i] = make(map[int]int) + batchUpdateLocks[i].Unlock() + } + t.Cleanup(func() { + common.BatchUpdateEnabled = oldBatchEnabled + for i := 0; i < BatchUpdateTypeCount; i++ { + batchUpdateLocks[i].Lock() + batchUpdateStores[i] = make(map[int]int) + batchUpdateLocks[i].Unlock() + } + }) +} + +func TestTryReserveQuotaWithoutRedis(t *testing.T) { + truncateTables(t) + resetBatchUpdateTestState(t) + + user := createReserveTestUser(t, 100) + reserved, err := TryReserveUserQuota(user.Id, 60) + require.NoError(t, err) + assert.True(t, reserved) + assert.Equal(t, 40, getUserQuotaFromDB(t, user.Id)) + + reserved, err = TryReserveUserQuota(user.Id, 41) + require.NoError(t, err) + assert.False(t, reserved) + assert.Equal(t, 40, getUserQuotaFromDB(t, user.Id)) + + token := createReserveTestToken(t, 80) + reserved, err = TryReserveTokenQuota(token.Id, token.Key, 25, false) + require.NoError(t, err) + assert.True(t, reserved) + reloaded := getTokenFromDB(t, token.Id) + assert.Equal(t, 55, reloaded.RemainQuota) + assert.Equal(t, 25, reloaded.UsedQuota) + + reserved, err = TryReserveTokenQuota(token.Id, token.Key, 56, false) + require.NoError(t, err) + assert.False(t, reserved) + assert.Equal(t, 55, getTokenFromDB(t, token.Id).RemainQuota) +} + +func TestRedisBatchReserveNeverFallsBackToStaleDatabaseBalance(t *testing.T) { + truncateTables(t) + resetBatchUpdateTestState(t) + useUserCacheMiniRedis(t) + common.BatchUpdateEnabled = true + + user := createReserveTestUser(t, 10) + reserved, err := TryReserveUserQuota(user.Id, 8) + require.NoError(t, err) + assert.True(t, reserved) + assert.Equal(t, 10, getUserQuotaFromDB(t, user.Id), "batch delta is not flushed yet") + + reserved, err = TryReserveUserQuota(user.Id, 3) + require.NoError(t, err) + assert.False(t, reserved, "stale DB balance must not authorize a second spend") + cachedUser, err := GetUserCache(user.Id) + require.NoError(t, err) + assert.Equal(t, 2, cachedUser.Quota) + + token := createReserveTestToken(t, 9) + reserved, err = TryReserveTokenQuota(token.Id, token.Key, 7, false) + require.NoError(t, err) + assert.True(t, reserved) + reserved, err = TryReserveTokenQuota(token.Id, token.Key, 3, false) + require.NoError(t, err) + assert.False(t, reserved) + assert.Equal(t, 9, getTokenFromDB(t, token.Id).RemainQuota) + + batchUpdate() + assert.Equal(t, 2, getUserQuotaFromDB(t, user.Id)) + reloadedToken := getTokenFromDB(t, token.Id) + assert.Equal(t, 2, reloadedToken.RemainQuota) + assert.Equal(t, 7, reloadedToken.UsedQuota) +} + +func TestReserveFallsBackToDatabaseWhenRedisIsUnavailable(t *testing.T) { + truncateTables(t) + resetBatchUpdateTestState(t) + server := useUserCacheMiniRedis(t) + + user := createReserveTestUser(t, 20) + require.NoError(t, populateUserCache(user)) + server.Close() + + // Redis 故障时降级为数据库条件更新：服务保持可用且不会超扣。 + reserved, err := TryReserveUserQuota(user.Id, 5) + require.NoError(t, err) + assert.True(t, reserved) + assert.Equal(t, 15, getUserQuotaFromDB(t, user.Id)) + + reserved, err = TryReserveUserQuota(user.Id, 16) + require.NoError(t, err) + assert.False(t, reserved) + assert.Equal(t, 15, getUserQuotaFromDB(t, user.Id)) +} + +func TestSynchronousReserveCompensatesCacheWhenPersistenceFails(t *testing.T) { + truncateTables(t) + resetBatchUpdateTestState(t) + useUserCacheMiniRedis(t) + + user := createReserveTestUser(t, 10) + require.NoError(t, populateUserCache(user)) + require.NoError(t, DB.Delete(&user).Error) + + reserved, err := TryReserveUserQuota(user.Id, 6) + assert.False(t, reserved) + assert.ErrorIs(t, err, gorm.ErrRecordNotFound) + cached, cacheErr := cacheGetUserBase(user.Id) + require.NoError(t, cacheErr) + assert.Equal(t, 10, cached.Quota) + + token := createReserveTestToken(t, 12) + _, err = GetTokenByKey(token.Key, true) + require.NoError(t, err) + require.NoError(t, DB.Delete(&token).Error) + reserved, err = TryReserveTokenQuota(token.Id, token.Key, 7, false) + assert.False(t, reserved) + assert.ErrorIs(t, err, gorm.ErrRecordNotFound) + cachedToken, cacheErr := cacheGetTokenByKey(token.Key) + require.NoError(t, cacheErr) + assert.Equal(t, 12, cachedToken.RemainQuota) + assert.Zero(t, cachedToken.UsedQuota) +} + +func TestTokenCacheInitPreservesLiveQuotaAndFenceBlocksStaleSnapshot(t *testing.T) { + truncateTables(t) + resetBatchUpdateTestState(t) + server := useUserCacheMiniRedis(t) + + token := createReserveTestToken(t, 100) + loaded, err := GetTokenByKey(token.Key, true) + require.NoError(t, err) + stale := *loaded + + result, err := cacheApplyTokenQuotaDelta(token.Id, token.Key, -70) + require.NoError(t, err) + require.Equal(t, cacheQuotaOK, result) + + // 已存在的哈希只刷新 TTL：数据库快照不得覆盖已被原子预扣的余额。 + code, err := cacheInitToken(stale) + require.NoError(t, err) + assert.Equal(t, 2, code) + cached, err := cacheGetTokenByKey(token.Key) + require.NoError(t, err) + assert.Equal(t, 30, cached.RemainQuota) + + // 变更期间：fence 删除缓存并拦截并发读者手中的过期快照。 + require.NoError(t, invalidateTokenCacheForMutation(token.Key)) + code, err = cacheInitToken(stale) + require.NoError(t, err) + assert.Zero(t, code, "the pre-mutation snapshot must not be published while fenced") + _, err = cacheGetTokenByKey(token.Key) + assert.Error(t, err) + + // fence 过期后可重新从数据库水合。 + server.FastForward(time.Duration(tokenCacheFenceSeconds+1) * time.Second) + fresh, err := GetTokenByKey(token.Key, false) + require.NoError(t, err) + assert.Equal(t, 100, fresh.RemainQuota) + cached, err = cacheGetTokenByKey(token.Key) + require.NoError(t, err) + assert.Equal(t, 100, cached.RemainQuota) +} diff --git a/model/subscription.go b/model/subscription.go index 497fea14..c89e63bf 100644 --- a/model/subscription.go +++ b/model/subscription.go @@ -600,6 +600,12 @@ func CompleteSubscriptionOrder(tradeNo string, providerPayload string, expectedP if !plan.Enabled { // still allow completion for already purchased orders } + // 锁定用户行：并发完成同一用户的不同订单（包括多实例部署下）时， + // 使 CreateUserSubscriptionFromPlanTx 的 MaxPurchasePerUser 检查按用户串行。 + var userRow User + if err := lockForUpdate(tx).Select("id").Where("id = ?", order.UserId).First(&userRow).Error; err != nil { + return err + } subscription, err := CreateUserSubscriptionFromPlanTx(tx, order.UserId, plan, "order") if err != nil { return err @@ -712,6 +718,11 @@ func AdminBindSubscription(userId int, planId int, sourceNote string) (string, e } groupChanged := false err = DB.Transaction(func(tx *gorm.DB) error { + // 与 CompleteSubscriptionOrder 一致：先锁用户行，再做购买次数检查。 + var userRow User + if err := lockForUpdate(tx).Select("id").Where("id = ?", userId).First(&userRow).Error; err != nil { + return err + } subscription, err := CreateUserSubscriptionFromPlanTx(tx, userId, plan, "admin") if err == nil { groupChanged = subscription.PrevUserGroup != "" @@ -737,9 +748,8 @@ func calcSubscriptionBalanceQuota(priceAmount float64) (int, error) { } quota := decimal.NewFromFloat(priceAmount). Mul(decimal.NewFromFloat(common.QuotaPerUnit)). - Ceil(). - IntPart() - return int(quota), nil + Ceil() + return common.QuotaFromDecimalStrict(quota) } // PurchaseSubscriptionWithBalance creates a subscription by deducting the user's wallet quota. diff --git a/model/token.go b/model/token.go index 5aa8b3d5..95cbabae 100644 --- a/model/token.go +++ b/model/token.go @@ -274,27 +274,10 @@ func GetTokenById(id int) (*Token, error) { token := Token{Id: id} var err error = nil err = DB.First(&token, "id = ?", id).Error - if shouldUpdateRedis(true, err) { - gopool.Go(func() { - if err := cacheSetToken(token); err != nil { - common.SysLog("failed to update user status cache: " + err.Error()) - } - }) - } return &token, err } func GetTokenByKey(key string, fromDB bool) (token *Token, err error) { - defer func() { - // Update Redis cache asynchronously on successful DB read - if shouldUpdateRedis(fromDB, err) && token != nil { - gopool.Go(func() { - if err := cacheSetToken(*token); err != nil { - common.SysLog("failed to update user status cache: " + err.Error()) - } - }) - } - }() if !fromDB && common.RedisEnabled { // Try Redis first token, err := cacheGetTokenByKey(key) @@ -303,9 +286,18 @@ func GetTokenByKey(key string, fromDB bool) (token *Token, err error) { } // Don't return error - fall through to DB } - fromDB = true - err = DB.Where(commonKeyCol+" = ?", key).First(&token).Error - return token, err + token = &Token{} + if err = DB.Where(commonKeyCol+" = ?", key).First(token).Error; err != nil { + return nil, err + } + if common.RedisEnabled { + // 冷缓存时用数据库快照初始化；已存在的哈希只刷新 TTL， + // 避免快照覆盖 Redis 中已被原子预扣的余额。初始化失败不影响本次读取。 + if _, cacheErr := cacheInitToken(*token); cacheErr != nil { + common.SysLog("failed to init token cache: " + cacheErr.Error()) + } + } + return token, nil } func (token *Token) Insert() error { @@ -316,47 +308,27 @@ func (token *Token) Insert() error { // Update Make sure your token's fields is completed, because this will update non-zero values func (token *Token) Update() (err error) { - err = DB.Model(token).Select("name", "status", "expired_time", "remain_quota", "unlimited_quota", - "model_limits_enabled", "model_limits", "allow_ips", "group", "cross_group_retry", "auto_groups").Updates(token).Error - if shouldUpdateRedis(true, err) { - if cacheErr := cacheSetToken(*token); cacheErr != nil { - common.SysLog("failed to update token cache: " + cacheErr.Error()) - if deleteErr := cacheDeleteToken(token.Key); deleteErr != nil { - common.SysLog("failed to invalidate token cache after update: " + deleteErr.Error()) - } - } + // 写库前失效缓存并设置 fence，防止并发读者把过期快照重新写回缓存。 + if cacheErr := invalidateTokenCacheForMutation(token.Key); cacheErr != nil { + common.SysLog("failed to invalidate token cache before update: " + cacheErr.Error()) } - return err + return DB.Model(token).Select("name", "status", "expired_time", "remain_quota", "unlimited_quota", + "model_limits_enabled", "model_limits", "allow_ips", "group", "cross_group_retry", "auto_groups").Updates(token).Error } func (token *Token) SelectUpdate() (err error) { - defer func() { - if shouldUpdateRedis(true, err) { - gopool.Go(func() { - err := cacheSetToken(*token) - if err != nil { - common.SysLog("failed to update token cache: " + err.Error()) - } - }) - } - }() + if cacheErr := invalidateTokenCacheForMutation(token.Key); cacheErr != nil { + common.SysLog("failed to invalidate token cache before status update: " + cacheErr.Error()) + } // This can update zero values return DB.Model(token).Select("accessed_time", "status").Updates(token).Error } func (token *Token) Delete() (err error) { - defer func() { - if shouldUpdateRedis(true, err) { - gopool.Go(func() { - err := cacheDeleteToken(token.Key) - if err != nil { - common.SysLog("failed to delete token cache: " + err.Error()) - } - }) - } - }() - err = DB.Delete(token).Error - return err + if cacheErr := invalidateTokenCacheForMutation(token.Key); cacheErr != nil { + common.SysLog("failed to invalidate token cache before delete: " + cacheErr.Error()) + } + return DB.Delete(token).Error } func (token *Token) IsModelLimitsEnabled() bool { @@ -408,8 +380,9 @@ func IncreaseTokenQuota(tokenId int, key string, quota int) (err error) { } if common.RedisEnabled { gopool.Go(func() { - err := cacheIncrTokenQuota(key, int64(quota)) - if err != nil { + // 守卫式增量：哈希不存在时跳过，由下次读取从数据库水合， + // 绝不创建只有配额字段的残缺哈希。 + if _, err := cacheApplyTokenQuotaDelta(tokenId, key, int64(quota)); err != nil { common.SysLog("failed to increase token quota: " + err.Error()) } }) @@ -438,8 +411,7 @@ func DecreaseTokenQuota(id int, key string, quota int) (err error) { } if common.RedisEnabled { gopool.Go(func() { - err := cacheDecrTokenQuota(key, int64(quota)) - if err != nil { + if _, err := cacheApplyTokenQuotaDelta(id, key, int64(-quota)); err != nil { common.SysLog("failed to decrease token quota: " + err.Error()) } }) @@ -482,6 +454,9 @@ func BatchDeleteTokens(ids []int, userId int) (int, error) { tx.Rollback() return 0, err } + if err := invalidateTokensCache(tokens); err != nil { + common.SysLog("failed to invalidate token cache before batch delete: " + err.Error()) + } if err := tx.Where("user_id = ? AND id IN (?)", userId, ids).Delete(&Token{}).Error; err != nil { tx.Rollback() @@ -492,14 +467,6 @@ func BatchDeleteTokens(ids []int, userId int) (int, error) { return 0, err } - if common.RedisEnabled { - gopool.Go(func() { - for _, t := range tokens { - _ = cacheDeleteToken(t.Key) - } - }) - } - return len(tokens), nil } @@ -540,7 +507,7 @@ func invalidateTokensCache(tokens []Token) error { if t.Key == "" { continue } - if err := cacheDeleteToken(t.Key); err != nil && firstErr == nil { + if err := invalidateTokenCacheForMutation(t.Key); err != nil && firstErr == nil { firstErr = err } } diff --git a/model/token_auto_groups_cache_test.go b/model/token_auto_groups_cache_test.go index 2018502d..d748c4f8 100644 --- a/model/token_auto_groups_cache_test.go +++ b/model/token_auto_groups_cache_test.go @@ -19,7 +19,7 @@ func TestTokenAutoGroupsRoundTripThroughRedisHashCache(t *testing.T) { AutoGroups: `["vip","default"]`, } - require.NoError(t, cacheSetToken(token)) + require.NoError(t, cacheSetTokenForTest(token)) cached, err := cacheGetTokenByKey(token.Key) require.NoError(t, err) assert.Equal(t, token.AutoGroups, cached.AutoGroups) @@ -43,7 +43,7 @@ func TestTokenUpdateSynchronouslyNarrowsPreheatedAutoGroupsCache(t *testing.T) { AutoGroups: `["default","vip"]`, } require.NoError(t, token.Insert()) - require.NoError(t, cacheSetToken(token)) + require.NoError(t, cacheSetTokenForTest(token)) preheated, err := cacheGetTokenByKey(token.Key) require.NoError(t, err) @@ -51,7 +51,18 @@ func TestTokenUpdateSynchronouslyNarrowsPreheatedAutoGroupsCache(t *testing.T) { require.NoError(t, token.SetAutoGroups([]string{"vip"})) require.NoError(t, token.Update()) - immediate, err := cacheGetTokenByKey(token.Key) + // Update 是限制性变更：写库前删除缓存并设置 fence。缓存不再提供旧的 + // 宽分组值，下一次读取必须看到收紧后的分组。 + _, cacheErr := cacheGetTokenByKey(token.Key) + require.Error(t, cacheErr, "the pre-update cache entry must be invalidated") + reloaded, err := GetTokenByKey(token.Key, false) require.NoError(t, err) - assert.JSONEq(t, `["vip"]`, immediate.AutoGroups) + assert.JSONEq(t, `["vip"]`, reloaded.AutoGroups) +} + +// cacheSetTokenForTest 以测试身份写入完整 token 缓存（含额度字段）， +// 模拟“已水合”的缓存状态。 +func cacheSetTokenForTest(token Token) error { + _, err := cacheInitToken(token) + return err } diff --git a/model/token_cache.go b/model/token_cache.go index 947f587d..c1cdf53f 100644 --- a/model/token_cache.go +++ b/model/token_cache.go @@ -1,65 +1,107 @@ package model import ( + "context" "fmt" + "strconv" "time" "github.com/QuantumNous/new-api/common" - "github.com/QuantumNous/new-api/constant" ) -func cacheSetToken(token Token) error { - key := common.GenerateHMAC(token.Key) - token.Clean() - err := common.RedisHSetObj(fmt.Sprintf("token:%s", key), &token, time.Duration(common.RedisKeyCacheSeconds())*time.Second) +func getTokenCacheKey(key string) string { + return fmt.Sprintf("token:%s", common.GenerateHMAC(key)) +} + +func getTokenCacheFenceKey(key string) string { + return fmt.Sprintf("token:fence:%s", common.GenerateHMAC(key)) +} + +func tokenCacheTTLSeconds() int { + ttl := common.RedisKeyCacheSeconds() + if ttl <= 0 { + return 60 + } + return ttl +} + +// tokenCacheFenceSeconds must outlive a token mutation's database write plus +// any in-flight reader's DB-read-to-cache-init gap. The fence is not deleted +// after commit; it expires naturally so a reader holding a pre-mutation +// snapshot cannot publish it right after the mutation cleared the cache. +// While the fence exists readers simply serve the database without caching. +const tokenCacheFenceSeconds = 10 + +// invalidateTokenCacheForMutation is called before a token metadata mutation +// writes to the database: it raises the fence and drops the cached hash so no +// reader can act on (or re-publish) the pre-mutation state. +func invalidateTokenCacheForMutation(key string) error { + if !common.RedisEnabled || key == "" { + return nil + } + ctx := context.Background() + err := common.RDB.Set(ctx, getTokenCacheFenceKey(key), 1, time.Duration(tokenCacheFenceSeconds)*time.Second).Err() if err != nil { return err } - return nil + return common.RDB.Del(ctx, getTokenCacheKey(key)).Err() } -func cacheDeleteToken(key string) error { - key = common.GenerateHMAC(key) - err := common.RedisDelKey(fmt.Sprintf("token:%s", key)) - if err != nil { - return err +// cacheInitToken publishes a database snapshot only when no mutation fence is +// active and the hash is cold. An existing hash only gets its TTL refreshed: +// its RemainQuota may already be ahead of this snapshot because atomic +// pre-consume decrements Redis first, so a snapshot must never overwrite any +// field of a live hash. +// 返回值：0=被 fence 拦截，1=完成初始化，2=哈希已存在，仅刷新 TTL。 +func cacheInitToken(token Token) (int, error) { + if !common.RedisEnabled { + return 0, nil } - return nil -} - -func cacheIncrTokenQuota(key string, increment int64) error { - key = common.GenerateHMAC(key) - err := common.RedisHIncrBy(fmt.Sprintf("token:%s", key), constant.TokenFiledRemainQuota, increment) - if err != nil { - return err + allowIps := "" + if token.AllowIps != nil { + allowIps = *token.AllowIps } - return nil + const script = ` +if redis.call('EXISTS', KEYS[2]) == 1 then + return 0 +end +if redis.call('EXISTS', KEYS[1]) == 1 then + redis.call('EXPIRE', KEYS[1], ARGV[17]) + return 2 +end +redis.call('HSET', KEYS[1], + 'Id', ARGV[1], 'UserId', ARGV[2], 'Status', ARGV[3], 'Name', ARGV[4], + 'CreatedTime', ARGV[5], 'AccessedTime', ARGV[6], 'ExpiredTime', ARGV[7], + 'UnlimitedQuota', ARGV[8], 'ModelLimitsEnabled', ARGV[9], 'ModelLimits', ARGV[10], + 'AllowIps', ARGV[11], 'Group', ARGV[12], 'CrossGroupRetry', ARGV[13], + 'AutoGroups', ARGV[14], 'RemainQuota', ARGV[15], 'UsedQuota', ARGV[16]) +redis.call('EXPIRE', KEYS[1], ARGV[17]) +return 1` + + return common.RDB.Eval(context.Background(), script, []string{ + getTokenCacheKey(token.Key), getTokenCacheFenceKey(token.Key), + }, + token.Id, token.UserId, token.Status, token.Name, + token.CreatedTime, token.AccessedTime, token.ExpiredTime, + strconv.FormatBool(token.UnlimitedQuota), strconv.FormatBool(token.ModelLimitsEnabled), + token.ModelLimits, allowIps, token.Group, strconv.FormatBool(token.CrossGroupRetry), + token.AutoGroups, token.RemainQuota, token.UsedQuota, + tokenCacheTTLSeconds(), + ).Int() } -func cacheDecrTokenQuota(key string, decrement int64) error { - return cacheIncrTokenQuota(key, -decrement) -} - -func cacheSetTokenField(key string, field string, value string) error { - key = common.GenerateHMAC(key) - err := common.RedisHSetField(fmt.Sprintf("token:%s", key), field, value) - if err != nil { - return err - } - return nil -} - -// CacheGetTokenByKey 从缓存中获取 token，如果缓存中不存在，则从数据库中获取 +// cacheGetTokenByKey 从缓存读取 token；不完整的哈希（如仅有配额字段）会被拒绝。 func cacheGetTokenByKey(key string) (*Token, error) { - hmacKey := common.GenerateHMAC(key) if !common.RedisEnabled { return nil, fmt.Errorf("redis is not enabled") } var token Token - err := common.RedisHGetObj(fmt.Sprintf("token:%s", hmacKey), &token) - if err != nil { + if err := common.RedisHGetObj(getTokenCacheKey(key), &token); err != nil { return nil, err } + if token.Id <= 0 { + return nil, fmt.Errorf("token cache is incomplete") + } token.Key = key return &token, nil } diff --git a/model/user.go b/model/user.go index 83d7aeec..0a9a42a7 100644 --- a/model/user.go +++ b/model/user.go @@ -547,7 +547,7 @@ func inviteUser(inviterId int) error { func (user *User) TransferAffQuotaToQuota(quota int) error { // 检查quota是否小于最小额度 if float64(quota) < common.QuotaPerUnit { - return fmt.Errorf("转移额度最小为%s！", logger.LogQuota(int(common.QuotaPerUnit))) + return fmt.Errorf("转移额度最小为%s！", logger.LogQuota(common.QuotaFromFloat(common.QuotaPerUnit))) } // 开始数据库事务 @@ -1180,24 +1180,9 @@ func ValidateAccessToken(token string) (*User, error) { // GetUserQuota gets quota from Redis first, falls back to DB if needed func GetUserQuota(id int, fromDB bool) (quota int, err error) { - defer func() { - // Update Redis cache asynchronously on successful DB read - if shouldUpdateRedis(fromDB, err) { - gopool.Go(func() { - if err := updateUserQuotaCache(id, quota); err != nil { - common.SysLog("failed to update user quota cache: " + err.Error()) - } - }) - } - }() if !fromDB && common.RedisEnabled { - quota, err := getUserQuotaCache(id) - if err == nil { - return quota, nil - } - // Don't return error - fall through to DB + return getUserQuotaCache(id) } - fromDB = true err = DB.Model(&User{}).Where("id = ?", id).Select("quota").Find("a).Error if err != nil { return 0, err @@ -1403,24 +1388,6 @@ func updateUserQuotaUsedQuotaAndRequestCount(id int, quota int, usedQuota int, r } } -func updateUserUsedQuota(id int, quota int) { - err := DB.Model(&User{}).Where("id = ?", id).Updates( - map[string]interface{}{ - "used_quota": gorm.Expr("used_quota + ?", quota), - }, - ).Error - if err != nil { - common.SysLog("failed to update user used quota: " + err.Error()) - } -} - -func updateUserRequestCount(id int, count int) { - err := DB.Model(&User{}).Where("id = ?", id).Update("request_count", gorm.Expr("request_count + ?", count)).Error - if err != nil { - common.SysLog("failed to update user request count: " + err.Error()) - } -} - // GetUsernameById gets username from Redis first, falls back to DB if needed func GetUsernameById(id int, fromDB bool) (username string, err error) { defer func() { diff --git a/model/user_cache.go b/model/user_cache.go index 7ca31497..e8996e9b 100644 --- a/model/user_cache.go +++ b/model/user_cache.go @@ -67,12 +67,6 @@ func invalidateUserCache(userId int) error { return common.RedisDelKey(getUserCacheKey(userId)) } -// InvalidateUserCache is the exported version of invalidateUserCache. -// 供 controller 等上层包在用户状态变更（如禁用、删除、角色变更）后主动清理缓存。 -func InvalidateUserCache(userId int) error { - return invalidateUserCache(userId) -} - func populateUserCache(user User) error { if !common.RedisEnabled { return nil @@ -187,14 +181,6 @@ func getUserQuotaCache(userId int) (int, error) { return cache.Quota, nil } -func getUserStatusCache(userId int) (int, error) { - cache, err := GetUserCache(userId) - if err != nil { - return 0, err - } - return cache.Status, nil -} - func getUserNameCache(userId int) (string, error) { cache, err := GetUserCache(userId) if err != nil { @@ -211,22 +197,6 @@ func getUserSettingCache(userId int) (dto.UserSetting, error) { return cache.GetSetting(), nil } -// New functions for individual field updates -func updateUserStatusCache(userId int, status bool) error { - statusInt := common.UserStatusEnabled - if !status { - statusInt = common.UserStatusDisabled - } - return updateUserCacheField(userId, "Status", statusInt) -} - -func updateUserQuotaCache(userId int, quota int) error { - if !common.RedisEnabled { - return nil - } - return common.RedisHSetField(getUserCacheKey(userId), "Quota", fmt.Sprintf("%d", quota)) -} - // RefreshUserGroupCache writes the database-authoritative group into an // existing user hash without changing the user's authentication version. func RefreshUserGroupCache(userId int) error { diff --git a/service/billing_session.go b/service/billing_session.go index 42b0ffbc..afc706a7 100644 --- a/service/billing_session.go +++ b/service/billing_session.go @@ -1,6 +1,7 @@ package service import ( + "errors" "fmt" "net/http" "strings" @@ -214,6 +215,16 @@ func (s *BillingSession) preConsume(c *gin.Context, quota int) *types.NewAPIErro s.tokenConsumed = 0 } // TODO: model 层应定义哨兵错误（如 ErrNoActiveSubscription），用 errors.Is 替代字符串匹配 + if errors.Is(err, ErrInsufficientWalletQuota) { + userQuota, quotaErr := model.GetUserQuota(s.relayInfo.UserId, false) + if quotaErr != nil { + userQuota = 0 + } + return types.NewErrorWithStatusCode( + fmt.Errorf("用户额度不足, 剩余额度: %s", logger.FormatQuota(userQuota)), + types.ErrorCodeInsufficientUserQuota, http.StatusForbidden, + types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog()) + } errMsg := err.Error() if strings.Contains(errMsg, "no active subscription") || strings.Contains(errMsg, "subscription quota insufficient") { return types.NewErrorWithStatusCode(fmt.Errorf("订阅额度不足或未配置订阅: %s", errMsg), types.ErrorCodeInsufficientUserQuota, http.StatusForbidden, types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog()) diff --git a/service/funding_source.go b/service/funding_source.go index fc629c83..7045704b 100644 --- a/service/funding_source.go +++ b/service/funding_source.go @@ -1,6 +1,7 @@ package service import ( + "errors" "time" "github.com/QuantumNous/new-api/model" @@ -26,6 +27,11 @@ type FundingSource interface { // WalletFunding — 钱包资金来源实现 // --------------------------------------------------------------------------- +// ErrInsufficientWalletQuota 钱包原子预扣失败（余额不足），未发生任何扣减。 +// BillingSession 据此映射为 ErrorCodeInsufficientUserQuota， +// 使 wallet_first 等计费偏好可以回退到订阅。 +var ErrInsufficientWalletQuota = errors.New("wallet quota insufficient") + type WalletFunding struct { userId int consumed int // 实际预扣的用户额度 @@ -37,9 +43,13 @@ func (w *WalletFunding) PreConsume(amount int) error { if amount <= 0 { return nil } - if err := model.DecreaseUserQuota(w.userId, amount, false); err != nil { + reserved, err := model.TryReserveUserQuota(w.userId, amount) + if err != nil { return err } + if !reserved { + return ErrInsufficientWalletQuota + } w.consumed = amount return nil } diff --git a/service/quota.go b/service/quota.go index 1f3a9ef2..0c0c2da4 100644 --- a/service/quota.go +++ b/service/quota.go @@ -391,19 +391,17 @@ func PreConsumeTokenQuota(relayInfo *relaycommon.RelayInfo, quota int) error { if relayInfo.IsPlayground { return nil } - //if relayInfo.TokenUnlimited { - // return nil - //} - token, err := model.GetTokenByKey(relayInfo.TokenKey, false) + // 原子预扣：检查与扣减在同一操作中完成，并发请求不可能同时通过检查后超扣。 + reserved, err := model.TryReserveTokenQuota(relayInfo.TokenId, relayInfo.TokenKey, quota, relayInfo.TokenUnlimited) if err != nil { return err } - if !relayInfo.TokenUnlimited && token.RemainQuota < quota { - return fmt.Errorf("token quota is not enough, token remain quota: %s, need quota: %s", logger.FormatQuota(token.RemainQuota), logger.FormatQuota(quota)) - } - err = model.DecreaseTokenQuota(relayInfo.TokenId, relayInfo.TokenKey, quota) - if err != nil { - return err + if !reserved { + remainQuota := 0 + if token, tokenErr := model.GetTokenByKey(relayInfo.TokenKey, false); tokenErr == nil && token != nil { + remainQuota = token.RemainQuota + } + return fmt.Errorf("token quota is not enough, token remain quota: %s, need quota: %s", logger.FormatQuota(remainQuota), logger.FormatQuota(quota)) } return nil }