From 62c88a84d1a7772078f4fe64894cac1b566a1cba Mon Sep 17 00:00:00 2001 From: fearlessfei <573088370@qq.com> Date: Wed, 17 Apr 2024 23:20:10 +0800 Subject: [PATCH] feat: migrate lua script to lua file (#4069) --- core/bloom/bloom.go | 27 ++++++++++------------- core/bloom/setscript.lua | 3 +++ core/bloom/testscript.lua | 6 +++++ core/limit/periodlimit.go | 21 +++++------------- core/limit/periodscript.lua | 14 ++++++++++++ core/limit/tokenlimit.go | 38 ++++++-------------------------- core/limit/tokenscript.lua | 31 ++++++++++++++++++++++++++ core/stores/redis/delscript.lua | 5 +++++ core/stores/redis/lockscript.lua | 6 +++++ core/stores/redis/redislock.go | 25 ++++++++++----------- 10 files changed, 102 insertions(+), 74 deletions(-) create mode 100644 core/bloom/setscript.lua create mode 100644 core/bloom/testscript.lua create mode 100644 core/limit/periodscript.lua create mode 100644 core/limit/tokenscript.lua create mode 100644 core/stores/redis/delscript.lua create mode 100644 core/stores/redis/lockscript.lua diff --git a/core/bloom/bloom.go b/core/bloom/bloom.go index 3a72cf84d..e47eb86f5 100644 --- a/core/bloom/bloom.go +++ b/core/bloom/bloom.go @@ -2,6 +2,7 @@ package bloom import ( "context" + _ "embed" "errors" "strconv" @@ -17,19 +18,15 @@ var ( // ErrTooLargeOffset indicates the offset is too large in bitset. ErrTooLargeOffset = errors.New("too large offset") - setScript = redis.NewScript(` -for _, offset in ipairs(ARGV) do - redis.call("setbit", KEYS[1], offset, 1) -end -`) - testScript = redis.NewScript(` -for _, offset in ipairs(ARGV) do - if tonumber(redis.call("getbit", KEYS[1], offset)) == 0 then - return false - end -end -return true -`) + //go:embed setscript.lua + setScript string + + scriptSet = redis.NewScript(setScript) + + //go:embed testscript.lua + testScript string + + scriptTest = redis.NewScript(testScript) ) type ( @@ -129,7 +126,7 @@ func (r *redisBitSet) check(ctx context.Context, offsets []uint) (bool, error) { return false, err } - resp, err := r.store.ScriptRunCtx(ctx, testScript, []string{r.key}, args) + resp, err := r.store.ScriptRunCtx(ctx, scriptTest, []string{r.key}, args) if errors.Is(err, redis.Nil) { return false, nil } else if err != nil { @@ -161,7 +158,7 @@ func (r *redisBitSet) set(ctx context.Context, offsets []uint) error { return err } - _, err = r.store.ScriptRunCtx(ctx, setScript, []string{r.key}, args) + _, err = r.store.ScriptRunCtx(ctx, scriptSet, []string{r.key}, args) if errors.Is(err, redis.Nil) { return nil } diff --git a/core/bloom/setscript.lua b/core/bloom/setscript.lua new file mode 100644 index 000000000..6d9d60f2d --- /dev/null +++ b/core/bloom/setscript.lua @@ -0,0 +1,3 @@ +for _, offset in ipairs(ARGV) do + redis.call("setbit", KEYS[1], offset, 1) +end diff --git a/core/bloom/testscript.lua b/core/bloom/testscript.lua new file mode 100644 index 000000000..2eda5784f --- /dev/null +++ b/core/bloom/testscript.lua @@ -0,0 +1,6 @@ +for _, offset in ipairs(ARGV) do + if tonumber(redis.call("getbit", KEYS[1], offset)) == 0 then + return false + end +end +return true diff --git a/core/limit/periodlimit.go b/core/limit/periodlimit.go index 332fe3c95..070bb7f81 100644 --- a/core/limit/periodlimit.go +++ b/core/limit/periodlimit.go @@ -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()), }) diff --git a/core/limit/periodscript.lua b/core/limit/periodscript.lua new file mode 100644 index 000000000..c351c8406 --- /dev/null +++ b/core/limit/periodscript.lua @@ -0,0 +1,14 @@ +-- to be compatible with aliyun redis, we cannot use `local key = KEYS[1]` to reuse the key +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 diff --git a/core/limit/tokenlimit.go b/core/limit/tokenlimit.go index e33b4b951..0070eca4a 100644 --- a/core/limit/tokenlimit.go +++ b/core/limit/tokenlimit.go @@ -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, diff --git a/core/limit/tokenscript.lua b/core/limit/tokenscript.lua new file mode 100644 index 000000000..b70ad1eef --- /dev/null +++ b/core/limit/tokenscript.lua @@ -0,0 +1,31 @@ +-- 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 +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 + +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 diff --git a/core/stores/redis/delscript.lua b/core/stores/redis/delscript.lua new file mode 100644 index 000000000..e9aa40a2a --- /dev/null +++ b/core/stores/redis/delscript.lua @@ -0,0 +1,5 @@ +if redis.call("GET", KEYS[1]) == ARGV[1] then + return redis.call("DEL", KEYS[1]) +else + return 0 +end diff --git a/core/stores/redis/lockscript.lua b/core/stores/redis/lockscript.lua new file mode 100644 index 000000000..0fd758386 --- /dev/null +++ b/core/stores/redis/lockscript.lua @@ -0,0 +1,6 @@ +if redis.call("GET", KEYS[1]) == ARGV[1] then + redis.call("SET", KEYS[1], ARGV[1], "PX", ARGV[2]) + return "OK" +else + return redis.call("SET", KEYS[1], ARGV[1], "NX", "PX", ARGV[2]) +end diff --git a/core/stores/redis/redislock.go b/core/stores/redis/redislock.go index ab26cc985..28b704b83 100644 --- a/core/stores/redis/redislock.go +++ b/core/stores/redis/redislock.go @@ -2,6 +2,7 @@ package redis import ( "context" + _ "embed" "errors" "math/rand" "strconv" @@ -20,17 +21,15 @@ const ( ) var ( - lockScript = NewScript(`if redis.call("GET", KEYS[1]) == ARGV[1] then - redis.call("SET", KEYS[1], ARGV[1], "PX", ARGV[2]) - return "OK" -else - return redis.call("SET", KEYS[1], ARGV[1], "NX", "PX", ARGV[2]) -end`) - delScript = NewScript(`if redis.call("GET", KEYS[1]) == ARGV[1] then - return redis.call("DEL", KEYS[1]) -else - return 0 -end`) + //go:embed lockscript.lua + lockScript string + + scriptLock = NewScript(lockScript) + + //go:embed delscript.lua + delScript string + + scriptDel = NewScript(delScript) ) // A RedisLock is a redis lock. @@ -62,7 +61,7 @@ func (rl *RedisLock) Acquire() (bool, error) { // AcquireCtx acquires the lock with the given ctx. func (rl *RedisLock) AcquireCtx(ctx context.Context) (bool, error) { seconds := atomic.LoadUint32(&rl.seconds) - resp, err := rl.store.ScriptRunCtx(ctx, lockScript, []string{rl.key}, []string{ + resp, err := rl.store.ScriptRunCtx(ctx, scriptLock, []string{rl.key}, []string{ rl.id, strconv.Itoa(int(seconds)*millisPerSecond + tolerance), }) if errors.Is(err, red.Nil) { @@ -90,7 +89,7 @@ func (rl *RedisLock) Release() (bool, error) { // ReleaseCtx releases the lock with the given ctx. func (rl *RedisLock) ReleaseCtx(ctx context.Context) (bool, error) { - resp, err := rl.store.ScriptRunCtx(ctx, delScript, []string{rl.key}, []string{rl.id}) + resp, err := rl.store.ScriptRunCtx(ctx, scriptDel, []string{rl.key}, []string{rl.id}) if err != nil { return false, err }