Repository: zeromicro/go-zero
Pull request #4069 (2024-04-11 09:59:56)

feat: migrate lua script to lua file
Reasons to move：
Better coding experience, syntax prompts highlighted, avoid some wrong.

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/zeromicro/go-zero/pull/4069?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=zeromicro) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 95.17%. Comparing base [(`8690859`)](https://app.codecov.io/gh/zeromicro/go-zero/commit/8690859c7d1c9461bdc422c3c65608bde30904db?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=zeromicro) to head [(`fa6c871`)](https://app.codecov.io/gh/zeromicro/go-zero/pull/4069?dropdown=coverage&src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=zeromicro).
> Report is 13 commits behind head on master.

<details><summary>Additional details and impacted files</summary>



| [Files](https://app.codecov.io/gh/zeromicro/go-zero/pull/4069?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=zeromicro) | Coverage Δ | |
|---|---|---|
| [core/bloom/bloom.go](https://app.codecov.io/gh/zeromicro/go-zero/pull/4069?src=pr&el=tree&filepath=core%2Fbloom%2Fbloom.go&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=zeromicro#diff-Y29yZS9ibG9vbS9ibG9vbS5nbw==) | `96.72% <100.00%> (+0.88%)` | :arrow_up: |
| [core/limit/periodlimit.go](https://app.codecov.io/gh/zeromicro/go-zero/pull/4069?src=pr&el=tree&filepath=core%2Flimit%2Fperiodlimit.go&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=zeromicro#diff-Y29yZS9saW1pdC9wZXJpb2RsaW1pdC5nbw==) | `90.00% <100.00%> (+0.63%)` | :arrow_up: |
| [core/limit/tokenlimit.go](https://app.codecov.io/gh/zeromicro/go-zero/pull/4069?src=pr&el=tree&filepath=core%2Flimit%2Ftokenlimit.go&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=zeromicro#diff-Y29yZS9saW1pdC90b2tlbmxpbWl0Lmdv) | `94.20% <100.00%> (+0.30%)` | :arrow_up: |
| [core/stores/redis/redislock.go](https://app.codecov.io/gh/zeromicro/go-zero/pull/4069?src=pr&el=tree&filepath=core%2Fstores%2Fredis%2Fredislock.go&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=zeromicro#diff-Y29yZS9zdG9yZXMvcmVkaXMvcmVkaXNsb2NrLmdv) | `84.21% <100.00%> (+2.81%)` | :arrow_up: |

... and [251 files with indirect coverage changes](https://app.codecov.io/gh/zeromicro/go-zero/pull/4069/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=zeromicro)

</details>
kevwan: Please rebase, included lots of previous commits.

Patch:
diff --git a/core/bloom/bloom.go b/core/bloom/bloom.go
index 3a72cf84..e47eb86f 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 00000000..6d9d60f2
--- /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/limit/periodlimit.go b/core/limit/periodlimit.go
index 332fe3c9..070bb7f8 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 00000000..c351c840
--- /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 e33b4b95..0070eca4 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
-
-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`)
+var (
+	//go:embed tokenscript.lua
+	tokenScript string
+
+	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 00000000..b70ad1ee
--- /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 00000000..e9aa40a2
--- /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 00000000..0fd75838
--- /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 ab26cc98..28b704b8 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
 	}


Test patch:
diff --git a/core/bloom/testscript.lua b/core/bloom/testscript.lua
new file mode 100644
index 00000000..2eda5784
--- /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