230 lines
9.9 KiB
Lua
230 lines
9.9 KiB
Lua
local sessions_key = KEYS[1]
|
|
local session_expiry_key = KEYS[2]
|
|
local runtime_key = KEYS[3]
|
|
local runtime_expiry_key = KEYS[4]
|
|
local owners_key = KEYS[5]
|
|
|
|
local operation = ARGV[1]
|
|
local ttl_ms = tonumber(ARGV[2])
|
|
local cleanup_limit = tonumber(ARGV[3])
|
|
local payload = ARGV[4]
|
|
local digest = ARGV[5]
|
|
|
|
local function reply(status, snapshots)
|
|
if snapshots then
|
|
return cjson.encode({status = status, snapshots = snapshots})
|
|
end
|
|
return '{"status":' .. cjson.encode(status) .. ',"snapshots":[]}'
|
|
end
|
|
|
|
local function now_ms()
|
|
local value = redis.call('TIME')
|
|
return tonumber(value[1]) * 1000 + math.floor(tonumber(value[2]) / 1000)
|
|
end
|
|
|
|
local function decode_table(value)
|
|
if not value then
|
|
return nil
|
|
end
|
|
local ok, decoded = pcall(cjson.decode, value)
|
|
if not ok or type(decoded) ~= 'table' then
|
|
return nil
|
|
end
|
|
return decoded
|
|
end
|
|
|
|
local function valid_uint(value)
|
|
return type(value) == 'string' and string.match(value, '^[0-9]+$') and
|
|
value ~= '0' and (string.len(value) == 1 or string.sub(value, 1, 1) ~= '0')
|
|
end
|
|
|
|
local function compare_uint(left, right)
|
|
if string.len(left) ~= string.len(right) then
|
|
return string.len(left) < string.len(right) and -1 or 1
|
|
end
|
|
if left == right then
|
|
return 0
|
|
end
|
|
return left < right and -1 or 1
|
|
end
|
|
|
|
local function cleanup(now)
|
|
local expired_sessions = redis.call('ZRANGEBYSCORE', session_expiry_key, '-inf', now, 'LIMIT', 0, cleanup_limit)
|
|
for _, worker_id in ipairs(expired_sessions) do
|
|
redis.call('HDEL', sessions_key, worker_id)
|
|
redis.call('HDEL', runtime_key, worker_id)
|
|
redis.call('ZREM', session_expiry_key, worker_id)
|
|
redis.call('ZREM', runtime_expiry_key, worker_id)
|
|
end
|
|
local expired_reports = redis.call('ZRANGEBYSCORE', runtime_expiry_key, '-inf', now, 'LIMIT', 0, cleanup_limit)
|
|
for _, worker_id in ipairs(expired_reports) do
|
|
redis.call('HDEL', runtime_key, worker_id)
|
|
redis.call('ZREM', runtime_expiry_key, worker_id)
|
|
end
|
|
end
|
|
|
|
local function valid_session(value)
|
|
return value and value.version == 1 and type(value.workerId) == 'string' and value.workerId ~= '' and
|
|
type(value.instanceId) == 'string' and value.instanceId ~= '' and
|
|
type(value.sessionId) == 'string' and value.sessionId ~= '' and
|
|
valid_uint(value.ackedSnapshotVersion) and valid_uint(value.ackedOwnershipEpoch)
|
|
end
|
|
|
|
local function valid_owner(value, worker_id, ownership_epoch, now)
|
|
return value and type(value.workerId) == 'string' and value.workerId == worker_id and
|
|
type(value.epoch) == 'number' and valid_uint(tostring(value.epoch)) and
|
|
compare_uint(ownership_epoch, tostring(value.epoch)) >= 0 and
|
|
type(value.expiresAtMs) == 'number' and value.expiresAtMs > now
|
|
end
|
|
|
|
local now = now_ms()
|
|
cleanup(now)
|
|
|
|
if operation == 'replace_session' then
|
|
if not ttl_ms or ttl_ms <= 0 then
|
|
return reply('invalid')
|
|
end
|
|
local session = decode_table(payload)
|
|
if not valid_session(session) then
|
|
return reply('invalid')
|
|
end
|
|
local current = decode_table(redis.call('HGET', sessions_key, session.workerId))
|
|
if current and valid_session(current) and current.sessionId == session.sessionId and
|
|
current.instanceId == session.instanceId and type(current.expiresAtMs) == 'number' and
|
|
current.expiresAtMs > now then
|
|
local epoch_order = compare_uint(session.ackedOwnershipEpoch, current.ackedOwnershipEpoch)
|
|
local version_order = compare_uint(session.ackedSnapshotVersion, current.ackedSnapshotVersion)
|
|
if epoch_order < 0 or (epoch_order == 0 and version_order < 0) then
|
|
return reply('stale')
|
|
end
|
|
if epoch_order > 0 or version_order > 0 then
|
|
redis.call('HDEL', runtime_key, session.workerId)
|
|
redis.call('ZREM', runtime_expiry_key, session.workerId)
|
|
end
|
|
else
|
|
redis.call('HDEL', runtime_key, session.workerId)
|
|
redis.call('ZREM', runtime_expiry_key, session.workerId)
|
|
end
|
|
session.expiresAtMs = now + ttl_ms
|
|
redis.call('HSET', sessions_key, session.workerId, cjson.encode(session))
|
|
redis.call('ZADD', session_expiry_key, session.expiresAtMs, session.workerId)
|
|
return reply('ok')
|
|
end
|
|
|
|
if operation == 'replace_report' then
|
|
if not ttl_ms or ttl_ms <= 0 or type(digest) ~= 'string' or digest == '' then
|
|
return reply('invalid')
|
|
end
|
|
local report = decode_table(payload)
|
|
if not report or report.version ~= 1 or type(report.workerId) ~= 'string' or report.workerId == '' or
|
|
type(report.sessionId) ~= 'string' or report.sessionId == '' or not valid_uint(report.sequence) or
|
|
not valid_uint(report.snapshotVersion) or not valid_uint(report.ownershipEpoch) or
|
|
type(report.observedAtMs) ~= 'number' or type(report.counters) ~= 'table' then
|
|
return reply('invalid')
|
|
end
|
|
local session = decode_table(redis.call('HGET', sessions_key, report.workerId))
|
|
if not valid_session(session) or session.workerId ~= report.workerId or session.sessionId ~= report.sessionId or
|
|
type(session.expiresAtMs) ~= 'number' or session.expiresAtMs <= now then
|
|
return reply('unavailable')
|
|
end
|
|
if report.snapshotVersion ~= session.ackedSnapshotVersion or
|
|
report.ownershipEpoch ~= session.ackedOwnershipEpoch then
|
|
return reply('stale')
|
|
end
|
|
local current = decode_table(redis.call('HGET', runtime_key, report.workerId))
|
|
if current and current.sessionId == report.sessionId and valid_uint(current.sequence) then
|
|
local ordering = compare_uint(report.sequence, current.sequence)
|
|
if ordering < 0 then
|
|
return reply('stale')
|
|
end
|
|
if ordering == 0 then
|
|
if current.digest == digest then
|
|
return reply('ok')
|
|
end
|
|
return reply('conflict')
|
|
end
|
|
end
|
|
local seen = {}
|
|
for _, counter in pairs(report.counters) do
|
|
if type(counter) ~= 'table' or type(counter.proxyId) ~= 'string' or counter.proxyId == '' or
|
|
type(counter.active) ~= 'number' or counter.active < 0 or counter.active ~= math.floor(counter.active) or
|
|
type(counter.reserved) ~= 'number' or counter.reserved < 0 or counter.reserved ~= math.floor(counter.reserved) or
|
|
type(counter.draining) ~= 'boolean' or seen[counter.proxyId] then
|
|
return reply('invalid')
|
|
end
|
|
seen[counter.proxyId] = true
|
|
local owner = decode_table(redis.call('HGET', owners_key, counter.proxyId))
|
|
if not valid_owner(owner, report.workerId, report.ownershipEpoch, now) then
|
|
return reply('stale')
|
|
end
|
|
end
|
|
report.digest = digest
|
|
report.expiresAtMs = now + ttl_ms
|
|
redis.call('HSET', runtime_key, report.workerId, cjson.encode(report))
|
|
redis.call('ZADD', runtime_expiry_key, report.expiresAtMs, report.workerId)
|
|
session.expiresAtMs = now + ttl_ms
|
|
redis.call('HSET', sessions_key, report.workerId, cjson.encode(session))
|
|
redis.call('ZADD', session_expiry_key, session.expiresAtMs, report.workerId)
|
|
return reply('ok')
|
|
end
|
|
|
|
if operation == 'read' then
|
|
local queries = decode_table(payload)
|
|
if not queries then
|
|
return reply('invalid')
|
|
end
|
|
if next(queries) == nil then
|
|
return '{"status":"ok","snapshots":[]}'
|
|
end
|
|
local snapshots = cjson.decode('[]')
|
|
local cache = {}
|
|
for _, query in ipairs(queries) do
|
|
if type(query) ~= 'table' or type(query.proxyId) ~= 'string' or query.proxyId == '' or
|
|
type(query.workerId) ~= 'string' or query.workerId == '' or not valid_uint(query.ownershipEpoch) then
|
|
return reply('invalid')
|
|
end
|
|
local snapshot = {proxyId = query.proxyId, active = 0, reserved = 0, draining = false, fresh = false}
|
|
local owner = decode_table(redis.call('HGET', owners_key, query.proxyId))
|
|
if valid_owner(owner, query.workerId, query.ownershipEpoch, now) and
|
|
compare_uint(query.ownershipEpoch, tostring(owner.epoch)) == 0 then
|
|
local cached = cache[query.workerId]
|
|
if not cached then
|
|
local session = decode_table(redis.call('HGET', sessions_key, query.workerId))
|
|
local report = decode_table(redis.call('HGET', runtime_key, query.workerId))
|
|
cached = {fresh = false, counters = {}}
|
|
if valid_session(session) and session.workerId == query.workerId and report and
|
|
report.workerId == query.workerId and report.sessionId == session.sessionId and
|
|
type(session.expiresAtMs) == 'number' and session.expiresAtMs > now and
|
|
type(report.expiresAtMs) == 'number' and report.expiresAtMs > now and
|
|
valid_uint(report.ownershipEpoch) and
|
|
report.snapshotVersion == session.ackedSnapshotVersion and
|
|
report.ownershipEpoch == session.ackedOwnershipEpoch then
|
|
cached.fresh = true
|
|
cached.ownershipEpoch = report.ownershipEpoch
|
|
if type(report.counters) == 'table' then
|
|
for _, counter in pairs(report.counters) do
|
|
if type(counter) == 'table' and type(counter.proxyId) == 'string' then
|
|
cached.counters[counter.proxyId] = counter
|
|
end
|
|
end
|
|
end
|
|
end
|
|
cache[query.workerId] = cached
|
|
end
|
|
if cached.fresh and compare_uint(cached.ownershipEpoch, query.ownershipEpoch) >= 0 then
|
|
snapshot.fresh = true
|
|
local counter = cached.counters[query.proxyId]
|
|
if counter then
|
|
snapshot.active = counter.active
|
|
snapshot.reserved = counter.reserved
|
|
snapshot.draining = counter.draining
|
|
end
|
|
end
|
|
end
|
|
snapshots[#snapshots + 1] = snapshot
|
|
end
|
|
return reply('ok', snapshots)
|
|
end
|
|
|
|
return reply('invalid')
|