330 lines
12 KiB
Lua
330 lines
12 KiB
Lua
local records_key = KEYS[1]
|
|
local unique_key = KEYS[2]
|
|
local idkeys_key = KEYS[3]
|
|
local expiry_key = KEYS[4]
|
|
local available_key = KEYS[5]
|
|
local inventory_key = KEYS[6]
|
|
local state_inventory_key = KEYS[7]
|
|
local owners_key = KEYS[8]
|
|
local owner_expiry_key = KEYS[9]
|
|
local epoch_key = KEYS[10]
|
|
local operation_key = KEYS[11]
|
|
|
|
local operation = ARGV[1]
|
|
local operation_ttl_ms = tonumber(ARGV[2])
|
|
local cleanup_limit = tonumber(ARGV[3])
|
|
local now_ms = tonumber(ARGV[4])
|
|
local proxy_id = ARGV[5]
|
|
local worker_id = ARGV[6]
|
|
local epoch = tonumber(ARGV[7])
|
|
local value = tonumber(ARGV[8])
|
|
local active = tonumber(ARGV[9])
|
|
local reserved = tonumber(ARGV[10])
|
|
local mutating = operation ~= 'get'
|
|
|
|
local function finish(reply)
|
|
local encoded = cjson.encode(reply)
|
|
if mutating then
|
|
redis.call('SET', operation_key, encoded, 'PX', operation_ttl_ms)
|
|
end
|
|
return encoded
|
|
end
|
|
|
|
if mutating then
|
|
local committed = redis.call('GET', operation_key)
|
|
if committed then
|
|
return committed
|
|
end
|
|
end
|
|
|
|
local function is_managed(state)
|
|
return state == 'FETCHED' or state == 'CHECKING' or state == 'AVAILABLE' or
|
|
state == 'SUSPECT' or state == 'DRAINING'
|
|
end
|
|
|
|
local function decrement_inventory(upstream)
|
|
if type(upstream) ~= 'string' or upstream == '' then
|
|
return
|
|
end
|
|
local count = redis.call('HINCRBY', inventory_key, upstream, -1)
|
|
if count < 0 then
|
|
redis.call('HSET', inventory_key, upstream, 0)
|
|
end
|
|
end
|
|
|
|
local function state_field(upstream, state)
|
|
return string.len(upstream) .. ':' .. upstream .. ':' .. state
|
|
end
|
|
|
|
local function is_counted(state)
|
|
return state == 'FETCHED' or state == 'CHECKING' or state == 'AVAILABLE' or
|
|
state == 'SUSPECT' or state == 'DRAINING' or state == 'UNHEALTHY' or state == 'EXTRACTED'
|
|
end
|
|
|
|
local function decrement_state(upstream, state)
|
|
if type(upstream) ~= 'string' or upstream == '' or not is_counted(state) then
|
|
return
|
|
end
|
|
local field = state_field(upstream, state)
|
|
local count = redis.call('HINCRBY', state_inventory_key, field, -1)
|
|
if count <= 0 then
|
|
redis.call('HDEL', state_inventory_key, field)
|
|
end
|
|
end
|
|
|
|
local function touch(key, expires_at_ms)
|
|
if redis.call('EXISTS', key) == 0 then
|
|
return
|
|
end
|
|
local current = redis.call('PEXPIRETIME', key)
|
|
if current < expires_at_ms then
|
|
redis.call('PEXPIREAT', key, expires_at_ms)
|
|
end
|
|
end
|
|
|
|
local function remove_available(id, record)
|
|
redis.call('ZREM', available_key, id)
|
|
local index_keys = record and record.indexKeys
|
|
if type(index_keys) == 'table' then
|
|
for _, index_key in ipairs(index_keys) do
|
|
if type(index_key) == 'string' and index_key ~= '' then
|
|
redis.call('ZREM', index_key, id)
|
|
end
|
|
end
|
|
end
|
|
end
|
|
|
|
local function remove_owned(id, record)
|
|
if record and type(record.ownerIndexKey) == 'string' and record.ownerIndexKey ~= '' then
|
|
redis.call('ZREM', record.ownerIndexKey, id)
|
|
end
|
|
end
|
|
|
|
local function add_available(id, record, at_ms)
|
|
local usable_until_ms = record and tonumber(record.usableUntilMs)
|
|
if not usable_until_ms or record.state ~= 'AVAILABLE' or usable_until_ms <= at_ms then
|
|
return
|
|
end
|
|
redis.call('ZADD', available_key, usable_until_ms, id)
|
|
touch(available_key, tonumber(record.expiresAtMs))
|
|
if type(record.indexKeys) == 'table' then
|
|
for _, index_key in ipairs(record.indexKeys) do
|
|
if type(index_key) == 'string' and index_key ~= '' then
|
|
redis.call('ZADD', index_key, usable_until_ms, id)
|
|
touch(index_key, tonumber(record.expiresAtMs))
|
|
end
|
|
end
|
|
end
|
|
end
|
|
|
|
local function remove_proxy(id)
|
|
local raw = redis.call('HGET', records_key, id)
|
|
local record = nil
|
|
if raw then
|
|
local decoded
|
|
decoded, record = pcall(cjson.decode, raw)
|
|
remove_available(id, decoded and record or nil)
|
|
remove_owned(id, decoded and record or nil)
|
|
if decoded and type(record) == 'table' and is_managed(record.state) then
|
|
decrement_inventory(record.sourceUpstream)
|
|
end
|
|
if decoded and type(record) == 'table' then
|
|
decrement_state(record.sourceUpstream, record.state)
|
|
end
|
|
else
|
|
redis.call('ZREM', available_key, id)
|
|
end
|
|
local digest = redis.call('HGET', idkeys_key, id)
|
|
if digest and redis.call('HGET', unique_key, digest) == id then
|
|
redis.call('HDEL', unique_key, digest)
|
|
end
|
|
redis.call('HDEL', idkeys_key, id)
|
|
redis.call('HDEL', records_key, id)
|
|
redis.call('ZREM', expiry_key, id)
|
|
redis.call('HDEL', owners_key, id)
|
|
redis.call('ZREM', owner_expiry_key, id)
|
|
end
|
|
|
|
local function cleanup_hard_expired(at_ms)
|
|
local expired = redis.call('ZRANGEBYSCORE', expiry_key, '-inf', at_ms, 'LIMIT', 0, cleanup_limit)
|
|
for _, id in ipairs(expired) do
|
|
remove_proxy(id)
|
|
end
|
|
end
|
|
|
|
local function decode_table(raw)
|
|
if not raw then
|
|
return nil
|
|
end
|
|
local decoded, value = pcall(cjson.decode, raw)
|
|
if not decoded or type(value) ~= 'table' then
|
|
return nil
|
|
end
|
|
return value
|
|
end
|
|
|
|
local function valid_assignment(assignment)
|
|
return assignment and assignment.version == 1 and type(assignment.proxyId) == 'string' and
|
|
assignment.proxyId ~= '' and type(assignment.workerId) == 'string' and assignment.workerId ~= '' and
|
|
tonumber(assignment.epoch) and tonumber(assignment.epoch) > 0 and
|
|
tonumber(assignment.assignmentVersion) and tonumber(assignment.assignmentVersion) > 0 and
|
|
tonumber(assignment.expiresAtMs) and tonumber(assignment.expiresAtMs) > 0 and
|
|
type(assignment.draining) == 'boolean'
|
|
end
|
|
|
|
local function clear_owner(id, assignment, at_ms, restore)
|
|
local raw_record = redis.call('HGET', records_key, id)
|
|
local record = decode_table(raw_record)
|
|
if record and (not assignment or record.ownerWorkerId == assignment.workerId) then
|
|
remove_owned(id, record)
|
|
record.ownerWorkerId = nil
|
|
redis.call('HSET', records_key, id, cjson.encode(record))
|
|
if restore then
|
|
add_available(id, record, at_ms)
|
|
end
|
|
end
|
|
redis.call('HDEL', owners_key, id)
|
|
redis.call('ZREM', owner_expiry_key, id)
|
|
end
|
|
|
|
if operation == 'assign' then
|
|
cleanup_hard_expired(now_ms)
|
|
local current_raw = redis.call('HGET', owners_key, proxy_id)
|
|
if current_raw then
|
|
local current = decode_table(current_raw)
|
|
if not valid_assignment(current) then
|
|
return finish({status = 'unavailable'})
|
|
end
|
|
if tonumber(current.expiresAtMs) > now_ms then
|
|
return finish({status = 'already_owned'})
|
|
end
|
|
clear_owner(proxy_id, current, now_ms, true)
|
|
end
|
|
|
|
local record = decode_table(redis.call('HGET', records_key, proxy_id))
|
|
if not record or record.state ~= 'AVAILABLE' or
|
|
type(record.ownerIndexKey) ~= 'string' or record.ownerIndexKey == '' or
|
|
(record.ownerWorkerId and record.ownerWorkerId ~= '') or
|
|
redis.call('HEXISTS', owners_key, proxy_id) == 1 or
|
|
not tonumber(record.usableUntilMs) or tonumber(record.usableUntilMs) <= now_ms then
|
|
return finish({status = 'unavailable'})
|
|
end
|
|
local expires_at_ms = now_ms + value
|
|
if tonumber(record.usableUntilMs) < expires_at_ms then
|
|
expires_at_ms = tonumber(record.usableUntilMs)
|
|
end
|
|
local next_epoch = redis.call('INCR', epoch_key)
|
|
redis.call('PERSIST', epoch_key)
|
|
local assignment = {
|
|
version = 1,
|
|
proxyId = proxy_id,
|
|
workerId = worker_id,
|
|
epoch = next_epoch,
|
|
assignmentVersion = 1,
|
|
expiresAtMs = expires_at_ms,
|
|
draining = false,
|
|
}
|
|
local encoded = cjson.encode(assignment)
|
|
redis.call('HSET', owners_key, proxy_id, encoded)
|
|
redis.call('ZADD', owner_expiry_key, expires_at_ms, proxy_id)
|
|
touch(owners_key, tonumber(record.expiresAtMs))
|
|
touch(owner_expiry_key, tonumber(record.expiresAtMs))
|
|
record.ownerWorkerId = worker_id
|
|
redis.call('HSET', records_key, proxy_id, cjson.encode(record))
|
|
remove_available(proxy_id, record)
|
|
redis.call('ZADD', record.ownerIndexKey, record.usableUntilMs, proxy_id)
|
|
touch(record.ownerIndexKey, tonumber(record.expiresAtMs))
|
|
return finish({status = 'ok', record = encoded})
|
|
end
|
|
|
|
if operation == 'renew' then
|
|
cleanup_hard_expired(now_ms)
|
|
local current = decode_table(redis.call('HGET', owners_key, proxy_id))
|
|
if not valid_assignment(current) or current.workerId ~= worker_id or tonumber(current.epoch) ~= epoch then
|
|
return finish({status = 'stale'})
|
|
end
|
|
if tonumber(current.expiresAtMs) <= now_ms then
|
|
clear_owner(proxy_id, current, now_ms, true)
|
|
return finish({status = 'stale'})
|
|
end
|
|
local record = decode_table(redis.call('HGET', records_key, proxy_id))
|
|
if not record or record.ownerWorkerId ~= worker_id or
|
|
not tonumber(record.usableUntilMs) or tonumber(record.usableUntilMs) <= now_ms then
|
|
clear_owner(proxy_id, current, now_ms, false)
|
|
return finish({status = 'stale'})
|
|
end
|
|
local expires_at_ms = now_ms + value
|
|
if tonumber(record.usableUntilMs) < expires_at_ms then
|
|
expires_at_ms = tonumber(record.usableUntilMs)
|
|
end
|
|
current.assignmentVersion = tonumber(current.assignmentVersion) + 1
|
|
current.expiresAtMs = expires_at_ms
|
|
local encoded = cjson.encode(current)
|
|
redis.call('HSET', owners_key, proxy_id, encoded)
|
|
redis.call('ZADD', owner_expiry_key, expires_at_ms, proxy_id)
|
|
touch(owners_key, tonumber(record.expiresAtMs))
|
|
touch(owner_expiry_key, tonumber(record.expiresAtMs))
|
|
return finish({status = 'ok', record = encoded})
|
|
end
|
|
|
|
if operation == 'begin_drain' then
|
|
local current = decode_table(redis.call('HGET', owners_key, proxy_id))
|
|
if not valid_assignment(current) or current.workerId ~= worker_id or tonumber(current.epoch) ~= epoch then
|
|
return finish({status = 'stale'})
|
|
end
|
|
if not current.draining then
|
|
current.draining = true
|
|
current.assignmentVersion = tonumber(current.assignmentVersion) + 1
|
|
local encoded = cjson.encode(current)
|
|
redis.call('HSET', owners_key, proxy_id, encoded)
|
|
local record = decode_table(redis.call('HGET', records_key, proxy_id))
|
|
remove_owned(proxy_id, record)
|
|
return finish({status = 'ok', record = encoded})
|
|
end
|
|
return finish({status = 'ok', record = cjson.encode(current)})
|
|
end
|
|
|
|
if operation == 'acknowledge_drain' then
|
|
local current = decode_table(redis.call('HGET', owners_key, proxy_id))
|
|
if not valid_assignment(current) or current.workerId ~= worker_id or tonumber(current.epoch) ~= epoch then
|
|
return finish({status = 'stale'})
|
|
end
|
|
if not current.draining then
|
|
return finish({status = 'not_draining'})
|
|
end
|
|
if active > 0 or reserved > 0 then
|
|
return finish({status = 'drain_not_ready'})
|
|
end
|
|
local server_time = redis.call('TIME')
|
|
local server_now_ms = tonumber(server_time[1]) * 1000 + math.floor(tonumber(server_time[2]) / 1000)
|
|
clear_owner(proxy_id, current, server_now_ms, true)
|
|
return finish({status = 'ok'})
|
|
end
|
|
|
|
if operation == 'get' then
|
|
local raw = redis.call('HGET', owners_key, proxy_id)
|
|
if not raw then
|
|
return finish({status = 'not_found'})
|
|
end
|
|
return finish({status = 'ok', record = raw})
|
|
end
|
|
|
|
if operation == 'expire' then
|
|
local ids = redis.call('ZRANGEBYSCORE', owner_expiry_key, '-inf', now_ms, 'LIMIT', 0, value)
|
|
local expired = {}
|
|
for _, id in ipairs(ids) do
|
|
local current = decode_table(redis.call('HGET', owners_key, id))
|
|
if valid_assignment(current) then
|
|
expired[#expired + 1] = current
|
|
end
|
|
clear_owner(id, current, now_ms, true)
|
|
end
|
|
local encoded = '[]'
|
|
if #expired > 0 then
|
|
encoded = cjson.encode(expired)
|
|
end
|
|
return finish({status = 'ok', record = encoded})
|
|
end
|
|
|
|
return finish({status = 'invalid'})
|