mirror of
https://github.com/zeromicro/go-zero.git
synced 2026-05-13 18:00:00 +08:00
feat: migrate lua script to lua file (#4069)
This commit is contained in:
@@ -2,6 +2,7 @@ package limit
|
||||
|
||||
import (
|
||||
"context"
|
||||
_ "embed"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strconv"
|
||||
@@ -20,37 +21,12 @@ const (
|
||||
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])
|
||||
local fill_time = capacity/rate
|
||||
local ttl = math.floor(fill_time*2)
|
||||
local last_tokens = tonumber(redis.call("get", KEYS[1]))
|
||||
if last_tokens == nil then
|
||||
last_tokens = capacity
|
||||
end
|
||||
var (
|
||||
//go:embed tokenscript.lua
|
||||
tokenScript string
|
||||
|
||||
local last_refreshed = tonumber(redis.call("get", KEYS[2]))
|
||||
if last_refreshed == nil then
|
||||
last_refreshed = 0
|
||||
end
|
||||
|
||||
local delta = math.max(0, now-last_refreshed)
|
||||
local filled_tokens = math.min(capacity, last_tokens+(delta*rate))
|
||||
local allowed = filled_tokens >= requested
|
||||
local new_tokens = filled_tokens
|
||||
if allowed then
|
||||
new_tokens = filled_tokens - requested
|
||||
end
|
||||
|
||||
redis.call("setex", KEYS[1], ttl, new_tokens)
|
||||
redis.call("setex", KEYS[2], ttl, now)
|
||||
|
||||
return allowed`)
|
||||
scriptToken = redis.NewScript(tokenScript)
|
||||
)
|
||||
|
||||
// A TokenLimiter controls how frequently events are allowed to happen with in one second.
|
||||
type TokenLimiter struct {
|
||||
@@ -112,7 +88,7 @@ func (lim *TokenLimiter) reserveN(ctx context.Context, now time.Time, n int) boo
|
||||
}
|
||||
|
||||
resp, err := lim.store.ScriptRunCtx(ctx,
|
||||
script,
|
||||
scriptToken,
|
||||
[]string{
|
||||
lim.tokenKey,
|
||||
lim.timestampKey,
|
||||
|
||||
Reference in New Issue
Block a user