github.com/binbinly/pkg@v0.0.11-0.20240321014439-f4fbf666eb0f/lock/redis.go (about)

     1  package lock
     2  
     3  import (
     4  	"context"
     5  	"strings"
     6  	"time"
     7  
     8  	"github.com/pkg/errors"
     9  	"github.com/redis/go-redis/v9"
    10  )
    11  
    12  // RedisLock is a redis lock.
    13  type RedisLock struct {
    14  	prefix string
    15  	key    string
    16  	token  string
    17  	rdb    *redis.Client
    18  	ttl    time.Duration
    19  }
    20  
    21  // Option RedisLock option
    22  type Option func(l *RedisLock)
    23  
    24  // WithPrefix with prefix
    25  func WithPrefix(prefix string) Option {
    26  	return func(l *RedisLock) {
    27  		l.prefix = prefix
    28  	}
    29  }
    30  
    31  // WithTTL with ttl
    32  func WithTTL(ttl time.Duration) Option {
    33  	return func(l *RedisLock) {
    34  		l.ttl = ttl
    35  	}
    36  }
    37  
    38  // NewRedisLock new a redis lock instance
    39  func NewRedisLock(rdb *redis.Client, key string, opts ...Option) *RedisLock {
    40  	opt := &RedisLock{
    41  		rdb:    rdb,
    42  		token:  genToken(),
    43  		prefix: _prefix,
    44  		ttl:    _ttl,
    45  	}
    46  	for _, f := range opts {
    47  		f(opt)
    48  	}
    49  	opt.key = strings.Join([]string{opt.prefix, key}, ":")
    50  	return opt
    51  }
    52  
    53  // Lock acquires the lock.
    54  func (l *RedisLock) Lock(ctx context.Context) (bool, error) {
    55  	isSet, err := l.rdb.SetNX(ctx, l.key, l.token, l.ttl).Result()
    56  	if err == redis.Nil {
    57  		return false, nil
    58  	} else if err != nil {
    59  		return false, errors.Wrapf(err, "[lock] acquires the lock err, key: %s", l.key)
    60  	}
    61  	return isSet, nil
    62  }
    63  
    64  // Unlock del the lock.
    65  // NOTE: token 一致才会执行删除，避免误删，这里用了lua脚本进行事务处理
    66  func (l *RedisLock) Unlock(ctx context.Context) (bool, error) {
    67  	luaScript := "if redis.call('GET',KEYS[1]) == ARGV[1] then return redis.call('DEL',KEYS[1]) else return 0 end"
    68  	ret, err := l.rdb.Eval(ctx, luaScript, []string{l.key}, l.token).Result()
    69  	if err != nil {
    70  		return false, err
    71  	}
    72  	reply, ok := ret.(int64)
    73  	if !ok {
    74  		return false, nil
    75  	}
    76  	return reply == 1, nil
    77  }