modelRT/distributedlock/luascript/rwlock_script.go

// Package luascript defines the lua script used for redis distributed lock
package luascript

// RLockScript is the lua script for the lock read lock command
/*
KEYS[1]:锁的键名（key）,通常是锁的唯一标识。
KEYS[2]:锁的超时键名前缀（rwTimeoutPrefix）,用于存储每个读锁的超时键。
ARGV[1]:锁的过期时间（lockLeaseTime）,单位为秒。
ARGV[2]:当前客户端的唯一标识（token）,用于区分不同的客户端。
*/
var RLockScript = `
local mode = redis.call('hget', KEYS[1], 'mode');
local lockKey = KEYS[2] .. ':' .. ARGV[2];
if (mode == false) then
    redis.call('hset', KEYS[1], 'mode', 'read');
    redis.call('hset', KEYS[1], lockKey, '1');
    redis.call('hpexpire', KEYS[1], ARGV[1], 'fields', '1', lockKey);
    redis.call('pexpire', KEYS[1], ARGV[1]);
    return 1;
end;

if (mode == 'write') then
    -- 放到 list 中等待写锁释放后再次尝试加锁并且订阅写锁释放的消息
    local waitKey = KEYS[1] .. ':read'; 
    redis.call('rpush', waitKey, ARGV[2]);
    return -1;
end;

if (mode == 'read') then
    if (redis.call('exists', KEYS[1], ARGV[2]) == 1) then
        redis.call('hincrby', KEYS[1], lockKey, '1');
        local remainTime = redis.call('hpttl', KEYS[1], 'fields', '1', lockKey);
        redis.call('hpexpire', KEYS[1], math.max(tonumber(remainTime[1]), ARGV[1]), 'fields', '1', lockKey);
    else
        redis.call('hset', KEYS[1], lockKey, '1');
        redis.call('hpexpire', KEYS[1], ARGV[1], 'fields', '1', lockKey);
    end;
    local cursor = 0;
    local maxRemainTime = tonumber(ARGV[1]);
    local pattern = KEYS[2] .. ':*';
    repeat
        local hscanResult = redis.call('hscan', KEYS[1], cursor, 'match', pattern, 'count', '100');
        cursor = tonumber(hscanResult[1]);
        local fields = hscanResult[2];

        for i = 1, #fields,2 do
            local field = fields[i];
            local remainTime = redis.call('hpttl', KEYS[1], 'fields', '1', field);
            maxRemainTime = math.max(tonumber(remainTime[1]), maxRemainTime);
        end;
    until cursor == 0;

    local remainTime = redis.call('pttl', KEYS[1]);
    redis.call('pexpire', KEYS[1], math.max(tonumber(remainTime),maxRemainTime));
    return 1;
end;
`

