woop/internal/tikv/redis_helper/lock.go

package redis_helper

import (
	"context"
	"crypto/rand"
	"encoding/base64"
	"errors"
	"github.com/go-redis/redis/v8"
)

type LockResult int

const (
	_ LockResult = iota
	LockResultSuccess
	LockResultRenewalSuccess
	LockResultFail
)

type RedisPreempt struct {
	key                      string
	password                 string
	lockScript, unlockScript *redis.Script
	lastLockStatus           bool
}

// CreatePreempt used to create a redis preempt instance
func CreatePreempt(key string) *RedisPreempt {
	p := &RedisPreempt{
		key: key,
	}
	p.init()

	return p
}

// init redis preempt instance and some script
func (p *RedisPreempt) init() {
	byt := make([]byte, 18)
	_, _ = rand.Read(byt)
	p.password = base64.StdEncoding.EncodeToString(byt)
	p.lockScript = redis.NewScript(`
		local val = redis.call('get', KEYS[1])
		if (val==nil or (type(val) == "boolean" and not val)) then
			redis.call('set', KEYS[1], ARGV[1], 'ex', ARGV[2])
			return 'LockResultSuccess'
		elseif (val == ARGV[1]) then
			redis.call('expire', KEYS[1], ARGV[2])
			return 'LockResultRenewalSuccess'
		else 
			return 'LockResultFail'
		end
	`)
	p.unlockScript = redis.NewScript(`
		if redis.call('get',KEYS[1]) == ARGV[1] then
			return redis.call('del', KEYS[1])
		else
			return 0
		end
	`)
}

// TryLock attempt to lock the master for ttlSecond
func (p *RedisPreempt) TryLock(ttlSecond int) (result LockResult, err error) {
	ret, err := p.lockScript.Run(context.Background(), redisInstance, []string{p.key}, p.password, ttlSecond).Text()
	if err != nil {
		return LockResultFail, err
	}
	switch ret {
	case "LockResultRenewalSuccess":
		p.lastLockStatus = true
		return LockResultRenewalSuccess, nil
	case "LockResultSuccess":
		p.lastLockStatus = true
		return LockResultSuccess, nil
	case "LockResultFail":
		return LockResultFail, nil
	default:
		return LockResultFail, errors.New("fixme: unknown return")
	}
}

// Unlock try to release the master permission
func (p *RedisPreempt) Unlock() (bool, error) {
	if p == nil {
		return false, nil
	}

	return p.unlockScript.Run(context.Background(), redisInstance, []string{p.key}, p.password).Bool()
}

// LastLockStatus get the last preempt status
func (p *RedisPreempt) LastLockStatus() bool {
	if p == nil {
		return false
	}

	return p.lastLockStatus
}