159 lines
5.4 KiB
Lua
159 lines
5.4 KiB
Lua
local operation = ARGV[1]
|
|
|
|
local function now_ms()
|
|
local value = redis.call('TIME')
|
|
return tonumber(value[1]) * 1000 + math.floor(tonumber(value[2]) / 1000)
|
|
end
|
|
|
|
local function reply(status, generation, epoch, wait_ms)
|
|
return cjson.encode({
|
|
status = status,
|
|
generation = generation,
|
|
epoch = tostring(epoch or '0'),
|
|
waitMs = wait_ms
|
|
})
|
|
end
|
|
|
|
local function read_leader()
|
|
local encoded = redis.call('GET', KEYS[3])
|
|
if not encoded then
|
|
return nil, nil
|
|
end
|
|
local ok, value = pcall(cjson.decode, encoded)
|
|
if not ok or type(value) ~= 'table' or type(value.generation) ~= 'string' or
|
|
type(value.holderId) ~= 'string' or type(value.token) ~= 'string' or
|
|
type(value.epoch) ~= 'string' then
|
|
return nil, 'invalid'
|
|
end
|
|
return value, nil
|
|
end
|
|
|
|
local function same_leader(value, generation, holder_id, token, epoch)
|
|
return value and value.generation == generation and value.holderId == holder_id and
|
|
value.token == token and value.epoch == tostring(epoch)
|
|
end
|
|
|
|
if operation == 'acquire_leader' then
|
|
local generation_candidate = ARGV[2]
|
|
local holder_id = ARGV[3]
|
|
local token = ARGV[4]
|
|
local lease_ttl = tonumber(ARGV[5])
|
|
if not lease_ttl or lease_ttl <= 0 then
|
|
return reply('invalid', '', 0, 0)
|
|
end
|
|
redis.call('SET', KEYS[1], generation_candidate, 'NX')
|
|
local generation = redis.call('GET', KEYS[1])
|
|
local current, current_error = read_leader()
|
|
if current_error then
|
|
return reply('unavailable', generation, 0, 0)
|
|
end
|
|
if current then
|
|
if current.generation == generation and current.holderId == holder_id and current.token == token then
|
|
current.expiresAtMs = now_ms() + lease_ttl
|
|
redis.call('SET', KEYS[3], cjson.encode(current), 'PX', lease_ttl)
|
|
return reply('ok', generation, current.epoch, 0)
|
|
end
|
|
local remaining = redis.call('PTTL', KEYS[3])
|
|
return reply('busy', generation, 0, math.max(remaining, 1))
|
|
end
|
|
redis.call('INCR', KEYS[2])
|
|
local epoch = redis.call('GET', KEYS[2])
|
|
local leader = {
|
|
version = 1,
|
|
generation = generation,
|
|
holderId = holder_id,
|
|
token = token,
|
|
epoch = epoch,
|
|
expiresAtMs = now_ms() + lease_ttl
|
|
}
|
|
redis.call('SET', KEYS[3], cjson.encode(leader), 'PX', lease_ttl)
|
|
return reply('ok', generation, epoch, 0)
|
|
end
|
|
|
|
if operation == 'renew_leader' then
|
|
local generation = ARGV[2]
|
|
local holder_id = ARGV[3]
|
|
local token = ARGV[4]
|
|
local epoch = ARGV[5]
|
|
local lease_ttl = tonumber(ARGV[6])
|
|
local current, current_error = read_leader()
|
|
if current_error then
|
|
return reply('unavailable', generation, epoch or 0, 0)
|
|
end
|
|
if not same_leader(current, generation, holder_id, token, epoch) then
|
|
return reply('stale', generation, epoch or 0, 0)
|
|
end
|
|
current.expiresAtMs = now_ms() + lease_ttl
|
|
redis.call('SET', KEYS[3], cjson.encode(current), 'PX', lease_ttl)
|
|
return reply('ok', generation, epoch, 0)
|
|
end
|
|
|
|
if operation == 'release_leader' then
|
|
local generation = ARGV[2]
|
|
local holder_id = ARGV[3]
|
|
local token = ARGV[4]
|
|
local epoch = ARGV[5]
|
|
local current, current_error = read_leader()
|
|
if current_error then
|
|
return reply('unavailable', generation, epoch or 0, 0)
|
|
end
|
|
if same_leader(current, generation, holder_id, token, epoch) then
|
|
redis.call('DEL', KEYS[3])
|
|
end
|
|
return reply('ok', generation, epoch or 0, 0)
|
|
end
|
|
|
|
if operation == 'acquire_fetch' then
|
|
local generation = ARGV[2]
|
|
local holder_id = ARGV[3]
|
|
local leader_token = ARGV[4]
|
|
local epoch = ARGV[5]
|
|
local permit_token = ARGV[6]
|
|
local request_interval = tonumber(ARGV[7])
|
|
local max_in_flight = tonumber(ARGV[8])
|
|
local permit_ttl = tonumber(ARGV[9])
|
|
local current, current_error = read_leader()
|
|
if current_error then
|
|
return reply('unavailable', generation, epoch or 0, 0)
|
|
end
|
|
if not same_leader(current, generation, holder_id, leader_token, epoch) then
|
|
return reply('stale', generation, epoch or 0, 0)
|
|
end
|
|
local now = now_ms()
|
|
redis.call('ZREMRANGEBYSCORE', KEYS[5], '-inf', now)
|
|
local existing = redis.call('ZSCORE', KEYS[5], permit_token)
|
|
if existing then
|
|
return reply('ok', generation, epoch, 0)
|
|
end
|
|
local next_request = redis.call('GET', KEYS[4])
|
|
if next_request and not tonumber(next_request) then
|
|
return reply('unavailable', generation, epoch, 0)
|
|
end
|
|
if next_request and tonumber(next_request) > now then
|
|
return reply('rate_limited', generation, epoch, tonumber(next_request) - now)
|
|
end
|
|
if redis.call('ZCARD', KEYS[5]) >= max_in_flight then
|
|
local earliest = redis.call('ZRANGE', KEYS[5], 0, 0, 'WITHSCORES')
|
|
local wait_ms = 1
|
|
if earliest[2] then
|
|
wait_ms = math.max(tonumber(earliest[2]) - now, 1)
|
|
end
|
|
return reply('at_capacity', generation, epoch, wait_ms)
|
|
end
|
|
redis.call('ZADD', KEYS[5], now + permit_ttl, permit_token)
|
|
redis.call('PEXPIRE', KEYS[5], permit_ttl + 1000)
|
|
if request_interval > 0 then
|
|
redis.call('SET', KEYS[4], now + request_interval, 'PX', request_interval)
|
|
else
|
|
redis.call('DEL', KEYS[4])
|
|
end
|
|
return reply('ok', generation, epoch, 0)
|
|
end
|
|
|
|
if operation == 'release_fetch' then
|
|
redis.call('ZREM', KEYS[5], ARGV[2])
|
|
return reply('ok', '', 0, 0)
|
|
end
|
|
|
|
return reply('invalid', '', 0, 0)
|