330 lines
13 KiB
Lua
330 lines
13 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 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 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.checkerLeaseKey) == 'string' and task.checkerLeaseKey ~= '' then
|
|
redis.call('ZREM', task.checkerLeaseKey, task_id)
|
|
end
|
|
if type(task.proxyId) == 'string' and redis.call('HGET', ref_task_key, task.proxyId) == task_id then
|
|
redis.call('HDEL', ref_task_key, task.proxyId)
|
|
end
|
|
if requeue and type(task.proxyId) == 'string' then
|
|
local _, record = live_record(task.proxyId)
|
|
if record then
|
|
redis.call('ZADD', due_key, now_ms, task.proxyId)
|
|
touch(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)
|
|
return task and task.taskId == fact.taskId and task.proxyId == fact.proxyId and task.level == fact.level and
|
|
(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 == 'due' then
|
|
local result = {}
|
|
local ids = redis.call('ZRANGEBYSCORE', due_key, '-inf', now_ms, 'LIMIT', 0, limit)
|
|
for _, proxy_id in ipairs(ids) do
|
|
local _, record = live_record(proxy_id)
|
|
if not record or not valid_state(record.state) then
|
|
redis.call('ZREM', due_key, proxy_id)
|
|
else
|
|
result[#result + 1] = {proxyId = proxy_id, state = record.state, dueAtMs = now_ms}
|
|
end
|
|
end
|
|
return finish({status = 'ok', candidates = result})
|
|
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' or task.state ~= 'QUEUED' or
|
|
type(task.taskId) ~= 'string' or type(task.proxyId) ~= 'string' or tonumber(task.deadlineMs or 0) <= now_ms or
|
|
tonumber(task.nextDueMs or 0) <= now_ms or tonumber(task.attempts or 0) <= 0 then
|
|
return finish({status = 'invalid'})
|
|
end
|
|
local score = redis.call('ZSCORE', due_key, task.proxyId)
|
|
local current = redis.call('HGET', ref_task_key, task.proxyId)
|
|
local raw, record = live_record(task.proxyId)
|
|
if score and tonumber(score) <= now_ms and not current and raw and valid_state(record.state) then
|
|
if 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', due_key, task.proxyId)
|
|
redis.call('HSET', tasks_key, task.taskId, cjson.encode(task))
|
|
redis.call('HSET', ref_task_key, task.proxyId, task.taskId)
|
|
redis.call('ZADD', queued_key, tonumber(task.priority), task.taskId)
|
|
redis.call('ZADD', task_expiry_key, 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)
|
|
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 = {}})
|
|
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 = {}
|
|
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.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))
|
|
if redis.call('HGET', ref_task_key, task.proxyId) == task.taskId then
|
|
redis.call('HDEL', ref_task_key, task.proxyId)
|
|
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
|
|
redis.call('ZADD', due_key, next_due_ms, task.proxyId)
|
|
touch(due_key, tonumber(record.expiresAtMs))
|
|
end
|
|
end
|
|
return finish({status = 'ok'})
|
|
end
|
|
|
|
return finish({status = 'invalid'})
|