github.com/lingyao2333/mo-zero@v1.4.1/core/stores/redis/redislock.go (about)

     1  package redis
     2  
     3  import (
     4  	"context"
     5  	"math/rand"
     6  	"strconv"
     7  	"sync/atomic"
     8  	"time"
     9  
    10  	red "github.com/go-redis/redis/v8"
    11  
    12  	"github.com/lingyao2333/mo-zero/core/logx"
    13  	"github.com/lingyao2333/mo-zero/core/stringx"
    14  )
    15  
    16  const (
    17  	randomLen       = 16
    18  	tolerance       = 500 // milliseconds
    19  	millisPerSecond = 1000
    20  	lockCommand     = `if redis.call("GET", KEYS[1]) == ARGV[1] then
    21      redis.call("SET", KEYS[1], ARGV[1], "PX", ARGV[2])
    22      return "OK"
    23  else
    24      return redis.call("SET", KEYS[1], ARGV[1], "NX", "PX", ARGV[2])
    25  end`
    26  	delCommand = `if redis.call("GET", KEYS[1]) == ARGV[1] then
    27      return redis.call("DEL", KEYS[1])
    28  else
    29      return 0
    30  end`
    31  )
    32  
    33  // A RedisLock is a redis lock.
    34  type RedisLock struct {
    35  	store   *Redis
    36  	seconds uint32
    37  	key     string
    38  	id      string
    39  }
    40  
    41  func init() {
    42  	rand.Seed(time.Now().UnixNano())
    43  }
    44  
    45  // NewRedisLock returns a RedisLock.
    46  func NewRedisLock(store *Redis, key string) *RedisLock {
    47  	return &RedisLock{
    48  		store: store,
    49  		key:   key,
    50  		id:    stringx.Randn(randomLen),
    51  	}
    52  }
    53  
    54  // Acquire acquires the lock.
    55  func (rl *RedisLock) Acquire() (bool, error) {
    56  	return rl.AcquireCtx(context.Background())
    57  }
    58  
    59  // AcquireCtx acquires the lock with the given ctx.
    60  func (rl *RedisLock) AcquireCtx(ctx context.Context) (bool, error) {
    61  	seconds := atomic.LoadUint32(&rl.seconds)
    62  	resp, err := rl.store.EvalCtx(ctx, lockCommand, []string{rl.key}, []string{
    63  		rl.id, strconv.Itoa(int(seconds)*millisPerSecond + tolerance),
    64  	})
    65  	if err == red.Nil {
    66  		return false, nil
    67  	} else if err != nil {
    68  		logx.Errorf("Error on acquiring lock for %s, %s", rl.key, err.Error())
    69  		return false, err
    70  	} else if resp == nil {
    71  		return false, nil
    72  	}
    73  
    74  	reply, ok := resp.(string)
    75  	if ok && reply == "OK" {
    76  		return true, nil
    77  	}
    78  
    79  	logx.Errorf("Unknown reply when acquiring lock for %s: %v", rl.key, resp)
    80  	return false, nil
    81  }
    82  
    83  // Release releases the lock.
    84  func (rl *RedisLock) Release() (bool, error) {
    85  	return rl.ReleaseCtx(context.Background())
    86  }
    87  
    88  // ReleaseCtx releases the lock with the given ctx.
    89  func (rl *RedisLock) ReleaseCtx(ctx context.Context) (bool, error) {
    90  	resp, err := rl.store.EvalCtx(ctx, delCommand, []string{rl.key}, []string{rl.id})
    91  	if err != nil {
    92  		return false, err
    93  	}
    94  
    95  	reply, ok := resp.(int64)
    96  	if !ok {
    97  		return false, nil
    98  	}
    99  
   100  	return reply == 1, nil
   101  }
   102  
   103  // SetExpire sets the expiration.
   104  func (rl *RedisLock) SetExpire(seconds int) {
   105  	atomic.StoreUint32(&rl.seconds, uint32(seconds))
   106  }