go-redis实现分布式锁

go-redis实现分布式锁

介绍

默认阻塞

在这种情况下只进行一次尝试获取锁，失败就停止了。

自旋锁

在这个模式下，会尝试获取锁，当失败后会尝试自旋不断的尝试，直到获取了锁。

ticker表示每次自旋的时间间隔，CAStime表示总共的自旋时间，超出后停止自旋。

在外部还有一个context用来控制整个goroutine运行时间

看门狗策略

原本我们设定了固定的redis锁时间，但有些任务时间长，有些任务时间短，看门狗机制下我们设定一个短TTL，启动一个 定时 goroutine，在锁即将过期时，检查自己是否仍然持有锁，并续约。

运行方式

• 在redis_lock目录下直接运行代码

go run .

• 在redis_lock目录下直接运行打包好的文件

./redis_lock

代码结构

.
├── go.mod
├── go.sum
├── README.md
└── redis_lock
    ├── lock.go
    ├── main.go
    ├── redis_lock
    └── watchdog.go

代码

main.go

package main

import (
	"context"
	"sync"
	"time"
)

var wait sync.WaitGroup

func test1(newlock *RedisLock) {
	C, cancel := context.WithTimeout(context.Background(), 10*time.Second) //超过该时间会停止goroutine运行
	defer cancel()
	wait.Add(1)
	go func() {
		islocked := newlock.TryLock(C) //尝试获取锁
		time.Sleep(5 * time.Second)    //模拟运行时间
		newlock.Unlock(islocked)       //解锁
		wait.Done()
	}()
	wait.Wait() //防止test退出导致C发送超时时间导致goroutine停止
}

// 测试用例，模拟两个goroutine同时争夺锁
func test2(newlock1, newlock2 *RedisLock) {
	C, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	wait.Add(2)
	go func() {
		islocked := newlock1.TryLock(C)
		time.Sleep(5 * time.Second)
		newlock1.Unlock(islocked)
		wait.Done()
	}()
	go func() {
		islocked := newlock2.TryLock(C)
		time.Sleep(5 * time.Second)
		newlock2.Unlock(islocked)
		wait.Done()
	}()
	wait.Wait()
}
func main() {
	RedisClient := InitRedis()            //初始化redis
	var model, watchdog bool = true, true //model[false表示默认阻塞，true表示自旋];watchdog[false关闭，true开启]
	lockKey := "my_lock"                  //锁的名称
	lockValue := GetCurrentID()           //锁的值[不同goroutine唯一身份标识]
	experation := 1 * time.Second         //TTL时间

	lock1 := NewRedisLock(RedisClient, lockKey, lockValue, experation, model, watchdog)
	defer lock1.RedisClient.Close()

	test1(lock1) //测试看门狗

	//测试多个goroutine对锁进行竞争
	lock2 := NewRedisLock(RedisClient, lockKey, lockValue, experation, model, watchdog)
	lock3 := NewRedisLock(RedisClient, lockKey, lockValue, experation, model, watchdog)

	defer lock2.RedisClient.Close()
	defer lock3.RedisClient.Close()

	test2(lock2, lock3)
}

lock.go

package main

import (
	"context"
	"log"
	"os"
	"strconv"
	"time"

	"github.com/redis/go-redis/v9"
)

var ctx = context.Background()

type RedisLock struct {
	RedisClient        *redis.Client
	LockKey, LockValue string
	TTL                time.Duration
	model, watchdog    bool
	StopChan           chan struct{}
}

// 初始化redis
func InitRedis() *redis.Client {
	RedisClient := redis.NewClient(&redis.Options{
		Addr:     "127.0.0.1:6379",
		Password: "",
		DB:       0,
	})
	_, err := RedisClient.Ping(ctx).Result()
	if err != nil {
		log.Fatalf("无法连接redis : %v", err)
	}
	return RedisClient
}

func NewRedisLock(client *redis.Client, key, value string, ttl time.Duration, model, watchdog bool) *RedisLock {
	return &RedisLock{
		RedisClient: client,
		LockKey:     key,
		LockValue:   value,
		TTL:         ttl,
		model:       model,
		watchdog:    watchdog,
		StopChan:    make(chan struct{}),
	}
}

