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"
"strconv"
"time"
@@ -28,20 +29,10 @@ 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`)
//go:embed periodscript.lua
periodScript string
scriptPeriod = redis.NewScript(periodScript)
)
type (
@@ -82,7 +73,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.ScriptRunCtx(ctx, periodScript, []string{h.keyPrefix + key}, []string{
resp, err := h.limitStore.ScriptRunCtx(ctx, scriptPeriod, []string{h.keyPrefix + key}, []string{
strconv.Itoa(h.quota),
strconv.Itoa(h.calcExpireSeconds()),
})