feat: migrate lua script to lua file (#4069)

This commit is contained in:
fearlessfei
2024-04-17 23:20:10 +08:00
committed by GitHub
parent 36088ea0d4
commit 62c88a84d1
10 changed files with 102 additions and 74 deletions

View File

@@ -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,