github.com/wfusion/gofusion@v1.1.14/lock/redis.go (about)

     1  package lock
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"strconv"
     7  	"time"
     8  
     9  	"github.com/pkg/errors"
    10  
    11  	"github.com/wfusion/gofusion/common/utils"
    12  	"github.com/wfusion/gofusion/config"
    13  	"github.com/wfusion/gofusion/redis"
    14  
    15  	rdsDrv "github.com/redis/go-redis/v9"
    16  )
    17  
    18  const (
    19  	redisLuaLockCommand = `
    20  if redis.call('EXISTS', KEYS[1]) == 0 or redis.call("HGET", KEYS[1], "holder") == ARGV[1] then
    21      local expired = redis.call("PTTL", KEYS[1])
    22  	if expired == -1 or expired == -2 then
    23  		redis.call('HMSET', KEYS[1], "count", 1, "holder", ARGV[1])
    24          redis.call("PEXPIRE", KEYS[1], ARGV[2])
    25  		return 1
    26      else
    27      	redis.call('HINCRBY', KEYS[1], "count", 1)
    28  		redis.call("PEXPIRE", KEYS[1], ARGV[2] + expired)
    29  		return 1
    30  	end
    31  else
    32  	return nil
    33  end`
    34  
    35  	redisLuaUnlockCommand = `
    36  if redis.call("HGET", KEYS[1], "holder") == ARGV[1] then
    37      if redis.call("HINCRBY", KEYS[1], "count", -1) <= 0 then
    38          redis.call("DEL", KEYS[1])
    39  		return 1
    40  	else
    41  		return 1
    42      end
    43  else
    44      return nil
    45  end`
    46  )
    47  
    48  type redisLuaLocker struct {
    49  	ctx       context.Context
    50  	appName   string
    51  	redisName string
    52  }
    53  
    54  func newRedisLuaLocker(ctx context.Context, appName, redisName string) ReentrantLockable {
    55  	return &redisLuaLocker{ctx: ctx, appName: appName, redisName: redisName}
    56  }
    57  
    58  func (r *redisLuaLocker) Lock(ctx context.Context, key string, opts ...utils.OptionExtender) (err error) {
    59  	opt := utils.ApplyOptions[lockOption](opts...)
    60  	expired := tolerance
    61  	if opt.expired > 0 {
    62  		expired = opt.expired
    63  	}
    64  	lockKey := r.formatLockKey(key)
    65  	err = redis.
    66  		Use(ctx, r.redisName, redis.AppName(r.appName)).
    67  		Eval(ctx, redisLuaLockCommand, []string{lockKey}, []string{
    68  			opt.reentrantKey, strconv.Itoa(int(expired / time.Millisecond)),
    69  		}).
    70  		Err()
    71  	if errors.Is(err, rdsDrv.Nil) {
    72  		return ErrTimeout
    73  	}
    74  	return
    75  }
    76  
    77  func (r *redisLuaLocker) Unlock(ctx context.Context, key string, opts ...utils.OptionExtender) (err error) {
    78  	opt := utils.ApplyOptions[lockOption](opts...)
    79  	lockKey := r.formatLockKey(key)
    80  	return redis.
    81  		Use(ctx, r.redisName, redis.AppName(r.appName)).
    82  		Eval(ctx, redisLuaUnlockCommand, []string{lockKey}, []string{opt.reentrantKey}).
    83  		Err()
    84  }
    85  
    86  func (r *redisLuaLocker) ReentrantLock(ctx context.Context, key, reentrantKey string,
    87  	opts ...utils.OptionExtender) (err error) {
    88  	opt := utils.ApplyOptions[lockOption](opts...)
    89  	if utils.IsStrBlank(opt.reentrantKey) {
    90  		return ErrReentrantKeyNotFound
    91  	}
    92  	return r.Lock(ctx, key, append(opts, ReentrantKey(reentrantKey))...)
    93  }
    94  
    95  func (r *redisLuaLocker) formatLockKey(key string) (format string) {
    96  	return fmt.Sprintf("%s:%s", config.Use(r.appName).AppName(), key)
    97  }
    98  
    99  type redisNXLocker struct {
   100  	ctx       context.Context
   101  	appName   string
   102  	redisName string
   103  }
   104  
   105  func newRedisNXLocker(ctx context.Context, appName, redisName string) Lockable {
   106  	return &redisNXLocker{ctx: ctx, appName: appName, redisName: redisName}
   107  }
   108  
   109  func (r *redisNXLocker) Lock(ctx context.Context, key string, opts ...utils.OptionExtender) (err error) {
   110  	opt := utils.ApplyOptions[lockOption](opts...)
   111  	expired := tolerance
   112  	if opt.expired > 0 {
   113  		expired = opt.expired
   114  	}
   115  	lockKey := r.formatLockKey(key)
   116  	cmd := redis.Use(ctx, r.redisName, redis.AppName(r.appName)).SetNX(ctx, lockKey, utils.UUID(), expired)
   117  	if err = cmd.Err(); err != nil {
   118  		return
   119  	}
   120  	if !cmd.Val() {
   121  		err = ErrTimeout
   122  		return
   123  	}
   124  	return
   125  }
   126  
   127  func (r *redisNXLocker) Unlock(ctx context.Context, key string, _ ...utils.OptionExtender) (err error) {
   128  	lockKey := r.formatLockKey(key)
   129  	return redis.Use(ctx, r.redisName, redis.AppName(r.appName)).Del(ctx, lockKey).Err()
   130  }
   131  
   132  func (r *redisNXLocker) formatLockKey(key string) (format string) {
   133  	return fmt.Sprintf("%s:%s", config.Use(r.appName).AppName(), key)
   134  }