473 lines
19 KiB
Lua
473 lines
19 KiB
Lua
local records_key = KEYS[1]
|
|
local expiry_key = KEYS[2]
|
|
local due_key = KEYS[3]
|
|
local queued_key = KEYS[4]
|
|
local leases_key = KEYS[5]
|
|
local tasks_key = KEYS[6]
|
|
local task_expiry_key = KEYS[7]
|
|
local ref_task_key = KEYS[8]
|
|
local checker_leases_key = KEYS[9]
|
|
local state_inventory_key = KEYS[10]
|
|
local egress_due_key = KEYS[11]
|
|
local target_due_key = KEYS[12]
|
|
|
|
local operation = ARGV[1]
|
|
local now_ms = tonumber(ARGV[2])
|
|
local decoded, payload = pcall(cjson.decode, ARGV[3])
|
|
if not decoded or type(payload) ~= 'table' or not now_ms then
|
|
return cjson.encode({status = 'invalid'})
|
|
end
|
|
|
|
local limit = tonumber(payload.limit or 0)
|
|
if not limit or limit <= 0 then
|
|
return cjson.encode({status = 'invalid'})
|
|
end
|
|
|
|
local function finish(reply)
|
|
return cjson.encode(reply)
|
|
end
|
|
|
|
local function live_record(proxy_id)
|
|
local raw = redis.call('HGET', records_key, proxy_id)
|
|
if not raw then
|
|
return nil, nil
|
|
end
|
|
local valid, record = pcall(cjson.decode, raw)
|
|
if not valid or type(record) ~= 'table' or tonumber(record.expiresAtMs) <= now_ms or record.state == 'EXTRACTED' then
|
|
return nil, nil
|
|
end
|
|
return raw, record
|
|
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 due_key_for_level(level)
|
|
if level == 'EGRESS' then
|
|
return egress_due_key
|
|
end
|
|
if level == 'TARGET' then
|
|
return target_due_key
|
|
end
|
|
return due_key
|
|
end
|
|
|
|
local function task_reference(level, proxy_id, routing_name, target_url)
|
|
if level == 'EGRESS' then
|
|
return proxy_id .. '\0EGRESS\0\0' .. target_url
|
|
end
|
|
if level == 'TARGET' then
|
|
return proxy_id .. '\0TARGET\0' .. routing_name .. '\0' .. target_url
|
|
end
|
|
return proxy_id
|
|
end
|
|
|
|
local function remove_task(task_id, requeue)
|
|
local raw = redis.call('HGET', tasks_key, task_id)
|
|
local task = nil
|
|
if raw then
|
|
local valid
|
|
valid, task = pcall(cjson.decode, raw)
|
|
if not valid or type(task) ~= 'table' then
|
|
task = nil
|
|
end
|
|
end
|
|
redis.call('ZREM', queued_key, task_id)
|
|
redis.call('ZREM', leases_key, task_id)
|
|
redis.call('ZREM', task_expiry_key, task_id)
|
|
redis.call('HDEL', tasks_key, task_id)
|
|
if task then
|
|
if type(task.upstreamTasksKey) == 'string' and task.upstreamTasksKey ~= '' then
|
|
redis.call('ZREM', task.upstreamTasksKey, task_id)
|
|
end
|
|
if type(task.checkerLeaseKey) == 'string' and task.checkerLeaseKey ~= '' then
|
|
redis.call('ZREM', task.checkerLeaseKey, task_id)
|
|
end
|
|
local reference = task.reference or task.proxyId
|
|
if type(reference) == 'string' and redis.call('HGET', ref_task_key, reference) == task_id then
|
|
redis.call('HDEL', ref_task_key, reference)
|
|
end
|
|
local due_member = task.dueMember or task.proxyId
|
|
if requeue and type(task.proxyId) == 'string' and type(due_member) == 'string' then
|
|
local _, record = live_record(task.proxyId)
|
|
if record then
|
|
local task_due_key = due_key_for_level(task.level)
|
|
redis.call('ZADD', task_due_key, now_ms, due_member)
|
|
touch(task_due_key, tonumber(record.expiresAtMs))
|
|
end
|
|
end
|
|
end
|
|
end
|
|
|
|
local function reap()
|
|
local expired_leases = redis.call('ZRANGEBYSCORE', leases_key, '-inf', now_ms, 'LIMIT', 0, limit)
|
|
for _, task_id in ipairs(expired_leases) do
|
|
local raw = redis.call('HGET', tasks_key, task_id)
|
|
if not raw then
|
|
redis.call('ZREM', leases_key, task_id)
|
|
else
|
|
local valid, task = pcall(cjson.decode, raw)
|
|
if not valid or type(task) ~= 'table' then
|
|
remove_task(task_id, true)
|
|
elseif task.state == 'LEASED' and tonumber(task.leaseExpiresAtMs or 0) <= now_ms then
|
|
redis.call('ZREM', leases_key, task_id)
|
|
if type(task.checkerLeaseKey) == 'string' and task.checkerLeaseKey ~= '' then
|
|
redis.call('ZREM', task.checkerLeaseKey, task_id)
|
|
end
|
|
task.state = 'QUEUED'
|
|
task.leaseCheckerId = ''
|
|
task.leaseToken = ''
|
|
task.leaseExpiresAtMs = 0
|
|
task.checkerLeaseKey = ''
|
|
redis.call('HSET', tasks_key, task_id, cjson.encode(task))
|
|
redis.call('ZADD', queued_key, tonumber(task.priority), task_id)
|
|
else
|
|
redis.call('ZREM', leases_key, task_id)
|
|
end
|
|
end
|
|
end
|
|
local expired_tasks = redis.call('ZRANGEBYSCORE', task_expiry_key, '-inf', now_ms, 'LIMIT', 0, limit)
|
|
for _, task_id in ipairs(expired_tasks) do
|
|
local raw = redis.call('HGET', tasks_key, task_id)
|
|
local requeue = true
|
|
if raw then
|
|
local valid, task = pcall(cjson.decode, raw)
|
|
if valid and type(task) == 'table' and task.state == 'DONE' then
|
|
requeue = false
|
|
end
|
|
end
|
|
remove_task(task_id, requeue)
|
|
end
|
|
end
|
|
|
|
local function valid_state(state)
|
|
return state == 'FETCHED' or state == 'CHECKING' or state == 'AVAILABLE' or
|
|
state == 'SUSPECT' or state == 'UNHEALTHY'
|
|
end
|
|
|
|
local function state_field(upstream, state)
|
|
return string.len(upstream) .. ':' .. upstream .. ':' .. state
|
|
end
|
|
|
|
local function decrement_state(upstream, state)
|
|
local field = state_field(upstream, state)
|
|
local value = redis.call('HINCRBY', state_inventory_key, field, -1)
|
|
if value <= 0 then
|
|
redis.call('HDEL', state_inventory_key, field)
|
|
end
|
|
end
|
|
|
|
local function increment_state(upstream, state)
|
|
redis.call('HINCRBY', state_inventory_key, state_field(upstream, state), 1)
|
|
end
|
|
|
|
local function matches_task(task, fact)
|
|
if not task or task.taskId ~= fact.taskId or task.proxyId ~= fact.proxyId or task.level ~= fact.level then
|
|
return false
|
|
end
|
|
if task.level == 'EGRESS' then
|
|
return (fact.routingName or '') == '' and (fact.targetUrl or '') == ''
|
|
end
|
|
return (task.routingName or '') == (fact.routingName or '') and (task.targetUrl or '') == (fact.targetUrl or '')
|
|
end
|
|
|
|
local function authorize(fact)
|
|
if type(fact) ~= 'table' or type(fact.taskId) ~= 'string' or type(fact.proxyId) ~= 'string' or
|
|
type(fact.level) ~= 'string' or type(fact.checkerId) ~= 'string' or type(fact.leaseToken) ~= 'string' then
|
|
return nil, 'invalid'
|
|
end
|
|
local raw = redis.call('HGET', tasks_key, fact.taskId)
|
|
if not raw then
|
|
return nil, 'not_found'
|
|
end
|
|
local valid, task = pcall(cjson.decode, raw)
|
|
if not valid or type(task) ~= 'table' then
|
|
return nil, 'invalid'
|
|
end
|
|
if not matches_task(task, fact) then
|
|
return nil, 'observation'
|
|
end
|
|
if task.state == 'DONE' then
|
|
if task.leaseCheckerId == fact.checkerId and task.leaseToken == fact.leaseToken then
|
|
return task, 'ok'
|
|
end
|
|
return nil, 'not_owned'
|
|
end
|
|
if task.state ~= 'LEASED' or tonumber(task.leaseExpiresAtMs or 0) <= now_ms then
|
|
return nil, 'lease_expired'
|
|
end
|
|
if task.leaseCheckerId ~= fact.checkerId or task.leaseToken ~= fact.leaseToken then
|
|
return nil, 'not_owned'
|
|
end
|
|
return task, 'ok'
|
|
end
|
|
|
|
reap()
|
|
|
|
if operation == 'inflight' then
|
|
return finish({status = 'ok', count = redis.call('ZCARD', queued_key) + redis.call('ZCARD', leases_key)})
|
|
end
|
|
|
|
if operation == 'upstream_inflight' then
|
|
if type(payload.upstreamTasksKey) ~= 'string' or payload.upstreamTasksKey == '' then
|
|
return finish({status = 'invalid'})
|
|
end
|
|
return finish({status = 'ok', count = redis.call('ZCARD', payload.upstreamTasksKey)})
|
|
end
|
|
|
|
if operation == 'due' then
|
|
local result = cjson.decode('[]')
|
|
local scan_limit = tonumber(payload.scanLimit or limit)
|
|
if not scan_limit or scan_limit < limit then
|
|
return finish({status = 'invalid'})
|
|
end
|
|
local level = payload.level or 'BASIC'
|
|
if level ~= 'BASIC' and level ~= 'EGRESS' and level ~= 'TARGET' then
|
|
return finish({status = 'invalid'})
|
|
end
|
|
if level == 'EGRESS' and (type(payload.targetUrl) ~= 'string' or payload.targetUrl == '') then
|
|
return finish({status = 'invalid'})
|
|
end
|
|
if level == 'TARGET' and (type(payload.routingName) ~= 'string' or payload.routingName == '' or
|
|
type(payload.targetUrl) ~= 'string' or payload.targetUrl == '') then
|
|
return finish({status = 'invalid'})
|
|
end
|
|
local result_members = {}
|
|
local function append_candidate(proxy_id, record, member)
|
|
if result_members[member] then
|
|
return
|
|
end
|
|
result_members[member] = true
|
|
result[#result + 1] = {
|
|
proxyId = proxy_id, upstreamId = record.sourceUpstream, state = record.state, level = level,
|
|
routingName = level == 'TARGET' and payload.routingName or '',
|
|
targetUrl = (level == 'EGRESS' or level == 'TARGET') and payload.targetUrl or '', dueAtMs = now_ms
|
|
}
|
|
end
|
|
if level == 'EGRESS' or level == 'TARGET' then
|
|
local basic_ids = redis.call('ZRANGEBYSCORE', due_key, '-inf', now_ms, 'LIMIT', 0, scan_limit)
|
|
for _, proxy_id in ipairs(basic_ids) do
|
|
local _, record = live_record(proxy_id)
|
|
if record and valid_state(record.state) and
|
|
(type(payload.upstreamId) ~= 'string' or payload.upstreamId == '' or record.sourceUpstream == payload.upstreamId) then
|
|
local member = payload.targetUrl .. '\0' .. proxy_id
|
|
if level == 'TARGET' then
|
|
member = payload.routingName .. '\0' .. payload.targetUrl .. '\0' .. proxy_id
|
|
end
|
|
local reference = task_reference(level, proxy_id, payload.routingName or '', payload.targetUrl)
|
|
if not redis.call('HGET', ref_task_key, reference) then
|
|
local task_due_key = due_key_for_level(level)
|
|
redis.call('ZADD', task_due_key, now_ms, member)
|
|
touch(task_due_key, tonumber(record.expiresAtMs))
|
|
end
|
|
end
|
|
end
|
|
end
|
|
local source_due_key = due_key_for_level(level)
|
|
local ids = redis.call('ZRANGEBYSCORE', source_due_key, '-inf', now_ms, 'LIMIT', 0, scan_limit)
|
|
for _, member in ipairs(ids) do
|
|
local proxy_id = member
|
|
local matching_target = true
|
|
if level == 'EGRESS' then
|
|
local separator = string.find(member, '\0', 1, true)
|
|
if not separator or string.sub(member, 1, separator - 1) ~= payload.targetUrl then
|
|
matching_target = false
|
|
else
|
|
proxy_id = string.sub(member, separator + 1)
|
|
end
|
|
elseif level == 'TARGET' then
|
|
local first_separator = string.find(member, '\0', 1, true)
|
|
local second_separator = first_separator and string.find(member, '\0', first_separator + 1, true) or nil
|
|
if not first_separator or not second_separator or string.sub(member, 1, first_separator - 1) ~= payload.routingName or
|
|
string.sub(member, first_separator + 1, second_separator - 1) ~= payload.targetUrl then
|
|
matching_target = false
|
|
else
|
|
proxy_id = string.sub(member, second_separator + 1)
|
|
end
|
|
end
|
|
if matching_target then
|
|
local _, record = live_record(proxy_id)
|
|
local reference = task_reference(level, proxy_id, payload.routingName or '', payload.targetUrl or '')
|
|
if not record or not valid_state(record.state) then
|
|
redis.call('ZREM', source_due_key, member)
|
|
elseif (level == 'EGRESS' or level == 'TARGET') and redis.call('HGET', ref_task_key, reference) then
|
|
redis.call('ZREM', source_due_key, member)
|
|
elseif type(payload.upstreamId) == 'string' and payload.upstreamId ~= '' and record.sourceUpstream ~= payload.upstreamId then
|
|
-- The ref remains due for the scheduler that owns this upstream.
|
|
else
|
|
append_candidate(proxy_id, record, member)
|
|
if #result >= limit then
|
|
break
|
|
end
|
|
end
|
|
end
|
|
end
|
|
local encoded_candidates = cjson.encode(result)
|
|
if #result == 0 then
|
|
encoded_candidates = '[]'
|
|
end
|
|
return finish({status = 'ok', candidatesJSON = encoded_candidates})
|
|
end
|
|
|
|
if operation == 'offer' then
|
|
if type(payload.tasks) ~= 'table' then
|
|
return finish({status = 'invalid'})
|
|
end
|
|
local offered = 0
|
|
for _, task in ipairs(payload.tasks) do
|
|
if type(task) ~= 'table' or task.version ~= 1 or (task.level ~= 'BASIC' and task.level ~= 'EGRESS' and task.level ~= 'TARGET') or task.state ~= 'QUEUED' or
|
|
type(task.taskId) ~= 'string' or type(task.proxyId) ~= 'string' or tonumber(task.deadlineMs or 0) <= now_ms or
|
|
type(task.upstreamId) ~= 'string' or task.upstreamId == '' or type(task.upstreamTasksKey) ~= 'string' or
|
|
task.upstreamTasksKey == '' or type(task.reference) ~= 'string' or task.reference == '' or
|
|
type(task.dueMember) ~= 'string' or task.dueMember == '' or tonumber(task.nextDueMs or 0) <= now_ms or tonumber(task.attempts or 0) <= 0 or
|
|
tonumber(task.maxInFlight or 0) <= 0 then
|
|
return finish({status = 'invalid'})
|
|
end
|
|
if task.level == 'BASIC' and ((task.routingName or '') ~= '' or (task.targetUrl or '') ~= '') then
|
|
return finish({status = 'invalid'})
|
|
end
|
|
if task.level == 'EGRESS' and ((task.routingName or '') ~= '' or type(task.targetUrl) ~= 'string' or task.targetUrl == '') then
|
|
return finish({status = 'invalid'})
|
|
end
|
|
if task.level == 'TARGET' and (type(task.routingName) ~= 'string' or task.routingName == '' or
|
|
type(task.targetUrl) ~= 'string' or task.targetUrl == '') then
|
|
return finish({status = 'invalid'})
|
|
end
|
|
local task_due_key = due_key_for_level(task.level)
|
|
local score = redis.call('ZSCORE', task_due_key, task.dueMember)
|
|
local current = redis.call('HGET', ref_task_key, task.reference)
|
|
local raw, record = live_record(task.proxyId)
|
|
local active = redis.call('ZCARD', task.upstreamTasksKey)
|
|
if score and tonumber(score) <= now_ms and not current and raw and record.sourceUpstream == task.upstreamId and
|
|
active < tonumber(task.maxInFlight) and valid_state(record.state) then
|
|
if task.level == 'BASIC' and (record.state == 'FETCHED' or record.state == 'UNHEALTHY') then
|
|
decrement_state(record.sourceUpstream, record.state)
|
|
increment_state(record.sourceUpstream, 'CHECKING')
|
|
record.state = 'CHECKING'
|
|
raw = cjson.encode(record)
|
|
redis.call('HSET', records_key, task.proxyId, raw)
|
|
end
|
|
redis.call('ZREM', task_due_key, task.dueMember)
|
|
redis.call('HSET', tasks_key, task.taskId, cjson.encode(task))
|
|
redis.call('HSET', ref_task_key, task.reference, task.taskId)
|
|
redis.call('ZADD', queued_key, tonumber(task.priority), task.taskId)
|
|
redis.call('ZADD', task_expiry_key, tonumber(task.deadlineMs), task.taskId)
|
|
redis.call('ZADD', task.upstreamTasksKey, tonumber(task.deadlineMs), task.taskId)
|
|
local expires_at_ms = tonumber(record.expiresAtMs)
|
|
touch(queued_key, expires_at_ms)
|
|
touch(tasks_key, expires_at_ms)
|
|
touch(task_expiry_key, expires_at_ms)
|
|
touch(ref_task_key, expires_at_ms)
|
|
touch(task.upstreamTasksKey, expires_at_ms)
|
|
offered = offered + 1
|
|
end
|
|
end
|
|
return finish({status = 'ok', count = offered})
|
|
end
|
|
|
|
if operation == 'claim' then
|
|
if type(payload.checkerId) ~= 'string' or type(payload.instanceId) ~= 'string' or
|
|
type(payload.levels) ~= 'table' or type(payload.tokens) ~= 'table' then
|
|
return finish({status = 'invalid'})
|
|
end
|
|
local max_in_flight = tonumber(payload.maxInFlight or 0)
|
|
local lease_ttl_ms = tonumber(payload.leaseTTLMS or 0)
|
|
if not max_in_flight or max_in_flight <= 0 or not lease_ttl_ms or lease_ttl_ms <= 0 then
|
|
return finish({status = 'invalid'})
|
|
end
|
|
redis.call('ZREMRANGEBYSCORE', checker_leases_key, '-inf', now_ms)
|
|
local capacity = max_in_flight - redis.call('ZCARD', checker_leases_key)
|
|
if capacity <= 0 then
|
|
return finish({status = 'ok', tasks = cjson.decode('[]')})
|
|
end
|
|
local supported = {}
|
|
for _, level in ipairs(payload.levels) do
|
|
supported[level] = true
|
|
end
|
|
local candidates = redis.call('ZRANGE', queued_key, 0, limit - 1)
|
|
local result = cjson.decode('[]')
|
|
for _, task_id in ipairs(candidates) do
|
|
if #result >= capacity or #result >= #payload.tokens then
|
|
break
|
|
end
|
|
local raw = redis.call('HGET', tasks_key, task_id)
|
|
local valid, task = pcall(cjson.decode, raw or '')
|
|
if not valid or type(task) ~= 'table' or task.state ~= 'QUEUED' or tonumber(task.deadlineMs or 0) <= now_ms then
|
|
remove_task(task_id, true)
|
|
elseif not supported[task.level] then
|
|
-- Preserve unsupported work for a Checker that advertises this level.
|
|
else
|
|
local record_raw, record = live_record(task.proxyId)
|
|
if not record_raw or not record then
|
|
remove_task(task_id, false)
|
|
else
|
|
local lease_expires_at_ms = now_ms + lease_ttl_ms
|
|
if tonumber(task.deadlineMs) < lease_expires_at_ms then
|
|
lease_expires_at_ms = tonumber(task.deadlineMs)
|
|
end
|
|
local token = payload.tokens[#result + 1]
|
|
if type(token) ~= 'string' or token == '' then
|
|
return finish({status = 'invalid'})
|
|
end
|
|
task.state = 'LEASED'
|
|
task.leaseCheckerId = payload.checkerId
|
|
task.leaseInstanceId = payload.instanceId
|
|
task.leaseToken = token
|
|
task.leaseExpiresAtMs = lease_expires_at_ms
|
|
task.checkerLeaseKey = checker_leases_key
|
|
local encoded = cjson.encode(task)
|
|
redis.call('HSET', tasks_key, task_id, encoded)
|
|
redis.call('ZREM', queued_key, task_id)
|
|
redis.call('ZADD', leases_key, lease_expires_at_ms, task_id)
|
|
redis.call('ZADD', checker_leases_key, lease_expires_at_ms, task_id)
|
|
touch(checker_leases_key, tonumber(record.expiresAtMs))
|
|
result[#result + 1] = {task = encoded, record = record_raw}
|
|
end
|
|
end
|
|
end
|
|
return finish({status = 'ok', tasks = result})
|
|
end
|
|
|
|
if operation == 'authorize' or operation == 'complete' then
|
|
local task, status = authorize(payload.fact)
|
|
if status ~= 'ok' then
|
|
return finish({status = status})
|
|
end
|
|
if operation == 'complete' and task.state ~= 'DONE' then
|
|
redis.call('ZREM', queued_key, task.taskId)
|
|
redis.call('ZREM', leases_key, task.taskId)
|
|
if type(task.upstreamTasksKey) == 'string' and task.upstreamTasksKey ~= '' then
|
|
redis.call('ZREM', task.upstreamTasksKey, task.taskId)
|
|
end
|
|
if type(task.checkerLeaseKey) == 'string' and task.checkerLeaseKey ~= '' then
|
|
redis.call('ZREM', task.checkerLeaseKey, task.taskId)
|
|
end
|
|
task.state = 'DONE'
|
|
redis.call('HSET', tasks_key, task.taskId, cjson.encode(task))
|
|
local reference = task.reference or task.proxyId
|
|
if redis.call('HGET', ref_task_key, reference) == task.taskId then
|
|
redis.call('HDEL', ref_task_key, reference)
|
|
end
|
|
local _, record = live_record(task.proxyId)
|
|
if record then
|
|
local next_due_ms = tonumber(task.nextDueMs or 0)
|
|
if next_due_ms <= now_ms then
|
|
next_due_ms = now_ms
|
|
end
|
|
local task_due_key = due_key_for_level(task.level)
|
|
redis.call('ZADD', task_due_key, next_due_ms, task.dueMember or task.proxyId)
|
|
touch(task_due_key, tonumber(record.expiresAtMs))
|
|
end
|
|
end
|
|
return finish({status = 'ok'})
|
|
end
|
|
|
|
return finish({status = 'invalid'})
|