178 lines
7.1 KiB
Lua
178 lines
7.1 KiB
Lua
local records_key = KEYS[1]
|
|
local inventory_key = KEYS[2]
|
|
local available_upstream_key = KEYS[3]
|
|
local owned_upstream_key = KEYS[4]
|
|
local owners_key = KEYS[5]
|
|
local sessions_key = KEYS[6]
|
|
local session_expiry_key = KEYS[7]
|
|
local runtime_key = KEYS[8]
|
|
local runtime_expiry_key = KEYS[9]
|
|
|
|
local upstream_id = ARGV[1]
|
|
local safety_margin_ms = tonumber(ARGV[2])
|
|
local scan_limit = tonumber(ARGV[3])
|
|
local cleanup_limit = tonumber(ARGV[4])
|
|
|
|
local function reply(status, managed, available_slots)
|
|
return cjson.encode({
|
|
status = status,
|
|
managed = managed or 0,
|
|
availableSlots = tostring(available_slots or 0)
|
|
})
|
|
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
|
|
|
|
if type(upstream_id) ~= 'string' or upstream_id == '' or not safety_margin_ms or safety_margin_ms < 0 or
|
|
not scan_limit or scan_limit <= 0 or not cleanup_limit or cleanup_limit <= 0 then
|
|
return reply('invalid', 0, 0)
|
|
end
|
|
|
|
local now = now_ms()
|
|
cleanup(now)
|
|
local threshold = now + safety_margin_ms
|
|
local available_ids = redis.call('ZRANGEBYSCORE', available_upstream_key, '(' .. threshold, '+inf',
|
|
'LIMIT', 0, scan_limit + 1)
|
|
if #available_ids > scan_limit then
|
|
return reply('unavailable', 0, 0)
|
|
end
|
|
local remaining = scan_limit - #available_ids
|
|
local owned_ids = redis.call('ZRANGEBYSCORE', owned_upstream_key, '(' .. threshold, '+inf',
|
|
'LIMIT', 0, remaining + 1)
|
|
if #owned_ids > remaining then
|
|
return reply('unavailable', 0, 0)
|
|
end
|
|
|
|
local managed = tonumber(redis.call('HGET', inventory_key, upstream_id) or '0')
|
|
if not managed or managed < 0 or managed ~= math.floor(managed) then
|
|
return reply('unavailable', 0, 0)
|
|
end
|
|
local available_slots = 0
|
|
local worker_cache = {}
|
|
local seen = {}
|
|
for _, proxy_id in ipairs(available_ids) do
|
|
seen[proxy_id] = true
|
|
local record = decode_table(redis.call('HGET', records_key, proxy_id))
|
|
if not record or type(record.sourceUpstream) ~= 'string' or type(record.state) ~= 'string' or
|
|
record.sourceUpstream ~= upstream_id or record.state ~= 'AVAILABLE' or
|
|
type(record.usableUntilMs) ~= 'number' or record.usableUntilMs <= threshold or
|
|
type(record.maxConcurrency) ~= 'number' or record.maxConcurrency < 0 or
|
|
record.maxConcurrency ~= math.floor(record.maxConcurrency) or
|
|
(record.ownerWorkerId and record.ownerWorkerId ~= '') or redis.call('HGET', owners_key, proxy_id) then
|
|
return reply('unavailable', 0, 0)
|
|
end
|
|
available_slots = available_slots + record.maxConcurrency
|
|
end
|
|
|
|
for _, proxy_id in ipairs(owned_ids) do
|
|
if seen[proxy_id] then
|
|
return reply('unavailable', 0, 0)
|
|
end
|
|
local record = decode_table(redis.call('HGET', records_key, proxy_id))
|
|
if not record or record.sourceUpstream ~= upstream_id or record.state ~= 'AVAILABLE' or
|
|
type(record.usableUntilMs) ~= 'number' or record.usableUntilMs <= threshold or
|
|
type(record.maxConcurrency) ~= 'number' or record.maxConcurrency < 0 or
|
|
record.maxConcurrency ~= math.floor(record.maxConcurrency) or
|
|
type(record.ownerWorkerId) ~= 'string' or record.ownerWorkerId == '' then
|
|
return reply('unavailable', 0, 0)
|
|
end
|
|
local owner_worker_id = record.ownerWorkerId
|
|
local owner = decode_table(redis.call('HGET', owners_key, proxy_id))
|
|
if not owner or owner.workerId ~= owner_worker_id or type(owner.epoch) ~= 'number' or
|
|
type(owner.expiresAtMs) ~= 'number' or owner.expiresAtMs <= now or
|
|
type(owner.draining) ~= 'boolean' then
|
|
return reply('unavailable', 0, 0)
|
|
end
|
|
local cached = worker_cache[owner_worker_id]
|
|
if not cached then
|
|
local session = decode_table(redis.call('HGET', sessions_key, owner_worker_id))
|
|
local report = decode_table(redis.call('HGET', runtime_key, owner_worker_id))
|
|
cached = {fresh = false, counters = {}}
|
|
if session and report and session.workerId == owner_worker_id and
|
|
report.workerId == owner_worker_id and session.sessionId == report.sessionId and
|
|
type(session.expiresAtMs) == 'number' and session.expiresAtMs > now and
|
|
type(report.expiresAtMs) == 'number' and report.expiresAtMs > now and
|
|
valid_uint(session.ackedSnapshotVersion) and valid_uint(session.ackedOwnershipEpoch) 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
|
|
worker_cache[owner_worker_id] = cached
|
|
end
|
|
local owner_epoch = tostring(owner.epoch)
|
|
if cached.fresh and valid_uint(owner_epoch) and
|
|
compare_uint(cached.ownershipEpoch, owner_epoch) >= 0 then
|
|
local counter = cached.counters[proxy_id]
|
|
local active = 0
|
|
local reserved = 0
|
|
local draining = false
|
|
if counter then
|
|
active = counter.active
|
|
reserved = counter.reserved
|
|
draining = counter.draining
|
|
end
|
|
if type(active) == 'number' and type(reserved) == 'number' and active >= 0 and reserved >= 0 and
|
|
active == math.floor(active) and reserved == math.floor(reserved) and
|
|
not draining and not owner.draining then
|
|
local slots = record.maxConcurrency - active - reserved
|
|
if slots > 0 then
|
|
available_slots = available_slots + slots
|
|
end
|
|
end
|
|
end
|
|
end
|
|
|
|
return reply('ok', managed, available_slots)
|