go-zero源码阅读-RedisLock分布式锁实现

一. 简介

go-zero分布式锁代码文件: core/stores/redislock.go。 该分布式锁通过lua脚本使用redis的set命令实现， 整体实现逻辑比较简单

二. 代码实现

一. 核心struct

// A RedisLock is a redis lock.
type RedisLock struct {
	store   *Redis // 依赖的redis
	seconds uint32 // key的过期时间
	key     string  // key
	id      string  // value
}

结构体字段围绕redis的set命令设计

二. 结构体初始化

func NewRedisLock(store *Redis, key string) *RedisLock {
	return &RedisLock{
		store: store,
		key:   key,
		id:    stringx.Randn(randomLen), // 生成16位字符串，作为set的value值，用来判断操作的用户是否为加锁的用户
	}
}

初始化时未指定过期时间，锁的过期时间默认500毫秒，同时提供SetExpire方法，来设置过期时间

三. 加锁

func (rl *RedisLock) Acquire() (bool, error) {
	return rl.AcquireCtx(context.Background())
}

// AcquireCtx acquires the lock with the given ctx.
func (rl *RedisLock) AcquireCtx(ctx context.Context) (bool, error) {
    // 获取过期时间
	seconds := atomic.LoadUint32(&rl.seconds)
    // 执行lua脚本进行加锁
    // 过期时间为设置的时间 加上 默认500ms, 秒为正整数同时加上默认的500， 防止过期时间为0，产生死锁
	resp, err := rl.store.EvalCtx(ctx, lockCommand, []string{rl.key}, []string{
		rl.id, strconv.Itoa(int(seconds)*millisPerSecond + tolerance),
	})

    // 判断加锁是否成功
	if err == red.Nil {
		return false, nil
	} else if err != nil {
		logx.Errorf("Error on acquiring lock for %s, %s", rl.key, err.Error())
		return false, err
	} else if resp == nil {
		return false, nil
	}

	reply, ok := resp.(string)
	if ok && reply == "OK" {
		return true, nil
	}

	logx.Errorf("Unknown reply when acquiring lock for %s: %v", rl.key, resp)
	return false, nil
}

四. 加锁lua脚本

lockCommand     = `if redis.call("GET", KEYS[1]) == ARGV[1] then
    redis.call("SET", KEYS[1], ARGV[1], "PX", ARGV[2])
    return "OK"
else
    return redis.call("SET", KEYS[1], ARGV[1], "NX", "PX", ARGV[2])
end`

首先获取锁的value, 判断和传入的随机数是否相同：

相同，则代表当前操作锁的是加锁的用户，直接设置过期时间

不相同，则代表当前不是加锁用户，使用set命令 nx(值不存在，才可以设置成功，排他性)，ex(过期时间), 进行加锁

五. 释放锁

func (rl *RedisLock) Release() (bool, error) {
	return rl.ReleaseCtx(context.Background())
}

// ReleaseCtx releases the lock with the given ctx.
func (rl *RedisLock) ReleaseCtx(ctx context.Context) (bool, error) {
    // 执行lua脚本，同加锁类似
	resp, err := rl.store.EvalCtx(ctx, delCommand, []string{rl.key}, []string{rl.id})
	if err != nil {
		return false, err
	}

	reply, ok := resp.(int64)
	if !ok {
		return false, nil
	}

	return reply == 1, nil
}

六. 释放锁lua脚本

delCommand = `if redis.call("GET", KEYS[1]) == ARGV[1] then
    return redis.call("DEL", KEYS[1])
else
    return 0
end`

同加锁类似，通过value判断，加锁的用户才可以删除锁