func (Lock *RedisLock) TryLock(c context.Context) bool {
	ticker := time.NewTicker(100 * time.Millisecond) //每次自旋的时间间隔0.1s
	defer ticker.Stop()

	CAStime := time.After(7 * time.Second) //自旋的总时间

	//默认阻塞模式和自旋模式
	if !Lock.model { //阻塞
		locked, err := acquireLock(Lock.RedisClient, Lock.LockKey, Lock.LockValue, Lock.TTL)
		if err != nil {
			log.Println("获取锁失败", err)
			return false
		}
		if locked {
			log.Println("获取锁成功")
			if Lock.watchdog {
				go renewLock(Lock)
			}
			return true
		}
	} else { //自旋
		for {
			select {
			case <-c.Done(): //context超时终止
				return false
			case <-CAStime: //自旋结束
				return false
			case <-ticker.C: //自旋
				locked, err := acquireLock(Lock.RedisClient, Lock.LockKey, Lock.LockValue, Lock.TTL)
				if err != nil {
					log.Println("redis出现bug:", err)
					return false
				}
				if locked {
					log.Println("获取锁成功")
					if Lock.watchdog {
						go renewLock(Lock)
					}
					return true
				} else {
					log.Println("锁已经被占用")
					continue
				}
			}
		}
	}
	return false
}

// 先关闭看门狗再释放redis锁，防止释放完锁以后看门狗发现锁没了报错
func (Lock *RedisLock) Unlock(locked bool) {
	if locked {
		if Lock.watchdog {
			close(Lock.StopChan)
		}
		unlocked, err := releaseLock(Lock.RedisClient, Lock.LockKey, Lock.LockValue)
		if err != nil {
			log.Println("释放锁失败", err)
		} else if unlocked {
			log.Println("释放锁成功")
		} else {
			log.Println("超时,锁已经被redis释放！")
		}
	} else {
		log.Println("锁被占用，获取失败！")
	}
}

// 获取锁，如果成功返回[true，nil],失败返回[false,nil],发生错误如数据库无法连接返回[false，err]
func acquireLock(
	client *redis.Client, lockkey, lockvalue string, experation time.Duration) (bool, error) {
	success, err := client.SetNX(ctx, lockkey, lockvalue, experation).Result()
	if err != nil {
		return false, err
	}
	return success, nil
}

// 释放锁
func releaseLock(
	client *redis.Client, lockKey, lockValue string) (bool, error) {
	luaScript := `
		if redis.call("GET", KEYS[1]) == ARGV[1] then
			return redis.call("DEL", KEYS[1])
		else
			return 0
		end`

	// 只有持有锁的客户端才可以删除
	result, err := client.Eval(ctx, luaScript, []string{lockKey}, lockValue).Int()
	if err != nil {
		return false, err
	}
	return result == 1, nil
}

// redis键值对的值，可自定义，比如用户的个人标识
func GetCurrentID() string {
	return strconv.Itoa(os.Getpid())
}

watchdog.go

package main

import (
	"fmt"
	"log"
	"time"

	"github.com/redis/go-redis/v9"
)

// 续约机制（看门狗）
func renewLock(Lock *RedisLock) {
	ticker := time.NewTicker(Lock.TTL / 2) // 在TTL的一半时间时尝试续约

	defer ticker.Stop()

	for {
		select {
		case <-Lock.StopChan:
			log.Println("关闭，看门狗")
			return
		case <-ticker.C:
			// 只续约自己持有的锁
			val, err := Lock.RedisClient.Get(ctx, Lock.LockKey).Result()
			if err == redis.Nil || err != nil || val != Lock.LockValue {
				fmt.Println("锁丢失，停止续时")
				return
			}
			// 续约锁
			Lock.RedisClient.PExpire(ctx, Lock.LockKey, Lock.TTL)
			fmt.Println("续时")
		}
	}
}