chore: refactor redislock (#2210)

* chore: refactor redislock

* chore: add more tests
This commit is contained in:
Kevin Wan
2022-07-30 19:46:10 +08:00
committed by GitHub
parent 4879d4dfcd
commit 91b8effb24
2 changed files with 32 additions and 22 deletions

View File

@@ -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()
}
}

View File

@@ -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)
})
}