439 lines
20 KiB
Lua
439 lines
20 KiB
Lua
local sessions_key = KEYS[1]
|
|
local session_expiry_key = KEYS[2]
|
|
local snapshots_key = KEYS[3]
|
|
local snapshot_expiry_key = KEYS[4]
|
|
local runtime_key = KEYS[5]
|
|
local runtime_expiry_key = KEYS[6]
|
|
local owners_key = KEYS[7]
|
|
local epoch_key = KEYS[8]
|
|
|
|
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, record)
|
|
if not snapshots then
|
|
local suffix = ''
|
|
if record then
|
|
suffix = ',"record":' .. cjson.encode(record)
|
|
end
|
|
return '{"status":' .. cjson.encode(status) .. ',"snapshots":[]' .. suffix .. '}'
|
|
end
|
|
return cjson.encode({status = status, snapshots = snapshots, record = record})
|
|
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
|
|
local expired_snapshots = redis.call('ZRANGEBYSCORE', snapshot_expiry_key, '-inf', now, 'LIMIT', 0, cleanup_limit)
|
|
for _, worker_id in ipairs(expired_snapshots) do
|
|
redis.call('HDEL', snapshots_key, worker_id)
|
|
redis.call('ZREM', snapshot_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
|
|
type(value.ackedSnapshotVersion) == 'string' and type(value.ackedOwnershipEpoch) == 'string' and
|
|
type(value.ackedChecksum) == 'string' and type(value.runtimeEnabled) == 'boolean'
|
|
end
|
|
|
|
local function valid_legacy_session(value)
|
|
return valid_session(value) and valid_uint(value.ackedSnapshotVersion) and valid_uint(value.ackedOwnershipEpoch)
|
|
end
|
|
|
|
local function valid_control_session(value)
|
|
if not valid_session(value) or type(value.zone) ~= 'string' or value.zone == '' or
|
|
type(value.protocolVersion) ~= 'number' or value.protocolVersion <= 0 or type(value.labels) ~= 'table' then
|
|
return false
|
|
end
|
|
if value.ackedSnapshotVersion == '0' and value.ackedOwnershipEpoch == '0' and value.ackedChecksum == '' then
|
|
return true
|
|
end
|
|
return valid_uint(value.ackedSnapshotVersion) and valid_uint(value.ackedOwnershipEpoch) and
|
|
string.len(value.ackedChecksum) == 64 and string.match(value.ackedChecksum, '^[0-9a-f]+$') ~= nil
|
|
end
|
|
|
|
local function valid_reference(value)
|
|
return value and value.version == 1 and type(value.workerId) == 'string' and value.workerId ~= '' and
|
|
valid_uint(value.snapshotVersion) and valid_uint(value.ownershipEpoch) and
|
|
type(value.checksum) == 'string' and string.len(value.checksum) == 64 and
|
|
string.match(value.checksum, '^[0-9a-f]+$') ~= nil
|
|
end
|
|
|
|
local function compare_reference(left, right)
|
|
local epoch_order = compare_uint(left.ownershipEpoch, right.ownershipEpoch)
|
|
if epoch_order ~= 0 then
|
|
return epoch_order
|
|
end
|
|
return compare_uint(left.snapshotVersion, right.snapshotVersion)
|
|
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 == 'current_epoch' then
|
|
local epoch = redis.call('GET', epoch_key)
|
|
if not epoch then
|
|
epoch = '1'
|
|
redis.call('SET', epoch_key, epoch)
|
|
end
|
|
redis.call('PERSIST', epoch_key)
|
|
if not valid_uint(epoch) then
|
|
return reply('invalid')
|
|
end
|
|
return reply('ok', nil, epoch)
|
|
end
|
|
|
|
if operation == 'open_session' then
|
|
if not ttl_ms or ttl_ms <= 0 then
|
|
return reply('invalid')
|
|
end
|
|
local session = decode_table(payload)
|
|
if not valid_control_session(session) or session.ackedSnapshotVersion ~= '0' or
|
|
session.ackedOwnershipEpoch ~= '0' or session.ackedChecksum ~= '' or session.runtimeEnabled then
|
|
return reply('invalid')
|
|
end
|
|
redis.call('HDEL', runtime_key, session.workerId)
|
|
redis.call('ZREM', runtime_expiry_key, session.workerId)
|
|
redis.call('HDEL', snapshots_key, session.workerId)
|
|
redis.call('ZREM', snapshot_expiry_key, session.workerId)
|
|
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 == 'validate_session' then
|
|
local validation = decode_table(payload)
|
|
if not validation or validation.version ~= 1 or type(validation.workerId) ~= 'string' or
|
|
validation.workerId == '' or type(validation.sessionId) ~= 'string' or validation.sessionId == '' then
|
|
return reply('invalid')
|
|
end
|
|
local session = decode_table(redis.call('HGET', sessions_key, validation.workerId))
|
|
if not valid_control_session(session) or session.workerId ~= validation.workerId or
|
|
session.sessionId ~= validation.sessionId or type(session.expiresAtMs) ~= 'number' or
|
|
session.expiresAtMs <= now then
|
|
return reply('unavailable')
|
|
end
|
|
return reply('ok')
|
|
end
|
|
|
|
if operation == 'record_snapshot' then
|
|
if not ttl_ms or ttl_ms <= 0 then
|
|
return reply('invalid')
|
|
end
|
|
local issue = decode_table(payload)
|
|
if not issue or issue.version ~= 1 or type(issue.sessionId) ~= 'string' or issue.sessionId == '' or
|
|
not valid_reference(issue.reference) then
|
|
return reply('invalid')
|
|
end
|
|
local reference = issue.reference
|
|
local session = decode_table(redis.call('HGET', sessions_key, reference.workerId))
|
|
if not valid_control_session(session) or session.workerId ~= reference.workerId or session.sessionId ~= issue.sessionId or
|
|
type(session.expiresAtMs) ~= 'number' or session.expiresAtMs <= now then
|
|
return reply('unavailable')
|
|
end
|
|
local epoch = redis.call('GET', epoch_key)
|
|
if not epoch then
|
|
epoch = '1'
|
|
redis.call('SET', epoch_key, epoch)
|
|
end
|
|
redis.call('PERSIST', epoch_key)
|
|
if not valid_uint(epoch) or compare_uint(reference.ownershipEpoch, epoch) ~= 0 then
|
|
return reply('snapshot_mismatch')
|
|
end
|
|
local current = decode_table(redis.call('HGET', snapshots_key, reference.workerId))
|
|
if current and type(current.expiresAtMs) == 'number' and current.expiresAtMs > now and valid_reference(current) then
|
|
local ordering = compare_reference(reference, current)
|
|
if ordering < 0 then
|
|
return reply('stale')
|
|
end
|
|
if ordering == 0 and reference.checksum ~= current.checksum then
|
|
return reply('conflict')
|
|
end
|
|
end
|
|
reference.expiresAtMs = now + ttl_ms
|
|
redis.call('HSET', snapshots_key, reference.workerId, cjson.encode(reference))
|
|
redis.call('ZADD', snapshot_expiry_key, reference.expiresAtMs, reference.workerId)
|
|
return reply('ok')
|
|
end
|
|
|
|
if operation == 'acknowledge_snapshot' then
|
|
if not ttl_ms or ttl_ms <= 0 then
|
|
return reply('invalid')
|
|
end
|
|
local acknowledgement = decode_table(payload)
|
|
if not acknowledgement or acknowledgement.version ~= 1 or type(acknowledgement.workerId) ~= 'string' or
|
|
acknowledgement.workerId == '' or type(acknowledgement.sessionId) ~= 'string' or acknowledgement.sessionId == '' or
|
|
type(acknowledgement.applied) ~= 'boolean' or type(acknowledgement.errorCode) ~= 'string' or
|
|
not valid_reference(acknowledgement.reference) or acknowledgement.reference.workerId ~= acknowledgement.workerId then
|
|
return reply('invalid')
|
|
end
|
|
local session = decode_table(redis.call('HGET', sessions_key, acknowledgement.workerId))
|
|
if not valid_control_session(session) or session.sessionId ~= acknowledgement.sessionId or
|
|
type(session.expiresAtMs) ~= 'number' or session.expiresAtMs <= now then
|
|
return reply('unavailable')
|
|
end
|
|
if session.ackedSnapshotVersion ~= '0' then
|
|
local previous = {ownershipEpoch = session.ackedOwnershipEpoch, snapshotVersion = session.ackedSnapshotVersion}
|
|
local acknowledged = compare_reference(acknowledgement.reference, previous)
|
|
if acknowledged < 0 then
|
|
return reply('stale_acknowledgement')
|
|
end
|
|
if acknowledged == 0 and acknowledgement.reference.checksum ~= session.ackedChecksum then
|
|
return reply('snapshot_mismatch')
|
|
end
|
|
end
|
|
local current = decode_table(redis.call('HGET', snapshots_key, acknowledgement.workerId))
|
|
if not valid_reference(current) or type(current.expiresAtMs) ~= 'number' or current.expiresAtMs <= now then
|
|
return reply('snapshot_mismatch')
|
|
end
|
|
local ordering = compare_reference(acknowledgement.reference, current)
|
|
if ordering < 0 then
|
|
return reply('stale_acknowledgement')
|
|
end
|
|
if ordering > 0 or acknowledgement.reference.checksum ~= current.checksum then
|
|
return reply('snapshot_mismatch')
|
|
end
|
|
if not acknowledgement.applied then
|
|
redis.call('HDEL', runtime_key, acknowledgement.workerId)
|
|
redis.call('ZREM', runtime_expiry_key, acknowledgement.workerId)
|
|
session.runtimeEnabled = false
|
|
session.expiresAtMs = now + ttl_ms
|
|
redis.call('HSET', sessions_key, acknowledgement.workerId, cjson.encode(session))
|
|
redis.call('ZADD', session_expiry_key, session.expiresAtMs, acknowledgement.workerId)
|
|
return reply('ok')
|
|
end
|
|
if session.ackedSnapshotVersion ~= '0' and
|
|
compare_reference(acknowledgement.reference, {ownershipEpoch = session.ackedOwnershipEpoch, snapshotVersion = session.ackedSnapshotVersion}) == 0 then
|
|
if not session.runtimeEnabled then
|
|
redis.call('HDEL', runtime_key, acknowledgement.workerId)
|
|
redis.call('ZREM', runtime_expiry_key, acknowledgement.workerId)
|
|
session.runtimeEnabled = true
|
|
end
|
|
session.expiresAtMs = now + ttl_ms
|
|
redis.call('HSET', sessions_key, acknowledgement.workerId, cjson.encode(session))
|
|
redis.call('ZADD', session_expiry_key, session.expiresAtMs, acknowledgement.workerId)
|
|
return reply('ok')
|
|
end
|
|
redis.call('HDEL', runtime_key, acknowledgement.workerId)
|
|
redis.call('ZREM', runtime_expiry_key, acknowledgement.workerId)
|
|
session.ackedSnapshotVersion = acknowledgement.reference.snapshotVersion
|
|
session.ackedOwnershipEpoch = acknowledgement.reference.ownershipEpoch
|
|
session.ackedChecksum = acknowledgement.reference.checksum
|
|
session.runtimeEnabled = true
|
|
session.expiresAtMs = now + ttl_ms
|
|
redis.call('HSET', sessions_key, acknowledgement.workerId, cjson.encode(session))
|
|
redis.call('ZADD', session_expiry_key, session.expiresAtMs, acknowledgement.workerId)
|
|
return reply('ok')
|
|
end
|
|
|
|
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_legacy_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 not valid_uint(session.ackedSnapshotVersion) or not valid_uint(session.ackedOwnershipEpoch) or
|
|
not session.runtimeEnabled or report.snapshotVersion ~= session.ackedSnapshotVersion or
|
|
report.ownershipEpoch ~= session.ackedOwnershipEpoch then
|
|
return reply('snapshot_mismatch')
|
|
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
|
|
current.expiresAtMs = now + ttl_ms
|
|
redis.call('HSET', runtime_key, report.workerId, cjson.encode(current))
|
|
redis.call('ZADD', runtime_expiry_key, current.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
|
|
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 valid_uint(session.ackedSnapshotVersion) and
|
|
valid_uint(session.ackedOwnershipEpoch) and session.runtimeEnabled 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')
|