From 278cd123c85ff902cdb3b920c1b671e90bc42853 Mon Sep 17 00:00:00 2001
From: Kevin Wan <wanjunfeng@gmail.com>
Date: Thu, 24 Mar 2022 16:17:01 +0800
Subject: [PATCH] feat: remove reentrance in redislock, timeout bug (#1704)

---
 core/stores/redis/redislock.go      | 39 ++++++++++++++++-------------
 core/stores/redis/redislock_test.go | 20 ---------------
 2 files changed, 21 insertions(+), 38 deletions(-)

diff --git a/core/stores/redis/redislock.go b/core/stores/redis/redislock.go
index daa9bdec..53c3aea6 100644
--- a/core/stores/redis/redislock.go
+++ b/core/stores/redis/redislock.go
@@ -2,6 +2,7 @@ package redis
 
 import (
 	"math/rand"
+	"strconv"
 	"sync/atomic"
 	"time"
 
@@ -11,19 +12,26 @@ import (
 )
 
 const (
+	randomLen       = 16
+	tolerance       = 500 // milliseconds
+	millisPerSecond = 1000
+	lockCommand     = `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`
 	delCommand = `if redis.call("GET", KEYS[1]) == ARGV[1] then
     return redis.call("DEL", KEYS[1])
 else
     return 0
 end`
-	randomLen = 16
 )
 
 // A RedisLock is a redis lock.
 type RedisLock struct {
 	store   *Redis
 	seconds uint32
-	count   int32
 	key     string
 	id      string
 }
@@ -43,35 +51,30 @@ func NewRedisLock(store *Redis, key string) *RedisLock {
 
 // Acquire acquires the lock.
 func (rl *RedisLock) Acquire() (bool, error) {
-	newCount := atomic.AddInt32(&rl.count, 1)
-	if newCount > 1 {
-		return true, nil
-	}
-
 	seconds := atomic.LoadUint32(&rl.seconds)
-	ok, err := rl.store.SetnxEx(rl.key, rl.id, int(seconds+1)) // +1s for tolerance
+	resp, err := rl.store.Eval(lockCommand, []string{rl.key}, []string{
+		rl.id, strconv.Itoa(int(seconds)*millisPerSecond + tolerance),
+	})
 	if err == red.Nil {
-		atomic.AddInt32(&rl.count, -1)
 		return false, nil
 	} else if err != nil {
-		atomic.AddInt32(&rl.count, -1)
 		logx.Errorf("Error on acquiring lock for %s, %s", rl.key, err.Error())
 		return false, err
-	} else if !ok {
-		atomic.AddInt32(&rl.count, -1)
+	} else if resp == nil {
 		return false, nil
 	}
 
-	return true, nil
+	reply, ok := resp.(string)
+	if ok && reply == "OK" {
+		return true, nil
+	}
+
+	logx.Errorf("Unknown reply when acquiring lock for %s: %v", rl.key, resp)
+	return false, nil
 }
 
 // Release releases the lock.
 func (rl *RedisLock) Release() (bool, error) {
-	newCount := atomic.AddInt32(&rl.count, -1)
-	if newCount > 0 {
-		return true, nil
-	}
-
 	resp, err := rl.store.Eval(delCommand, []string{rl.key}, []string{rl.id})
 	if err != nil {
 		return false, err
diff --git a/core/stores/redis/redislock_test.go b/core/stores/redis/redislock_test.go
index 2e945f8e..55ec9ef7 100644
--- a/core/stores/redis/redislock_test.go
+++ b/core/stores/redis/redislock_test.go
@@ -29,25 +29,5 @@ func TestRedisLock(t *testing.T) {
 		endAcquire, err := secondLock.Acquire()
 		assert.Nil(t, err)
 		assert.True(t, endAcquire)
-
-		endAcquire, err = secondLock.Acquire()
-		assert.Nil(t, err)
-		assert.True(t, endAcquire)
-
-		release, err = secondLock.Release()
-		assert.Nil(t, err)
-		assert.True(t, release)
-
-		againAcquire, err = firstLock.Acquire()
-		assert.Nil(t, err)
-		assert.False(t, againAcquire)
-
-		release, err = secondLock.Release()
-		assert.Nil(t, err)
-		assert.True(t, release)
-
-		firstAcquire, err = firstLock.Acquire()
-		assert.Nil(t, err)
-		assert.True(t, firstAcquire)
 	})
 }