69 lines
2.1 KiB
Go
69 lines
2.1 KiB
Go
package model
|
|
|
|
import (
|
|
"context"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
)
|
|
|
|
type cacheQuotaResult int
|
|
|
|
const (
|
|
cacheQuotaInsufficient cacheQuotaResult = iota
|
|
cacheQuotaOK
|
|
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])
|
|
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`
|
|
|
|
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)
|
|
}
|