Pārlūkot izejas kodu

fix redis try-lock bug (#1366)

#issue_id: 1338

Co-authored-by: zhangwei <>
种豆得豆 3 gadi atpakaļ
vecāks
revīzija
836726e710
2 mainītis faili ar 38 papildinājumiem un 17 dzēšanām
  1. 17 17
      core/stores/redis/redislock.go
  2. 21 0
      core/stores/redis/redislock_test.go

+ 17 - 17
core/stores/redis/redislock.go

@@ -2,7 +2,6 @@ package redis
 
 import (
 	"math/rand"
-	"strconv"
 	"sync/atomic"
 	"time"
 
@@ -12,12 +11,6 @@ import (
 )
 
 const (
-	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
@@ -32,6 +25,7 @@ end`
 type RedisLock struct {
 	store   *Redis
 	seconds uint32
+	count   int32
 	key     string
 	id      string
 }
@@ -51,30 +45,36 @@ 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)
-	resp, err := rl.store.Eval(lockCommand, []string{rl.key}, []string{
-		rl.id, strconv.Itoa(int(seconds)*millisPerSecond + tolerance),
-	})
+	ok, err := rl.store.SetnxEx(rl.key, rl.id, int(seconds))
 	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 resp == nil {
+	} else if !ok {
+		atomic.AddInt32(&rl.count, -1)
 		return false, nil
 	}
 
-	reply, ok := resp.(string)
-	if ok && reply == "OK" {
-		return true, nil
-	}
+	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

+ 21 - 0
core/stores/redis/redislock_test.go

@@ -29,5 +29,26 @@ 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)
+
 	})
 }