优化lua脚本管理 (#16)

lock.go

@@ -2,22 +2,23 @@ package go_redislock
 
 import (
 	"context"
+	_ "embed"
 	"fmt"
 	"github.com/go-redis/redis/v8"
 	"sync"
 	"time"
 )
 
+// 默认锁超时时间
+const lockTime = 5 * time.Second
+
 type RedisLockInter interface {
 	// Lock 加锁
 	Lock() error
-
 	// UnLock 解锁
 	UnLock() error
-
 	// SpinLock 自旋锁
 	SpinLock(timeout time.Duration) error
-
 	// Renew 手动续期
 	Renew() error
 }
@@ -34,9 +35,6 @@ type RedisLock struct {
 	mutex           sync.Mutex
 }
 
-// 默认锁超时时间
-const lockTime = 5 * time.Second
-
 type Option func(lock *RedisLock)
 
 func New(ctx context.Context, redisClient *redis.Client, lockKey string, options ...Option) RedisLockInter {

lock_redis.go

@@ -2,12 +2,22 @@ package go_redislock
 
 import (
 	"context"
+	_ "embed"
 	"errors"
 	"fmt"
 	"log"
 	"time"
 )
 
+var (
+	//go:embed lua/lock.lua
+	lockScript string
+	//go:embed lua/unLock.lua
+	unLockScript string
+	//go:embed lua/renew.lua
+	renewScript string
+)
+
 // Lock 加锁
 func (lock *RedisLock) Lock() error {
 	lock.mutex.Lock()

lock_redis_test.go

@@ -445,11 +445,41 @@ func getRedisClient() *redis.Client {
 	if os.Getenv("GITHUB_ACTIONS") == "true" {
 		return nil
 	}
-	return redis.NewClient(&redis.Options{
-		Addr:     "localhost:6379",
+
+	rdb := redis.NewClient(&redis.Options{
+		Addr:     "127.0.0.1:6379",
 		Password: "",
 		DB:       0,
 	})
+
+	// 尝试执行PING命令
+	// 如果执行PING命令出错，则表明连接失败
+	ctx := context.TODO()
+	_, err := rdb.Ping(ctx).Result()
+	if err != nil {
+		log.Fatalf("Failed to ping Redis: %v", err)
+	}
+
+	return rdb
+}
+
+func TestSevLock(t *testing.T) {
+	redisClient := getRedisClient()
+	if redisClient == nil {
+		log.Println("Github actions skip this test")
+		return
+	}
+
+	ctx := context.TODO()
+	key := "test_key_TestSevLock"
+	lock := New(ctx, redisClient, key)
+	defer lock.UnLock()
+
+	err := lock.Lock()
+	if err != nil {
+		t.Errorf("Lock() returned unexpected error: %v", err)
+		return
+	}
 }
 
 // 测试加锁成功

lua/lock.lua

@@ -0,0 +1,20 @@
+local lock_key = KEYS[1]
+local lock_value = ARGV[1]
+local lock_ttl = tonumber(ARGV[2])
+local reentrant_key = lock_key .. ':count:' .. lock_value
+local reentrant_count = tonumber(redis.call('GET', reentrant_key) or '0')
+
+if reentrant_count > 0 then
+    redis.call('INCR', reentrant_key)
+    redis.call('EXPIRE', lock_key, lock_ttl)
+    redis.call('EXPIRE', reentrant_key, lock_ttl)
+    return "OK"
+end
+
+if redis.call('SET', lock_key, lock_value, 'NX', 'EX', lock_ttl) then
+    redis.call('SET', reentrant_key, 1)
+    redis.call('EXPIRE', reentrant_key, lock_ttl)
+    return "OK"
+end
+
+return nil

lua/renew.lua

@@ -0,0 +1,13 @@
+local lock_key = KEYS[1]
+local lock_value = ARGV[1]
+local lock_ttl = tonumber(ARGV[2])
+local reentrant_key = lock_key .. ':count:' .. lock_value
+local reentrant_count = tonumber(redis.call('GET', reentrant_key) or '0')
+
+if reentrant_count > 0 or redis.call('GET', lock_key) == lock_value then
+    redis.call('EXPIRE', lock_key, lock_ttl)
+    redis.call('EXPIRE', reentrant_key, lock_ttl)
+    return "OK"
+end
+
+return nil

lua/unLock.lua

@@ -0,0 +1,20 @@
+local lock_key = KEYS[1]
+local lock_value = ARGV[1]
+local reentrant_key = lock_key .. ':count:' .. lock_value
+local reentrant_count = tonumber(redis.call('GET', reentrant_key) or '0')
+
+if reentrant_count > 1 then
+    redis.call('DECR', reentrant_key)
+    return "OK"
+elseif reentrant_count == 1 then
+    redis.call('DEL', reentrant_key)
+    redis.call('DEL', lock_key)
+    return "OK"
+end
+
+if redis.call('GET', lock_key) == lock_value then
+    redis.call('DEL', lock_key)
+    return "OK"
+else
+    return nil
+end