284 lines
10 KiB
Lua
284 lines
10 KiB
Lua
local operation = ARGV[1]
|
|
local max_safe_integer = 9007199254740991
|
|
|
|
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
|
|
|
|
local function read_counter(key)
|
|
local encoded = redis.call('GET', key)
|
|
if not encoded then
|
|
return 0, nil
|
|
end
|
|
local value = tonumber(encoded)
|
|
if not value or value < 0 or value ~= math.floor(value) then
|
|
return nil, 'invalid'
|
|
end
|
|
return value, nil
|
|
end
|
|
|
|
local function read_permit(token)
|
|
local encoded = redis.call('HGET', KEYS[8], token)
|
|
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.state) ~= 'string' or
|
|
type(value.expected) ~= 'number' or value.expected <= 0 or
|
|
value.expected ~= math.floor(value.expected) then
|
|
return nil, 'invalid'
|
|
end
|
|
return value, nil
|
|
end
|
|
|
|
local function cleanup_expired(now)
|
|
local expired = redis.call('ZRANGEBYSCORE', KEYS[9], '-inf', now, 'LIMIT', 0, 256)
|
|
if #expired == 0 then
|
|
return nil
|
|
end
|
|
local fetched, fetched_error = read_counter(KEYS[6])
|
|
local pending, pending_error = read_counter(KEYS[7])
|
|
if fetched_error or pending_error then
|
|
return 'invalid'
|
|
end
|
|
for _, permit_token in ipairs(expired) do
|
|
local permit, permit_error = read_permit(permit_token)
|
|
if permit_error then
|
|
return 'invalid'
|
|
end
|
|
if permit and permit.state == 'reserved' then
|
|
if pending < permit.expected or fetched > max_safe_integer - permit.expected then
|
|
return 'invalid'
|
|
end
|
|
pending = pending - permit.expected
|
|
fetched = fetched + permit.expected
|
|
end
|
|
redis.call('HDEL', KEYS[8], permit_token)
|
|
redis.call('ZREM', KEYS[5], permit_token)
|
|
redis.call('ZREM', KEYS[9], permit_token)
|
|
end
|
|
redis.call('SET', KEYS[6], fetched)
|
|
redis.call('SET', KEYS[7], pending)
|
|
return nil
|
|
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
|
|
local created = redis.call('SET', KEYS[1], generation_candidate, 'NX')
|
|
if created then
|
|
redis.call('DEL', KEYS[2], KEYS[3], KEYS[4], KEYS[5], KEYS[6], KEYS[7], KEYS[8], KEYS[9])
|
|
end
|
|
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 expected = tonumber(ARGV[10])
|
|
local max_total = tonumber(ARGV[11])
|
|
if not expected or expected <= 0 or expected ~= math.floor(expected) or
|
|
expected > max_safe_integer or not max_total or max_total < 0 or
|
|
max_total > max_safe_integer or max_total ~= math.floor(max_total) then
|
|
return reply('invalid', generation, epoch or 0, 0)
|
|
end
|
|
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()
|
|
if cleanup_expired(now) then
|
|
return reply('unavailable', generation, epoch, 0)
|
|
end
|
|
local existing, existing_error = read_permit(permit_token)
|
|
if existing_error then
|
|
return reply('unavailable', generation, epoch, 0)
|
|
end
|
|
if existing and existing.state == 'reserved' then
|
|
return reply('ok', generation, epoch, 0)
|
|
end
|
|
if existing then
|
|
return reply('unavailable', generation, epoch, 0)
|
|
end
|
|
local fetched, fetched_error = read_counter(KEYS[6])
|
|
local pending, pending_error = read_counter(KEYS[7])
|
|
if fetched_error or pending_error then
|
|
return reply('unavailable', generation, epoch, 0)
|
|
end
|
|
if max_total > 0 and fetched + pending + expected > max_total then
|
|
return reply('quota_exhausted', 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('ZADD', KEYS[9], now + permit_ttl, permit_token)
|
|
redis.call('HSET', KEYS[8], permit_token, cjson.encode({
|
|
version = 1,
|
|
state = 'reserved',
|
|
expected = expected
|
|
}))
|
|
redis.call('SET', KEYS[7], pending + expected)
|
|
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 == 'complete_fetch' or operation == 'cancel_fetch' then
|
|
local permit_token = ARGV[2]
|
|
local fetched_count = tonumber(ARGV[3])
|
|
local settlement_ttl = tonumber(ARGV[4])
|
|
if not fetched_count or fetched_count < 0 or fetched_count ~= math.floor(fetched_count) or
|
|
fetched_count > max_safe_integer or not settlement_ttl or settlement_ttl <= 0 then
|
|
return reply('invalid', '', 0, 0)
|
|
end
|
|
local now = now_ms()
|
|
if cleanup_expired(now) then
|
|
return reply('unavailable', '', 0, 0)
|
|
end
|
|
local permit, permit_error = read_permit(permit_token)
|
|
if permit_error then
|
|
return reply('unavailable', '', 0, 0)
|
|
end
|
|
if not permit or permit.state ~= 'reserved' then
|
|
return reply('ok', '', 0, 0)
|
|
end
|
|
local fetched, fetched_error = read_counter(KEYS[6])
|
|
local pending, pending_error = read_counter(KEYS[7])
|
|
if fetched_error or pending_error or pending < permit.expected or
|
|
fetched > max_safe_integer - fetched_count then
|
|
return reply('unavailable', '', 0, 0)
|
|
end
|
|
pending = pending - permit.expected
|
|
if operation == 'complete_fetch' then
|
|
fetched = fetched + fetched_count
|
|
permit.state = 'completed'
|
|
else
|
|
permit.state = 'cancelled'
|
|
end
|
|
redis.call('SET', KEYS[6], fetched)
|
|
redis.call('SET', KEYS[7], pending)
|
|
redis.call('HSET', KEYS[8], permit_token, cjson.encode(permit))
|
|
redis.call('ZREM', KEYS[5], permit_token)
|
|
redis.call('ZADD', KEYS[9], now + settlement_ttl, permit_token)
|
|
return reply('ok', '', 0, 0)
|
|
end
|
|
|
|
return reply('invalid', '', 0, 0)
|