Repository: zeromicro/go-zero
Pull request #2208 (2022-07-29 08:03:32)

feat(redislock): support set context


Patch:
diff --git a/core/stores/redis/redislock.go b/core/stores/redis/redislock.go
index 53c3aea6..e53c70e2 100644
--- a/core/stores/redis/redislock.go
+++ b/core/stores/redis/redislock.go
@@ -1,12 +1,14 @@
 package redis
 
 import (
+	"context"
 	"math/rand"
 	"strconv"
 	"sync/atomic"
 	"time"
 
 	red "github.com/go-redis/redis/v8"
+
 	"github.com/zeromicro/go-zero/core/logx"
 	"github.com/zeromicro/go-zero/core/stringx"
 )
@@ -34,6 +36,7 @@ type RedisLock struct {
 	seconds uint32
 	key     string
 	id      string
+	ctx     context.Context
 }
 
 func init() {
@@ -51,8 +54,9 @@ func NewRedisLock(store *Redis, key string) *RedisLock {
 
 // Acquire acquires the lock.
 func (rl *RedisLock) Acquire() (bool, error) {
+	rl.fillCtx()
 	seconds := atomic.LoadUint32(&rl.seconds)
-	resp, err := rl.store.Eval(lockCommand, []string{rl.key}, []string{
+	resp, err := rl.store.EvalCtx(rl.ctx, lockCommand, []string{rl.key}, []string{
 		rl.id, strconv.Itoa(int(seconds)*millisPerSecond + tolerance),
 	})
 	if err == red.Nil {
@@ -75,7 +79,8 @@ func (rl *RedisLock) Acquire() (bool, error) {
 
 // Release releases the lock.
 func (rl *RedisLock) Release() (bool, error) {
-	resp, err := rl.store.Eval(delCommand, []string{rl.key}, []string{rl.id})
+	rl.fillCtx()
+	resp, err := rl.store.EvalCtx(rl.ctx, delCommand, []string{rl.key}, []string{rl.id})
 	if err != nil {
 		return false, err
 	}
@@ -92,3 +97,14 @@ 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/go.mod b/go.mod
index c76614b0..92425a9f 100644
--- a/go.mod
+++ b/go.mod
@@ -36,7 +36,7 @@ require (
 	golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a
 	golang.org/x/time v0.0.0-20220411224347-583f2d630306
 	google.golang.org/grpc v1.48.0
-	google.golang.org/protobuf v1.28.0
+	google.golang.org/protobuf v1.28.1
 	gopkg.in/cheggaaa/pb.v1 v1.0.28
 	gopkg.in/h2non/gock.v1 v1.1.2
 	gopkg.in/yaml.v2 v2.4.0
diff --git a/go.sum b/go.sum
index 6f39ce1b..71b36eb7 100644
--- a/go.sum
+++ b/go.sum
@@ -830,8 +830,9 @@ google.golang.org/protobuf v1.25.1-0.20200805231151-a709e31e5d12/go.mod h1:9JNX7
 google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw=
 google.golang.org/protobuf v1.26.0/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQnmE0givc=
 google.golang.org/protobuf v1.27.1/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQnmE0givc=
-google.golang.org/protobuf v1.28.0 h1:w43yiav+6bVFTBQFZX0r7ipe9JQ1QsbMgHwbBziscLw=
 google.golang.org/protobuf v1.28.0/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I=
+google.golang.org/protobuf v1.28.1 h1:d0NfwRgPtno5B1Wa6L2DAG+KivqkdutMf1UhdNx175w=
+google.golang.org/protobuf v1.28.1/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I=
 gopkg.in/alecthomas/kingpin.v2 v2.2.6/go.mod h1:FMv+mEhP44yOT+4EoQTLFTRgOQ1FBLkstjWtayDeSgw=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=


Test patch:
diff --git a/core/stores/redis/redislock_test.go b/core/stores/redis/redislock_test.go
index 55ec9ef7..410b46a7 100644
--- a/core/stores/redis/redislock_test.go
+++ b/core/stores/redis/redislock_test.go
@@ -1,33 +1,51 @@
 package redis
 
 import (
+	"context"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+
 	"github.com/zeromicro/go-zero/core/stringx"
 )
 
 func TestRedisLock(t *testing.T) {
-	runOnRedis(t, func(client *Redis) {
-		key := stringx.Rand()
-		firstLock := NewRedisLock(client, key)
-		firstLock.SetExpire(5)
-		firstAcquire, err := firstLock.Acquire()
-		assert.Nil(t, err)
-		assert.True(t, firstAcquire)
-
-		secondLock := NewRedisLock(client, key)
-		secondLock.SetExpire(5)
-		againAcquire, err := secondLock.Acquire()
-		assert.Nil(t, err)
-		assert.False(t, againAcquire)
-
-		release, err := firstLock.Release()
-		assert.Nil(t, err)
-		assert.True(t, release)
-
-		endAcquire, err := secondLock.Acquire()
-		assert.Nil(t, err)
-		assert.True(t, endAcquire)
+	testFn := func(ctx context.Context) func(client *Redis) {
+		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)
+			assert.False(t, againAcquire)
+
+			release, err := firstLock.Release()
+			assert.Nil(t, err)
+			assert.True(t, release)
+
+			endAcquire, err := secondLock.Acquire()
+			assert.Nil(t, err)
+			assert.True(t, endAcquire)
+		}
+	}
+
+	t.Run("normal", func(t *testing.T) {
+		runOnRedis(t, testFn(nil))
+	})
+
+	t.Run("withContext", func(t *testing.T) {
+		runOnRedis(t, testFn(context.Background()))
 	})
 }