// UnRLockScript is the lua script for the unlock read lock command
/*
KEYS[1]:锁的键名（key）,通常是锁的唯一标识。
KEYS[2]:锁的超时键名前缀（rwTimeoutPrefix）,用于存储每个读锁的超时键。
KEYS[3]:锁的释放通知写频道（chankey）,用于通知其他写等待客户端锁已释放。
ARGV[1]:解锁消息（unlockMessage）,用于通知其他客户端锁已释放。
ARGV[2]:当前客户端的唯一标识（token）,用于区分不同的客户端。
*/
var UnRLockScript = `
local lockKey = KEYS[2] .. ':' .. ARGV[2];
local mode = redis.call('hget', KEYS[1], 'mode');
if (mode == false) then
    local writeWait = KEYS[1] .. ':write';
    -- 优先写锁加锁
    local counter = redis.call('llen',writeWait);
    if (counter >= 1) then
        redis.call('publish', KEYS[3], ARGV[1]);
    end;
    return 1;
elseif (mode == 'write') then
    return -2;
end;

-- 判断当前的确是读模式但是当前 token 并没有加读锁的情况，返回 0
local lockExists = redis.call('hexists', KEYS[1], lockKey);
if ((mode == 'read') and (lockExists == 0)) then
    return 0;
end;

local counter = redis.call('hincrby', KEYS[1], lockKey, -1);
local delTTLs = redis.call('hpttl', KEYS[1], 'fields', '1', lockKey);
local delTTL = tonumber(delTTLs[1]);
if (counter == 0) then
    redis.call('hdel', KEYS[1], lockKey);
end;

if (redis.call('hlen', KEYS[1]) > 1) then
    local cursor = 0;
    local maxRemainTime = 0;
    local pattern = KEYS[2] .. ':*';
    repeat
        local hscanResult = redis.call('hscan', KEYS[1], cursor, 'match', pattern, 'count', '100');
        cursor = tonumber(hscanResult[1]);
        local fields = hscanResult[2];

        for i = 1, #fields,2 do
            local field = fields[i];
            local remainTime = redis.call('hpttl', KEYS[1], 'fields', '1', field);
            maxRemainTime = math.max(tonumber(remainTime[1]), maxRemainTime);
        end;
    until cursor == 0;

    if (maxRemainTime > 0) then
        if (delTTL > maxRemainTime) then
            redis.call('pexpire', KEYS[1], maxRemainTime);
        else
            local remainTime = redis.call('pttl', KEYS[1]);
            redis.call('pexpire', KEYS[1], math.max(tonumber(remainTime),maxRemainTime));
        end;
    end;
else
    redis.call('del', KEYS[1]);
    local writeWait = KEYS[1] .. ':write';
    -- 优先写锁加锁
    local counter = redis.call('llen',writeWait);
    if (counter >= 1) then
        redis.call('publish', KEYS[3], ARGV[1]);
    end;
    return 1;
end;
`

// WLockScript is the lua script for the lock write lock command
/*
KEYS[1]:锁的键名（key）,通常是锁的唯一标识。
KEYS[2]:锁的超时键名前缀（rwTimeoutPrefix）,用于存储每个读锁的超时键。
ARGV[1]:锁的过期时间（lockLeaseTime）,单位为秒。
ARGV[2]:当前客户端的唯一标识（token）,用于区分不同的客户端。
*/
var WLockScript = `
local mode = redis.call('hget', KEYS[1], 'mode');
local lockKey = KEYS[2] .. ':' .. ARGV[2];
local waitKey = KEYS[1] .. ':write';
if (mode == false) then
    local waitListLen = redis.call('llen', waitKey);
    if (waitListLen > 0) then
        local firstToken = redis.call('lindex', waitKey,'0');
        if (firstToken ~= ARGV[2]) then
            return -7;
        end;
    end;
    redis.call('hset', KEYS[1], 'mode', 'write');
	redis.call('hset', KEYS[1], lockKey, 1);
    redis.call('hpexpire', KEYS[1], ARGV[1], 'fields', '1', lockKey);
	redis.call('pexpire', KEYS[1], ARGV[1]);
    redis.call('lpop', waitKey, '1');
	return 1;
elseif (mode == 'read') then
    -- 放到 list 中等待读锁释放后再次尝试加锁并且订阅读锁释放的消息
    redis.call('rpush', waitKey, ARGV[2]);
    return -3;
else
    -- 可重入写锁逻辑
    local lockKey = KEYS[2] .. ':' .. ARGV[2];
    local lockExists = redis.call('hexists', KEYS[1], lockKey);
    if (lockExists == 1) then
        redis.call('hincrby', KEYS[1], lockKey, 1);
        redis.call('hpexpire', KEYS[1], ARGV[1], 'fields', '1', lockKey);
        redis.call('pexpire', KEYS[1], ARGV[1]);
        return 1;
    end;
    -- 放到 list 中等待写锁释放后再次尝试加锁并且订阅写锁释放的消息
    local key = KEYS[1] .. ':write';
    redis.call('rpush', key, ARGV[2]);
    return -4;
end;
`

