Merge branch 'master' of github.com:go-eagle/eagle

SonderEASE · Oct 15, 2022 · 4d5cfe2 · 4d5cfe2
2 parents a87573f + 4f7a342
commit 4d5cfe2
Show file tree

Hide file tree

Showing 4 changed files with 34 additions and 18 deletions.
diff --git a/pkg/lock/etcd.go b/pkg/lock/etcd.go
@@ -11,34 +11,47 @@ import (
 
 // EtcdLock define a etcd lock
 type EtcdLock struct {
-	sess *concurrency.Session
-	mu   *concurrency.Mutex
+	sess       *concurrency.Session
+	mu         *concurrency.Mutex
+	cancelFunc context.CancelFunc
 }
 
 // NewEtcdLock create a etcd lock
 // ttl for lease
 func NewEtcdLock(client *v3.Client, key string, ttl int) (mutex *EtcdLock, err error) {
 	mutex = &EtcdLock{}
+
+	// get lock timeout
+	// set lease ttl == request timeout
+	expiration := time.Duration(ttl) * time.Second
+	ctx, cancel := context.WithTimeout(context.Background(), expiration)
+	mutex.cancelFunc = cancel
+
 	// default session ttl = 60s
-	mutex.sess, err = concurrency.NewSession(client, concurrency.WithTTL(ttl))
+	mutex.sess, err = concurrency.NewSession(
+		client,
+		concurrency.WithTTL(ttl),
+		concurrency.WithContext(ctx),
+	)
 	if err != nil {
 		return
 	}
+
 	mutex.mu = concurrency.NewMutex(mutex.sess, getEtcdKey(key))
+
 	return
 }
 
 // Lock acquires the lock.
-func (l *EtcdLock) Lock(ctx context.Context, timeout time.Duration) (b bool, err error) {
-	// get lock timeout
-	ctx, cancel := context.WithTimeout(ctx, timeout)
-	defer cancel()
+func (l *EtcdLock) Lock(ctx context.Context) (b bool, err error) {
 	// NOTE: ignore bool value
 	return true, l.mu.Lock(ctx)
 }
 
 // Unlock release a lock.
 func (l *EtcdLock) Unlock(ctx context.Context) (b bool, err error) {
+	defer l.cancelFunc()
+
 	err = l.mu.Unlock(ctx)
 	if err != nil {
 		return

diff --git a/pkg/lock/lock.go b/pkg/lock/lock.go
@@ -2,7 +2,6 @@ package lock
 
 import (
 	"context"
-	"time"
 )
 
 const (
@@ -14,6 +13,6 @@ const (
 
 // Lock define common func
 type Lock interface {
-	Lock(ctx context.Context, timeout time.Duration) (bool, error)
+	Lock(ctx context.Context) (bool, error)
 	Unlock(ctx context.Context) (bool, error)
 }
diff --git a/pkg/lock/lock_test.go b/pkg/lock/lock_test.go
@@ -12,9 +12,10 @@ import (
 
 func TestLockWithDefaultTimeout(t *testing.T) {
 	redis.InitTestRedis()
+	expiration := 2 * time.Second
 
-	lock := NewRedisLock(redis.RedisClient, "lock1")
-	ok, err := lock.Lock(context.Background(), 2*time.Second)
+	lock := NewRedisLock(redis.RedisClient, "lock1", expiration)
+	ok, err := lock.Lock(context.Background())
 	if err != nil {
 		t.Error(err)
 	}
@@ -35,11 +36,12 @@ func TestLockWithDefaultTimeout(t *testing.T) {
 
 func TestLockWithTimeout(t *testing.T) {
 	redis.InitTestRedis()
+	expiration := 2 * time.Second
 
 	t.Run("should lock/unlock success", func(t *testing.T) {
 		ctx := context.Background()
-		lock1 := NewRedisLock(redis.RedisClient, "lock2")
-		ok, err := lock1.Lock(ctx, 2*time.Second)
+		lock1 := NewRedisLock(redis.RedisClient, "lock2", expiration)
+		ok, err := lock1.Lock(ctx)
 		assert.Nil(t, err)
 		assert.True(t, ok)
 
@@ -50,8 +52,8 @@ func TestLockWithTimeout(t *testing.T) {
 
 	t.Run("should unlock failed", func(t *testing.T) {
 		ctx := context.Background()
-		lock2 := NewRedisLock(redis.RedisClient, "lock3")
-		ok, err := lock2.Lock(ctx, 2*time.Second)
+		lock2 := NewRedisLock(redis.RedisClient, "lock3", expiration)
+		ok, err := lock2.Lock(ctx)
 		assert.Nil(t, err)
 		assert.True(t, ok)
 

diff --git a/pkg/lock/redis.go b/pkg/lock/redis.go
@@ -15,22 +15,24 @@ type RedisLock struct {
 	key         string
 	redisClient *redis.Client
 	token       string
+	expiration  time.Duration
 }
 
 // NewRedisLock new a redis lock instance
 // nolint
-func NewRedisLock(rdb *redis.Client, key string) *RedisLock {
+func NewRedisLock(rdb *redis.Client, key string, expiration time.Duration) *RedisLock {
 	opt := &RedisLock{
 		key:         getRedisKey(key),
 		redisClient: rdb,
 		token:       genToken(),
+		expiration:  expiration,
 	}
 	return opt
 }
 
 // Lock acquires the lock.
-func (l *RedisLock) Lock(ctx context.Context, timeout time.Duration) (bool, error) {
-	isSet, err := l.redisClient.SetNX(ctx, l.key, l.token, timeout).Result()
+func (l *RedisLock) Lock(ctx context.Context) (bool, error) {
+	isSet, err := l.redisClient.SetNX(ctx, l.key, l.token, l.expiration).Result()
 	if err == redis.Nil {
 		return false, nil
 	} else if err != nil {