go-zero/core/stores/redis/redislock.go

package redis

import (
	"math/rand"
	"sync/atomic"
	"time"

	red "github.com/go-redis/redis"
	"github.com/tal-tech/go-zero/core/logx"
	"github.com/tal-tech/go-zero/core/stringx"
)

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

// A RedisLock is a redis lock.
type RedisLock struct {
	store   *Redis
	seconds uint32
	count   int32
	key     string
	id      string
}

func init() {
	rand.Seed(time.Now().UnixNano())
}

// NewRedisLock returns a RedisLock.
func NewRedisLock(store *Redis, key string) *RedisLock {
	return &RedisLock{
		store: store,
		key:   key,
		id:    stringx.Randn(randomLen),
	}
}

// Acquire acquires the lock.
func (rl *RedisLock) Acquire() (bool, error) {
	newCount := atomic.AddInt32(&rl.count, 1)
	if newCount > 1 {
		return true, nil
	}

	seconds := atomic.LoadUint32(&rl.seconds)
	ok, err := rl.store.SetnxEx(rl.key, rl.id, int(seconds+1)) // +1s for tolerance
	if err == red.Nil {
		atomic.AddInt32(&rl.count, -1)
		return false, nil
	} else if err != nil {
		atomic.AddInt32(&rl.count, -1)
		logx.Errorf("Error on acquiring lock for %s, %s", rl.key, err.Error())
		return false, err
	} else if !ok {
		atomic.AddInt32(&rl.count, -1)
		return false, nil
	}

	return true, nil
}

// Release releases the lock.
func (rl *RedisLock) Release() (bool, error) {
	newCount := atomic.AddInt32(&rl.count, -1)
	if newCount > 0 {
		return true, nil
	}

	resp, err := rl.store.Eval(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
}

// SetExpire sets the expiration.
func (rl *RedisLock) SetExpire(seconds int) {
	atomic.StoreUint32(&rl.seconds, uint32(seconds))
}