// UnWLockScript is the lua script for the unlock write lock command
/*
KEYS[1]:锁的键名（key）,通常是锁的唯一标识。
KEYS[2]:锁的超时键名前缀（rwTimeoutPrefix）,用于存储每个读锁的超时键。
KEYS[3]:锁的释放通知写频道（writeChankey）,用于通知其他写等待客户端锁已释放。
KEYS[4]:锁的释放通知读频道（readChankey）,用于通知其他读等待客户端锁已释放。
ARGV[1]:解锁消息（unlockMessage）,用于通知其他客户端锁已释放。
ARGV[2]:当前客户端的唯一标识（token）,用于区分不同的客户端。
*/
var UnWLockScript = `
local mode = redis.call('hget', KEYS[1], 'mode');
local writeWait = KEYS[1] .. ':write';
if (mode == false) then
    -- 优先写锁加锁,无写锁的情况通知读锁加锁
    local counter = redis.call('llen',writeWait);
    if (counter >= 1) then
        redis.call('publish', KEYS[3], ARGV[1]);
    else
        redis.call('publish', KEYS[4], ARGV[1]);
    end;
	return 1;
elseif (mode == 'read') then
    return -5;
else
    local lockKey = KEYS[2] .. ':' .. ARGV[2];
    local lockExists = redis.call('hexists', KEYS[1], lockKey);
    if (lockExists >= 1) then
        -- 可重入写锁逻辑
        local incrRes = redis.call('hincrby', KEYS[1], lockKey, -1);
        if (incrRes == 0) then
            redis.call('del', KEYS[1]);
            local counter = redis.call('llen',writeWait);
            if (counter >= 1) then
                redis.call('publish', KEYS[3], ARGV[1]);
            else
                redis.call('publish', KEYS[4], ARGV[1]);
            end;
            return 1;
        end;
        return 0;
    else
        return -6;
    end;
end;
`

// RefreshRWLockScript is the lua script for the refresh lock command
/*
KEYS[1]:锁的键名（key）,通常是锁的唯一标识。
KEYS[2]:锁的超时键名前缀（rwTimeoutPrefix）,用于存储每个读锁的超时键。
ARGV[1]:锁的过期时间（lockLeaseTime）,单位为秒。
ARGV[2]:当前客户端的唯一标识（token）,用于区分不同的客户端。
*/
var RefreshRWLockScript = `
local lockKey = KEYS[2] .. ':' .. ARGV[2];
local lockExists = redis.call('hexists', KEYS[1], lockKey);
local mode = redis.call('hget', KEYS[1], 'mode');
local maxRemainTime = tonumber(ARGV[1]);
if (lockExists == 1) then
    redis.call('hpexpire', KEYS[1], ARGV[1], 'fields', '1', lockKey);
    if (mode == 'read') then
        local cursor = 0;
        local pattern = KEYS[2] .. ':*';
        repeat
            local hscanResult = redis.call('hscan', KEYS[1], cursor, 'match', pattern, 'count', '100');
            cursor = tonumber(hscanResult[1]);
            local fields = hscanResult[2];

            for i = 1, #fields,2 do
                local field = fields[i];
                local remainTime = redis.call('hpttl', KEYS[1], 'fields', '1', field);
                maxRemainTime = math.max(tonumber(remainTime[1]), maxRemainTime);
            end;
        until cursor == 0;
        
        if (maxRemainTime > 0) then
            local remainTime = redis.call('pttl', KEYS[1]);
            redis.call('pexpire', KEYS[1], math.max(tonumber(remainTime),maxRemainTime));
        end;
    elseif (mode == 'write') then
        redis.call('pexpire', KEYS[1], ARGV[1]);
    end;
    -- return redis.call('pttl',KEYS[1]);
    return 1;
end;
return -8;
`