From 91b8effb24b04c32e7a116254d777292030e4f58 Mon Sep 17 00:00:00 2001
From: Kevin Wan <wanjunfeng@gmail.com>
Date: Sat, 30 Jul 2022 19:46:10 +0800
Subject: [PATCH] chore: refactor redislock (#2210)

* chore: refactor redislock

* chore: add more tests
---
 core/stores/redis/redislock.go      | 28 ++++++++++++----------------
 core/stores/redis/redislock_test.go | 26 ++++++++++++++++++++------
 2 files changed, 32 insertions(+), 22 deletions(-)

diff --git a/core/stores/redis/redislock.go b/core/stores/redis/redislock.go
index e53c70e2..de66e97e 100644
--- a/core/stores/redis/redislock.go
+++ b/core/stores/redis/redislock.go
@@ -36,7 +36,6 @@ type RedisLock struct {
 	seconds uint32
 	key     string
 	id      string
-	ctx     context.Context
 }
 
 func init() {
@@ -54,9 +53,13 @@ func NewRedisLock(store *Redis, key string) *RedisLock {
 
 // Acquire acquires the lock.
 func (rl *RedisLock) Acquire() (bool, error) {
-	rl.fillCtx()
+	return rl.AcquireCtx(context.Background())
+}
+
+// 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.EvalCtx(rl.ctx, lockCommand, []string{rl.key}, []string{
+	resp, err := rl.store.EvalCtx(ctx, lockCommand, []string{rl.key}, []string{
 		rl.id, strconv.Itoa(int(seconds)*millisPerSecond + tolerance),
 	})
 	if err == red.Nil {
@@ -79,8 +82,12 @@ func (rl *RedisLock) Acquire() (bool, error) {
 
 // Release releases the lock.
 func (rl *RedisLock) Release() (bool, error) {
-	rl.fillCtx()
-	resp, err := rl.store.EvalCtx(rl.ctx, delCommand, []string{rl.key}, []string{rl.id})
+	return rl.ReleaseCtx(context.Background())
+}
+
+// ReleaseCtx releases the lock with the given ctx.
+func (rl *RedisLock) ReleaseCtx(ctx context.Context) (bool, error) {
+	resp, err := rl.store.EvalCtx(ctx, delCommand, []string{rl.key}, []string{rl.id})
 	if err != nil {
 		return false, err
 	}
@@ -97,14 +104,3 @@ func (rl *RedisLock) Release() (bool, error) {
 func (rl *RedisLock) SetExpire(seconds int) {
 	atomic.StoreUint32(&rl.seconds, uint32(seconds))
 }
-
-// WithContext set context.
-func (rl *RedisLock) WithContext(ctx context.Context) {
-	rl.ctx = ctx
-}
-
-func (rl *RedisLock) fillCtx() {
-	if rl.ctx == nil {
-		rl.ctx = context.Background()
-	}
-}
diff --git a/core/stores/redis/redislock_test.go b/core/stores/redis/redislock_test.go
index 410b46a7..4f1d535d 100644
--- a/core/stores/redis/redislock_test.go
+++ b/core/stores/redis/redislock_test.go
@@ -14,18 +14,12 @@ func TestRedisLock(t *testing.T) {
 		return func(client *Redis) {
 			key := stringx.Rand()
 			firstLock := NewRedisLock(client, key)
-			if ctx != nil {
-				firstLock.WithContext(ctx)
-			}
 			firstLock.SetExpire(5)
 			firstAcquire, err := firstLock.Acquire()
 			assert.Nil(t, err)
 			assert.True(t, firstAcquire)
 
 			secondLock := NewRedisLock(client, key)
-			if ctx != nil {
-				secondLock.WithContext(ctx)
-			}
 			secondLock.SetExpire(5)
 			againAcquire, err := secondLock.Acquire()
 			assert.Nil(t, err)
@@ -49,3 +43,23 @@ func TestRedisLock(t *testing.T) {
 		runOnRedis(t, testFn(context.Background()))
 	})
 }
+
+func TestRedisLock_Expired(t *testing.T) {
+	runOnRedis(t, func(client *Redis) {
+		key := stringx.Rand()
+		redisLock := NewRedisLock(client, key)
+		ctx, cancel := context.WithCancel(context.Background())
+		cancel()
+		_, err := redisLock.AcquireCtx(ctx)
+		assert.NotNil(t, err)
+	})
+
+	runOnRedis(t, func(client *Redis) {
+		key := stringx.Rand()
+		redisLock := NewRedisLock(client, key)
+		ctx, cancel := context.WithCancel(context.Background())
+		cancel()
+		_, err := redisLock.ReleaseCtx(ctx)
+		assert.NotNil(t, err)
+	})
+}