浏览代码

chore: refactor redislock (#2210)

* chore: refactor redislock

* chore: add more tests
Kevin Wan 2 年之前
父节点
当前提交
91b8effb24
共有 2 个文件被更改,包括 32 次插入22 次删除
  1. 12 16
      core/stores/redis/redislock.go
  2. 20 6
      core/stores/redis/redislock_test.go

+ 12 - 16
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()
-	}
-}

+ 20 - 6
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)
+	})
+}