github.com/shuguocloud/go-zero@v1.3.0/core/stores/redis/redislock.go (about)

     1  package redis
     2  
     3  import (
     4  	"math/rand"
     5  	"sync/atomic"
     6  	"time"
     7  
     8  	red "github.com/go-redis/redis"
     9  	"github.com/shuguocloud/go-zero/core/logx"
    10  	"github.com/shuguocloud/go-zero/core/stringx"
    11  )
    12  
    13  const (
    14  	delCommand = `if redis.call("GET", KEYS[1]) == ARGV[1] then
    15      return redis.call("DEL", KEYS[1])
    16  else
    17      return 0
    18  end`
    19  	randomLen = 16
    20  )
    21  
    22  // A RedisLock is a redis lock.
    23  type RedisLock struct {
    24  	store   *Redis
    25  	seconds uint32
    26  	count   int32
    27  	key     string
    28  	id      string
    29  }
    30  
    31  func init() {
    32  	rand.Seed(time.Now().UnixNano())
    33  }
    34  
    35  // NewRedisLock returns a RedisLock.
    36  func NewRedisLock(store *Redis, key string) *RedisLock {
    37  	return &RedisLock{
    38  		store: store,
    39  		key:   key,
    40  		id:    stringx.Randn(randomLen),
    41  	}
    42  }
    43  
    44  // Acquire acquires the lock.
    45  func (rl *RedisLock) Acquire() (bool, error) {
    46  	newCount := atomic.AddInt32(&rl.count, 1)
    47  	if newCount > 1 {
    48  		return true, nil
    49  	}
    50  
    51  	seconds := atomic.LoadUint32(&rl.seconds)
    52  	ok, err := rl.store.SetnxEx(rl.key, rl.id, int(seconds+1)) // +1s for tolerance
    53  	if err == red.Nil {
    54  		atomic.AddInt32(&rl.count, -1)
    55  		return false, nil
    56  	} else if err != nil {
    57  		atomic.AddInt32(&rl.count, -1)
    58  		logx.Errorf("Error on acquiring lock for %s, %s", rl.key, err.Error())
    59  		return false, err
    60  	} else if !ok {
    61  		atomic.AddInt32(&rl.count, -1)
    62  		return false, nil
    63  	}
    64  
    65  	return true, nil
    66  }
    67  
    68  // Release releases the lock.
    69  func (rl *RedisLock) Release() (bool, error) {
    70  	newCount := atomic.AddInt32(&rl.count, -1)
    71  	if newCount > 0 {
    72  		return true, nil
    73  	}
    74  
    75  	resp, err := rl.store.Eval(delCommand, []string{rl.key}, []string{rl.id})
    76  	if err != nil {
    77  		return false, err
    78  	}
    79  
    80  	reply, ok := resp.(int64)
    81  	if !ok {
    82  		return false, nil
    83  	}
    84  
    85  	return reply == 1, nil
    86  }
    87  
    88  // SetExpire sets the expiration.
    89  func (rl *RedisLock) SetExpire(seconds int) {
    90  	atomic.StoreUint32(&rl.seconds, uint32(seconds))
    91  }