diff --git a/core/stores/redis/redislock.go b/core/stores/redis/redislock.go index 5ed5f24e..3e85d972 100644 --- a/core/stores/redis/redislock.go +++ b/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 diff --git a/core/stores/redis/redislock_test.go b/core/stores/redis/redislock_test.go index 588cf25a..d12ed03e 100644 --- a/core/stores/redis/redislock_test.go +++ b/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) + }) }