Go语言实现分布式锁

package redis_lock
import (
	"context"
	_ "embed"
	"errors"
	"github.com/go-redis/redis/v9"
	"github.com/google/uuid"
	"golang.org/x/sync/singleflight"
	"time"
)
// go:embed 可以直接解析出文件中的字符串
var (
	//go:embed lua_unlock.lua
	luaUnlock string
	//go:embed refresh.lua
	luaRefresh string
	//go:embed lock.lua
	luaLock string
	//定义好两个异常信息
	ErrLockNotHold         = errors.New("未持有锁")
	ErrFailedToPreemptLock = errors.New("加锁失败")
)
type Client struct {
	//采用公共的接口，后续实例通过传入的方式
	client redis.Cmdable
	// singleflight 用于在一个实例的多个携程中只需要竞争出一个携程
	s singleflight.Group
}
func NewClient(c redis.Cmdable) *Client {
	return &Client{
		client: c,
	}
}
func (c *Client) SingleflightLock(ctx context.Context,
	key string,
	expire time.Duration,
	retry RetryStrategy,
	timeout time.Duration) (*Lock, error) {
	for {
		flag := false
		resCh := c.s.DoChan(key, func() (interface{}, error) {
			flag = true
			return c.Lock(ctx, key, expire, retry, timeout)
		})
		select {
		case res := <-resCh:
			if flag {
				if res.Err != nil {
					return nil, res.Err
				}
				//返回锁对象
				return res.Val.(*Lock), nil
			}
		case <-ctx.Done():
			return nil, ctx.Err()
		}
	}
}
//Lock 加锁方法，根据重试机制进行重新获取
func (c *Client) Lock(ctx context.Context,
	key string,
	expire time.Duration,
	retry RetryStrategy,
	timeout time.Duration) (*Lock, error) {
	var timer *time.Timer
	defer func() {
		if timer != nil {
			timer.Stop()
		}
	}()
	for {
		//设置超时
		lct, cancel := context.WithTimeout(ctx, timeout)
		//获取到uuid
		value := uuid.New().String()
		//执行lua脚本进行加锁
		result, err := c.client.Eval(lct, luaLock, []string{key}, value, expire).Bool()
		//用于主动释放资源
		cancel()
		if err != nil && !errors.Is(err, context.DeadlineExceeded) {
			return nil, err
		}
		if result {
			return newLock(c.client, key, value), nil
		}
		//可以不传重试机制
		if retry != nil {
			//通过重试机制获取重试的策略
			interval, ok := retry.Next()
			if !ok {
				//不用重试
				return nil, ErrFailedToPreemptLock
			}
			if timer == nil {
				timer = time.NewTimer(interval)
			}
			timer.Reset(interval)
			select {
			case <-timer.C: //睡眠时间超时了
				return nil, ctx.Err()
			case <-ctx.Done(): //整个调用的超时
				return nil, ctx.Err()
			}
		}
	}
}
// TryLock 尝试加锁
func (c *Client) TryLock(ctx context.Context, key string, expire time.Duration) (*Lock, error) {
	return c.Lock(ctx, key, expire, nil, 0)
}
// NewLock 创建一个锁结构体
func newLock(client redis.Cmdable, key string, value string) *Lock {
	return &Lock{
		client:     client,
		key:        key,
		value:      value,
		unLockChan: make(chan struct{}, 1), //设置1个缓存数据，用于解锁的信号量
	}
}
// Lock 结构体对象
type Lock struct {
	client redis.Cmdable
	key    string
	value  string
	expire time.Duration
	//在解锁成功之后发送信号来取消续约
	unLockChan chan struct{}
}
// AutoRefresh 自动续约
func (l *Lock) AutoRefresh(interval time.Duration, timeout time.Duration) error {
	//设计一个管道，如果失败了，就发送数据到管道之中，通知进行重试
	retry := make(chan struct{}, 1)
	//方法返回时关闭close
	defer close(retry)
	ticker := time.NewTicker(interval)
	for {
		select {
		//接收到结束的信号时，直接return
		case <-l.unLockChan:
			return nil
		//监听重试的管道
		case <-retry:
			ctx, cancel := context.WithTimeout(context.Background(), timeout)
			err := l.Refresh(ctx)
			//主动调用释放资源
			cancel()
			if err == context.DeadlineExceeded {
				// 执行重试往管道中发送一个信号
				retry <- struct{}{}
				continue
			}
			if err != nil {
				return err
			}
		case <-ticker.C:
			ctx, cancel := context.WithTimeout(context.Background(), timeout)
			err := l.Refresh(ctx)
			//主动调用释放资源
			cancel()
			if err == context.DeadlineExceeded {
				// 执行重试往管道中发送一个信号
				retry <- struct{}{}
				continue
			}
			if err != nil {
				return err
			}
		}
	}
}
// Refresh 续约
func (l *Lock) Refresh(ctx context.Context) error {
	//执行lua脚本，对锁进行续约
	i, err := l.client.Eval(ctx, luaRefresh, []string{l.key}, l.value, l.expire.Milliseconds()).Int64()
	if err == redis.Nil {
		return ErrLockNotHold
	}
	if err != nil {
		return err
	}
	if i == 0 {
		return ErrLockNotHold
	}
	return nil
}
// Unlock 解锁
func (l *Lock) Unlock(ctx context.Context) error {
	//解锁时，退出方法需要发送一个信号让自动续约的goroutine停止
	defer func() {
		l.unLockChan <- struct{}{}
		close(l.unLockChan)
	}()
	//判断返回的结果
	result, err := l.client.Eval(ctx, luaUnlock, []string{l.key}, l.value).Int64()
	if err == redis.Nil {
		return ErrLockNotHold
	}
	if err != nil {
		return err
	}
	//lua脚本返回的结果如果为0，也是代表当前锁不是自己的
	if result == 0 {
		return ErrLockNotHold
	}
	return nil
}

