SET key value [EX seconds] [PX milliseconds] [NX|XX]

-- Lua 腳本實現(xiàn) Redis 分布式鎖

-- 生成唯一標識
local requestId = ARGV[1]

-- 嘗試獲取鎖
local lockKey = KEYS[1]
local lockValue = requestId
local lockExpireTime = tonumber(ARGV[2])

local result = redis.call('SET', lockKey, lockValue, 'NX', 'PX', lockExpireTime)

-- 判斷獲取鎖的結(jié)果
if result == 'OK' then
    -- 獲取鎖成功，設(shè)置鎖的過期時間
    return 'OK'
else
    -- 獲取鎖失敗
    return 'FAIL'
end

if redis.call("get", KEYS[1]) == ARGV[1] then

    return redis.call("del", KEYS[1])

else

    return 0

end

package redis

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

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

const (
    letters  = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
    randomLen = 16

    // 默認超時時間，用來防止死鎖
    tolerance       = 300 // milliseconds
    millisPerSecond = 800

    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`

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

)

type redisLock struct {
    // redis客戶端
    store *Redis
    // 超時時間
    seconds uint32
    // 鎖key
    keys string
    // 鎖value，防止鎖被別人獲取到
    value string
}

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

// NewRedisLock returns a RedisLock.
func NewRedisLock(store *Redis, keys string) *RedisLock {
    return &RedisLock{
        store: store,
        keys:   keys,
        // 獲取鎖時，鎖的值通過隨機字符串生成
        // 實際上go-zero提供更加高效的隨機字符串生成方式
        // 見core/stringx/random.go：Randn
        value:    randomStr(randomLen),
    }
}

// Acquire acquires the lock.
// 加鎖
func (rl *RedisLock) Acquire() (bool, error) {
    // 獲取過期時間
    seconds := atomic.LoadUint32(&rl.seconds)
    // 默認鎖過期時間為500ms，防止死鎖
    resp, err := rl.store.Eval(lockCommand, []string{rl.keys}, []string{
        rl.value, strconv.Itoa(int(seconds)*millisPerSecond + tolerance),
    })
    if err == red.Nil {
        return false, nil
    } else if err != nil {
        logx.Errorf("Error on 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 lock for %s: %v", rl.keys, resp)
    return false, nil
}

// Release releases the lock.
// 釋放鎖
func (rl *RedisLock) Release() (bool, error) {
    resp, err := rl.store.Eval(delCommand, []string{rl.keys}, []string{rl.value})
    if err != nil {
        return false, err
    }

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

    return reply == 1, nil
}

func randomStr(n int) string {
    b := make([]byte, n)
    for i := range b {
        b[i] = letters[rand.Intn(len(letters))]
    }
    return string(b)
}

// SetExpire sets the expire.
// 需要注意的是需要在Acquire()之前調(diào)用
// 不然默認為300ms自動釋放
func (rl *RedisLock) SetExpire(seconds int) {
    atomic.StoreUint32(&rl.seconds, uint32(seconds))
}