【Go】使用lua脚本实现redis分布式锁

分布式锁的要求

实现分布式锁之前要明确一下分布式锁的要求

• 互斥性，在任意时刻，只能有一个进程持有锁
• 防死锁，不能因为持有锁的客户端宕机而使其他进程无法获取到锁。
• 加锁和解锁的必须是同一个进程。
• 保证锁的续租。

redis分布式锁的优缺点

redis实现的分布式锁性能会比zookeeper、etcd等实现的要好，但因为单点故障或者主备异步复制的问题，可能会出现当master宕机crash会导致多个client同时持有分布式锁。（这里的方案会出现这些问题）
因为etcd的高可靠、强一致存储，可以避免故障时出现的问题，保证锁的安全性。可以参考我的这篇 文章 里介绍的etcd分布式锁的实现。

本文的Redis分布式锁的实现步骤

这里参考了JAVA Redisson分布式锁的实现，对redis的操作使用lua脚本，具体步骤如下：

• 首次加锁：如果客户端线程第⼀次去加锁的话，会在key对应的hash数据结构中添加线程标识uuid 1，指定该线程当前对这个key加锁⼀次了。
• 再次加锁：根据uuid判断如果是该线程的再次加锁，则key中的uuid值加1，表示持有锁的次数加1。
• 锁的维持：第一次加锁成功后，会在后台启动一个goroutine定时去续约。
• 加锁失败的阻塞：根据失败时返回的锁的过期时间阻塞，时间到了就去抢锁，可以设置等待超时时间。
• 锁的释放：判断如果key中对应的uuid不存在了，则直接返回。否则对uuid值进行减1操作，判断如果uuid <= 0 则直接将key删除。并取消续租的goroutine。

具体代码实现

package main
import (
	"context"
	"fmt"
	"github.com/go-redis/redis"
	"github.com/google/uuid"
	"log"
	"strconv"
	"testing"
	"time"
)

const (
    //加锁的脚本
	LockScript =  `if (redis.call('exists', KEYS[1]) == 0) then
redis.call('hincrby', KEYS[1], ARGV[2], 1)
redis.call('pexpire', KEYS[1], ARGV[1])
return "OK"
end

if (redis.call('hexists', KEYS[1], ARGV[2]) == 1) then
redis.call('hincrby', KEYS[1], ARGV[2], 1)
redis.call('pexpire', KEYS[1], ARGV[1])
return "OK"
end
return redis.call('pttl', KEYS[1])
`
    //维持锁的脚本
	KeepAliveScript = `if (redis.call('hexists', KEYS[1], ARGV[2]) == 1) then
redis.call('pexpire', KEYS[1], ARGV[1])
return "OK"
end
return "FIAL"
`

    //解锁的脚本
	UnlockScript = `if (redis.call('hexists', KEYS[1], ARGV[1]) == 0) then
return "OK"
end
local counter = redis.call('hincrby', KEYS[1], ARGV[1], -1)
if (counter <= 0) then
redis.call('del', KEYS[1])
return "OK"
end
return "reduce"
`
)

type RedisMutex struct {
	//redis 客户端实例
	Redis *redis.Client
	//锁的存活时间
	TTL time.Duration
	//间隔多久去给锁续期(要比TTL短，不然就没意义了)
	PerKeepAlive time.Duration
	//锁超时时间，避免持有锁的goroutine长时间阻塞（需要做取舍）
	TimeOut time.Duration
	//最大等待锁多久
	MaxWait time.Duration
	//要加锁的key
	Key string
	//作为goroutine的标识
	Id string
	//控制续期任务的关闭
	ctx context.Context
	cancel context.CancelFunc
}

//加锁 如果成功直接返回，如果失败则循环阻塞重试
func (r *RedisMutex)Lock () (err error) {
	res := r.Redis.Eval(LockScript,[]string{r.Key},[]string{strconv.Itoa(int(r.TTL.Milliseconds())),r.Id })
	if res.Err() != nil {
		return res.Err()
	}
	val ,ok := res.Val().(string)
	if ok && val == "OK" {
		if r.cancel == nil {
			ctx, cancel := context.WithTimeout(r.ctx,r.TimeOut)
			r.cancel = cancel
			go KeepAliveWorker(ctx,r.Redis,r.PerKeepAlive,r.TTL,r.Key,r.Id)
		}
		return
	}
	//没有获取到锁则返回结果是锁的过期时间，重复循环获取
	ttl := time.Duration(res.Val().(int64))
	ctx, _ := context.WithTimeout(r.ctx,r.MaxWait)
	for  {
		waiterAfter := time.After(time.Duration(ttl) * time.Millisecond)
		select {
		case <- ctx.Done():
			return fmt.Errorf("等待锁锁超时")
		case <- waiterAfter :
			//尝试获取锁
			res := r.Redis.Eval(LockScript,[]string{r.Key},[]string{strconv.Itoa(int(r.TTL.Milliseconds())),r.Id })
			if res.Err() != nil {
				return res.Err()
			}
			val,ok := res.Val().(string)
			if ok && val == "OK" {
				ctx, cancel := context.WithTimeout(r.ctx,r.TimeOut)
				r.cancel = cancel
				go KeepAliveWorker(ctx,r.Redis,r.PerKeepAlive,r.TTL,r.Key,r.Id)
				return nil
			}
			ttl = time.Duration(res.Val().(int64))
		}
	}
}

//锁的释放，当key不在了的时候，取消续租任务
func (r *RedisMutex)Unlock()(err error) {
	res := r.Redis.Eval(UnlockScript,[]string{r.Key},[]string{r.Id})
	if res.Err() != nil {
		return res.Err()
	}
	val := res.Val().(string)
	//锁不存在或者被删除了
	if val == "OK" && r.cancel != nil{
		//取消续租任务
		r.cancel()
		r.cancel = nil
	}

	return
}

//续租任务 通过超时和主动取消控制
func KeepAliveWorker(ctx context.Context,redis *redis.Client, PerKeepAlive,ttl time.Duration,key ,id string ) {
	ticker := time.Tick(PerKeepAlive)
	for {
		select {
		case <- ctx.Done():
			return
		case <- ticker:
			// 执行续租任务
			res := redis.Eval(KeepAliveScript,[]string{key},[]string{strconv.Itoa(int(ttl.Milliseconds())),id })
			if res.Err() != nil {
				log.Fatal(res.Err().Error())
			}
			if res.Val().(string) != "OK" {
				return
			}
		}
	}
}

总结

这里写的比较粗糙，虽然我已经测试过了，但也可能会有一些问题，欢迎大家指出问题共同学习。