mirror of
https://github.com/zeromicro/go-zero.git
synced 2026-05-14 10:20: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"
|
||||
"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()),
|
||||
})
|
||||
|
||||
Reference in New Issue
Block a user