package model import ( "context" "errors" "fmt" "github.com/QuantumNous/new-api/common" "gorm.io/gorm" ) type cacheQuotaResult int const ( cacheQuotaInsufficient cacheQuotaResult = iota cacheQuotaOK cacheQuotaMiss ) 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]) or redis.call('HEXISTS', KEYS[1], 'Quota') == 0 then return -1 end local quota = tonumber(redis.call('HGET', KEYS[1], 'Quota')) if quota == nil or quota < tonumber(ARGV[1]) then return 0 end redis.call('HINCRBY', KEYS[1], 'Quota', -tonumber(ARGV[1])) return 1` const userQuotaDeltaScript = ` 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]) or redis.call('HEXISTS', KEYS[1], 'Quota') == 0 then return -1 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 } switch result { case 1: return cacheQuotaOK, nil case 0: return cacheQuotaInsufficient, nil default: return cacheQuotaMiss, nil } } func cacheTryReserveUserQuota(userID int, amount int64) (cacheQuotaResult, error) { result, err := common.RDB.Eval(context.Background(), userQuotaReserveScript, []string{getUserCacheKey(userID)}, amount, userID, userCacheSchemaVersion).Int() return quotaResultFromLua(result, err) } func cacheApplyUserQuotaDelta(userID int, delta int64) (cacheQuotaResult, error) { result, err := common.RDB.Eval(context.Background(), userQuotaDeltaScript, []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 }