package redis_lock
import "time"
// RetryStrategy 重试策略
type RetryStrategy interface {
	// Next 下一次重试的时间是多久，返回两个参数 time 时间，bool 是否直接重试
	Next() (time.Duration, bool)
}

--[[ 获取到对应的value是否跟当前的一样 ]]
if redis.call("get", KEYS[1]) == ARGV[1]
then
-- 如果一样直接对其时间进行续约
	return redis.call("pexpire", KEYS[1], ARGV[2])
else
-- 如果不一样调用setnx命令对其进行设置值
	return redis.call("set", KEYS[1], ARGV[1], "NX", "PX", ARGV[2])

if redis.call("get", KEYS[1]) == ARGV[1] then
	-- 返回0，代表key不在
	return redis.call("del", KEYS[1])
else
	-- key在，但是值不对
	return 0
end

if redis.call("get", KEYS[1]) == ARGV[1] then
	-- 返回0，代表key不在
	return redis.call("pexpire", KEYS[1], ARGV[2])
else
	-- key在，但是值不对
	return 0
end

func TestClient_TryLock(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()
	testCase := []struct {
		//测试的场景
		name string
		//输入
		key        string
		expiration time.Duration
		//返回一个mock数据
		mock func() redis.Cmdable
		//期望的返回的错误值
		wantError error
		//期望返回的锁
		wantLock *Lock
	}{
		{
			name:       "locked",
			key:        "locked-key",
			expiration: time.Minute,
			mock: func() redis.Cmdable {
				rdb := mocks.NewMockCmdable(ctrl)
				res := redis.NewBoolResult(true, nil)
				i := []interface{}{gomock.Any(), time.Minute}
				rdb.EXPECT().Eval(gomock.Any(), luaLock, []string{"locked-key"}, i...).Return(res)
				return rdb
			},
			wantLock: &Lock{
				key: "locked-key",
			},
		},
	}
	for _, tc := range testCase {
		t.Run(tc.name, func(t *testing.T) {
			var c = NewClient(tc.mock())
			l, err := c.TryLock(context.Background(), tc.key, tc.expiration)
			assert.Equal(t, tc.wantError, err)
			if err != nil {
				return
			}
			//判断返回的key是否跟期望的一样
			assert.Equal(t, tc.key, l.key)
			assert.Equal(t, tc.wantLock.key, l.key)
			assert.NotEmpty(t, l.value)
		})
	}
}