zeromicro/go-zero
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat(redis): add ScriptRun API and migrate EvalCtx to ScriptRun for limit, lock and bloom (#3087)

Browse Source
This commit is contained in:
cong
2023-04-02 11:28:25 +08:00
committed by GitHub
parent b49fc81618
commit 5da8a93c75
6 changed files with 96 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files

38
core/limit/periodlimit.go
View File

@@ -9,21 +9,6 @@ import (
"github.com/zeromicro/go-zero/core/stores/redis"
)
// to be compatible with aliyun redis, we cannot use `local key = KEYS[1]` to reuse the key
const periodScript = `local limit = tonumber(ARGV[1])
local window = tonumber(ARGV[2])
local current = redis.call("INCRBY", KEYS[1], 1)
if current == 1 then
redis.call("expire", KEYS[1], window)
end
if current < limit then
return 1
elseif current == limit then
return 2
else
return 0
end`
const (
// Unknown means not initialized state.
Unknown = iota
@@ -39,8 +24,25 @@ const (
internalHitQuota = 2
)
// ErrUnknownCode is an error that represents unknown status code.
var ErrUnknownCode = errors.New("unknown status code")
var (
// ErrUnknownCode is an error that represents unknown status code.
ErrUnknownCode = errors.New("unknown status code")
// to be compatible with aliyun redis, we cannot use `local key = KEYS[1]` to reuse the key
periodScript = redis.NewScript(`local limit = tonumber(ARGV[1])
local window = tonumber(ARGV[2])
local current = redis.call("INCRBY", KEYS[1], 1)
if current == 1 then
redis.call("expire", KEYS[1], window)
end
if current < limit then
return 1
elseif current == limit then
return 2
else
return 0
end`)
)
type (
// PeriodOption defines the method to customize a PeriodLimit.
@@ -80,7 +82,7 @@ func (h *PeriodLimit) Take(key string) (int, error) {
// TakeCtx requests a permit with context, it returns the permit state.
func (h *PeriodLimit) TakeCtx(ctx context.Context, key string) (int, error) {
resp, err := h.limitStore.EvalCtx(ctx, periodScript, []string{h.keyPrefix + key}, []string{
resp, err := h.limitStore.ScriptRunCtx(ctx, periodScript, []string{h.keyPrefix + key}, []string{
strconv.Itoa(h.quota),
strconv.Itoa(h.calcExpireSeconds()),
})

21
core/limit/tokenlimit.go
View File

@@ -15,10 +15,15 @@ import (
)
const (
// to be compatible with aliyun redis, we cannot use `local key = KEYS[1]` to reuse the key
// KEYS[1] as tokens_key
// KEYS[2] as timestamp_key
script = `local rate = tonumber(ARGV[1])
tokenFormat = "{%s}.tokens"
timestampFormat = "{%s}.ts"
pingInterval = time.Millisecond * 100
)
// to be compatible with aliyun redis, we cannot use `local key = KEYS[1]` to reuse the key
// KEYS[1] as tokens_key
// KEYS[2] as timestamp_key
var script = redis.NewScript(`local rate = tonumber(ARGV[1])
local capacity = tonumber(ARGV[2])
local now = tonumber(ARGV[3])
local requested = tonumber(ARGV[4])
@@ -45,11 +50,7 @@ end
redis.call("setex", KEYS[1], ttl, new_tokens)
redis.call("setex", KEYS[2], ttl, now)
return allowed`
tokenFormat = "{%s}.tokens"
timestampFormat = "{%s}.ts"
pingInterval = time.Millisecond * 100
)
return allowed`)
// A TokenLimiter controls how frequently events are allowed to happen with in one second.
type TokenLimiter struct {
@@ -110,7 +111,7 @@ func (lim *TokenLimiter) reserveN(ctx context.Context, now time.Time, n int) boo
return lim.rescueLimiter.AllowN(now, n)
}
resp, err := lim.store.EvalCtx(ctx,
resp, err := lim.store.ScriptRunCtx(ctx,
script,
[]string{
lim.tokenKey,