github.com/songzhibin97/gkit@v1.2.13/distributed/locker/lock_ridis/redis.go (about)

     1  package lock_ridis
     2  
     3  import (
     4  	"errors"
     5  	"strconv"
     6  	"time"
     7  
     8  	"github.com/go-redis/redis/v8"
     9  	"github.com/songzhibin97/gkit/distributed/locker"
    10  	"github.com/songzhibin97/gkit/options"
    11  )
    12  
    13  var (
    14  	ErrLockFailed   = errors.New("获取锁失败")
    15  	ErrUnLockFailed = errors.New("释放锁失败")
    16  	CMDLock         = `if redis.call("GET", KEYS[1]) == ARGV[1] then
    17      redis.call("SET", KEYS[1], ARGV[1], "PX", ARGV[2])
    18      return "OK"
    19  else
    20      return redis.call("SET", KEYS[1], ARGV[1], "NX", "PX", ARGV[2])
    21  end`
    22  	CMDUnlock = `if redis.call("GET", KEYS[1]) == ARGV[1] then
    23      return redis.call("DEL", KEYS[1])
    24  else
    25      return 0
    26  end`
    27  )
    28  
    29  type Lock struct {
    30  	config
    31  	client redis.UniversalClient
    32  }
    33  
    34  func (l *Lock) lock(key string, expire int, mark string) error {
    35  	ctx := l.client.Context()
    36  	resp, err := l.client.Eval(ctx, CMDLock, []string{key}, []string{mark, strconv.Itoa(expire)}).Result()
    37  	if err != nil && !errors.Is(err, redis.Nil) {
    38  		return err
    39  	}
    40  	if errors.Is(err, redis.Nil) || resp == nil {
    41  		return ErrLockFailed
    42  	}
    43  	reply, ok := resp.(string)
    44  	if !ok || reply != "OK" {
    45  		return ErrUnLockFailed
    46  	}
    47  	return nil
    48  }
    49  
    50  func (l *Lock) Lock(key string, expire int, mark string) error {
    51  	var err error
    52  	for i := 0; i < l.retries+1; i++ {
    53  		err = l.lock(key, expire, mark)
    54  		if err == nil {
    55  			break
    56  		}
    57  		if l.interval > 0 {
    58  			time.Sleep(l.interval)
    59  		}
    60  	}
    61  	return err
    62  }
    63  
    64  func (l *Lock) UnLock(key string, mark string) error {
    65  	ctx := l.client.Context()
    66  	resp, err := l.client.Eval(ctx, CMDUnlock, []string{key}, []string{mark}).Result()
    67  	if err != nil && !errors.Is(err, redis.Nil) {
    68  		return err
    69  	}
    70  	if errors.Is(err, redis.Nil) || resp == nil {
    71  		return ErrUnLockFailed
    72  	}
    73  	reply, ok := resp.(int64)
    74  	if !ok || reply != 1 {
    75  		return ErrUnLockFailed
    76  	}
    77  	return nil
    78  }
    79  
    80  func NewRedisLock(client redis.UniversalClient, opts ...options.Option) locker.Locker {
    81  	o := &config{
    82  		interval: 0,
    83  		retries:  0,
    84  	}
    85  	for _, opt := range opts {
    86  		opt(o)
    87  	}
    88  	if o.interval <= 0 || o.retries <= 0 {
    89  		o.interval = 0
    90  		o.retries = 0
    91  	}
    92  	return &Lock{
    93  		client: client,
    94  		config: *o,
    95  	}
    96  }