268 lines
11 KiB
Lua
268 lines
11 KiB
Lua
local records_key = KEYS[1]
|
|
local unique_key = KEYS[2]
|
|
local idkeys_key = KEYS[3]
|
|
local expiry_key = KEYS[4]
|
|
local available_key = KEYS[5]
|
|
local inventory_key = KEYS[6]
|
|
local state_inventory_key = KEYS[7]
|
|
local owners_key = KEYS[8]
|
|
local owner_expiry_key = KEYS[9]
|
|
local operation_key = KEYS[10]
|
|
local health_due_key = KEYS[11]
|
|
local health_queued_key = KEYS[12]
|
|
local health_leases_key = KEYS[13]
|
|
local health_tasks_key = KEYS[14]
|
|
local health_task_expiry_key = KEYS[15]
|
|
local health_ref_task_key = KEYS[16]
|
|
local health_unhealthy_key = KEYS[17]
|
|
local owned_upstream_key = KEYS[18]
|
|
local upstream_drain_policies_key = KEYS[19]
|
|
|
|
local operation = ARGV[1]
|
|
local now_ms = tonumber(ARGV[2])
|
|
local limit = tonumber(ARGV[3])
|
|
local upstream_id = ARGV[4]
|
|
local operation_ttl_ms = tonumber(ARGV[5])
|
|
local unhealthy_policies_json = ARGV[6]
|
|
|
|
local function finish(reply)
|
|
local encoded = cjson.encode(reply)
|
|
redis.call('SET', operation_key, encoded, 'PX', operation_ttl_ms)
|
|
return encoded
|
|
end
|
|
|
|
local committed = redis.call('GET', operation_key)
|
|
if committed then
|
|
return committed
|
|
end
|
|
|
|
local function is_managed(state)
|
|
return state == 'FETCHED' or state == 'CHECKING' or state == 'AVAILABLE' or
|
|
state == 'SUSPECT' or state == 'DRAINING'
|
|
end
|
|
|
|
local function decrement_inventory(upstream)
|
|
if type(upstream) ~= 'string' or upstream == '' then
|
|
return
|
|
end
|
|
local count = redis.call('HINCRBY', inventory_key, upstream, -1)
|
|
if count < 0 then
|
|
redis.call('HSET', inventory_key, upstream, 0)
|
|
end
|
|
end
|
|
|
|
local function state_field(upstream, state)
|
|
return string.len(upstream) .. ':' .. upstream .. ':' .. state
|
|
end
|
|
|
|
local function is_counted(state)
|
|
return state == 'FETCHED' or state == 'CHECKING' or state == 'AVAILABLE' or
|
|
state == 'SUSPECT' or state == 'DRAINING' or state == 'UNHEALTHY' or state == 'EXTRACTED'
|
|
end
|
|
|
|
local function decrement_state(upstream, state)
|
|
if type(upstream) ~= 'string' or upstream == '' or not is_counted(state) then
|
|
return
|
|
end
|
|
local field = state_field(upstream, state)
|
|
local count = redis.call('HINCRBY', state_inventory_key, field, -1)
|
|
if count <= 0 then
|
|
redis.call('HDEL', state_inventory_key, field)
|
|
end
|
|
end
|
|
|
|
local function remove_available(proxy_id, record)
|
|
redis.call('ZREM', available_key, proxy_id)
|
|
local index_keys = record and record.indexKeys
|
|
if type(index_keys) == 'table' then
|
|
for _, index_key in ipairs(index_keys) do
|
|
if type(index_key) == 'string' and index_key ~= '' then
|
|
redis.call('ZREM', index_key, proxy_id)
|
|
end
|
|
end
|
|
end
|
|
end
|
|
|
|
local function remove_owned(proxy_id, record)
|
|
if record and type(record.ownerIndexKey) == 'string' and record.ownerIndexKey ~= '' then
|
|
redis.call('ZREM', record.ownerIndexKey, proxy_id)
|
|
end
|
|
end
|
|
|
|
local function remove_worker_owned(proxy_id)
|
|
local raw = redis.call('HGET', owners_key, proxy_id)
|
|
if not raw then
|
|
return
|
|
end
|
|
local decoded, assignment = pcall(cjson.decode, raw)
|
|
if decoded and type(assignment) == 'table' and
|
|
type(assignment.workerIndexKey) == 'string' and assignment.workerIndexKey ~= '' then
|
|
redis.call('ZREM', assignment.workerIndexKey, proxy_id)
|
|
end
|
|
end
|
|
|
|
local function remove_health_task(proxy_id)
|
|
redis.call('ZREM', health_due_key, proxy_id)
|
|
local task_id = redis.call('HGET', health_ref_task_key, proxy_id)
|
|
if not task_id then
|
|
return
|
|
end
|
|
local raw = redis.call('HGET', health_tasks_key, task_id)
|
|
if raw then
|
|
local decoded, task = pcall(cjson.decode, raw)
|
|
if decoded and type(task) == 'table' then
|
|
if type(task.checkerLeaseKey) == 'string' and task.checkerLeaseKey ~= '' then
|
|
redis.call('ZREM', task.checkerLeaseKey, task_id)
|
|
end
|
|
if type(task.upstreamTasksKey) == 'string' and task.upstreamTasksKey ~= '' then
|
|
redis.call('ZREM', task.upstreamTasksKey, task_id)
|
|
end
|
|
end
|
|
end
|
|
redis.call('ZREM', health_queued_key, task_id)
|
|
redis.call('ZREM', health_leases_key, task_id)
|
|
redis.call('ZREM', health_task_expiry_key, task_id)
|
|
redis.call('HDEL', health_tasks_key, task_id)
|
|
redis.call('HDEL', health_ref_task_key, proxy_id)
|
|
end
|
|
|
|
local function remove_proxy(proxy_id)
|
|
local raw = redis.call('HGET', records_key, proxy_id)
|
|
local record = nil
|
|
if raw then
|
|
local decoded
|
|
decoded, record = pcall(cjson.decode, raw)
|
|
remove_available(proxy_id, decoded and record or nil)
|
|
remove_owned(proxy_id, decoded and record or nil)
|
|
if decoded and type(record) == 'table' and is_managed(record.state) then
|
|
decrement_inventory(record.sourceUpstream)
|
|
end
|
|
if decoded and type(record) == 'table' then
|
|
decrement_state(record.sourceUpstream, record.state)
|
|
end
|
|
else
|
|
redis.call('ZREM', available_key, proxy_id)
|
|
end
|
|
local digest = redis.call('HGET', idkeys_key, proxy_id)
|
|
if digest and redis.call('HGET', unique_key, digest) == proxy_id then
|
|
redis.call('HDEL', unique_key, digest)
|
|
end
|
|
redis.call('HDEL', idkeys_key, proxy_id)
|
|
redis.call('HDEL', records_key, proxy_id)
|
|
redis.call('ZREM', expiry_key, proxy_id)
|
|
remove_worker_owned(proxy_id)
|
|
redis.call('HDEL', owners_key, proxy_id)
|
|
redis.call('ZREM', owner_expiry_key, proxy_id)
|
|
redis.call('ZREM', health_unhealthy_key, proxy_id)
|
|
remove_health_task(proxy_id)
|
|
end
|
|
|
|
local expired = redis.call('ZRANGEBYSCORE', expiry_key, '-inf', now_ms, 'LIMIT', 0, limit)
|
|
for _, proxy_id in ipairs(expired) do
|
|
remove_proxy(proxy_id)
|
|
end
|
|
|
|
if operation == 'sweep' then
|
|
return finish({status = 'ok', count = #expired})
|
|
end
|
|
if operation == 'unhealthy' then
|
|
local decoded, policies = pcall(cjson.decode, unhealthy_policies_json or '')
|
|
if not decoded or type(policies) ~= 'table' then
|
|
return finish({status = 'invalid', count = 0, deferredOwned = 0, drainCandidates = {}})
|
|
end
|
|
local scan_limit = limit * 4
|
|
if scan_limit > 1024 then
|
|
scan_limit = 1024
|
|
end
|
|
local candidates = redis.call('ZRANGEBYSCORE', health_unhealthy_key, '-inf', now_ms, 'LIMIT', 0, scan_limit)
|
|
local removed = 0
|
|
local deferred_owned = 0
|
|
local drain_candidates = {}
|
|
for _, proxy_id in ipairs(candidates) do
|
|
local raw = redis.call('HGET', records_key, proxy_id)
|
|
if not raw then
|
|
redis.call('ZREM', health_unhealthy_key, proxy_id)
|
|
else
|
|
local valid, record = pcall(cjson.decode, raw)
|
|
local remove_after_ms = valid and type(record) == 'table' and tonumber(policies[record.sourceUpstream]) or nil
|
|
local unhealthy_since_ms = valid and type(record) == 'table' and tonumber(record.unhealthySinceMs or '0') or nil
|
|
if not valid or type(record) ~= 'table' or record.state ~= 'UNHEALTHY' or not remove_after_ms or
|
|
remove_after_ms <= 0 or not unhealthy_since_ms or unhealthy_since_ms <= 0 then
|
|
redis.call('ZREM', health_unhealthy_key, proxy_id)
|
|
elseif unhealthy_since_ms + remove_after_ms > now_ms then
|
|
redis.call('ZADD', health_unhealthy_key, unhealthy_since_ms + remove_after_ms, proxy_id)
|
|
else
|
|
local owner_raw = redis.call('HGET', owners_key, proxy_id)
|
|
local owned = owner_raw or (type(record.ownerWorkerId) == 'string' and record.ownerWorkerId ~= '')
|
|
if owned then
|
|
local retry_at_ms = now_ms + 1000
|
|
redis.call('ZADD', health_unhealthy_key, retry_at_ms, proxy_id)
|
|
deferred_owned = deferred_owned + 1
|
|
local owner_valid = false
|
|
local owner = nil
|
|
if owner_raw then
|
|
owner_valid, owner = pcall(cjson.decode, owner_raw)
|
|
end
|
|
if owner_valid and type(owner) == 'table' and owner.draining ~= true and
|
|
type(owner.workerId) == 'string' and owner.workerId ~= '' and tonumber(owner.epoch) and
|
|
tonumber(owner.epoch) > 0 and #drain_candidates < limit then
|
|
table.insert(drain_candidates, {
|
|
proxyId = proxy_id,
|
|
workerId = owner.workerId,
|
|
assignmentEpoch = tonumber(owner.epoch),
|
|
unhealthySinceMs = unhealthy_since_ms,
|
|
})
|
|
end
|
|
elseif removed < limit then
|
|
remove_proxy(proxy_id)
|
|
removed = removed + 1
|
|
end
|
|
end
|
|
end
|
|
end
|
|
return finish({status = 'ok', count = removed, deferredOwned = deferred_owned, drainCandidates = drain_candidates})
|
|
end
|
|
if operation == 'disabled_upstream' then
|
|
local raw_policy = redis.call('HGET', upstream_drain_policies_key, upstream_id)
|
|
local policy = raw_policy and decode_table(raw_policy) or nil
|
|
if not policy or policy.version ~= 1 or policy.upstreamId ~= upstream_id or policy.enabled ~= false or
|
|
type(policy.revision) ~= 'string' or not string.match(policy.revision, '^[0-9]+$') or policy.revision == '0' then
|
|
return finish({status = 'stale', count = 0, disabledDrainCandidates = {}})
|
|
end
|
|
local scan_limit = limit * 4
|
|
if scan_limit > 1024 then
|
|
scan_limit = 1024
|
|
end
|
|
local proxy_ids = redis.call('ZRANGEBYSCORE', owned_upstream_key, '(' .. now_ms, '+inf', 'LIMIT', 0, scan_limit)
|
|
local candidates = {}
|
|
for _, proxy_id in ipairs(proxy_ids) do
|
|
local record = decode_table(redis.call('HGET', records_key, proxy_id))
|
|
local owner = decode_table(redis.call('HGET', owners_key, proxy_id))
|
|
if record and owner and record.sourceUpstream == upstream_id and record.state == 'AVAILABLE' and
|
|
record.ownerWorkerId == owner.workerId and owner.draining == false and
|
|
type(owner.workerId) == 'string' and owner.workerId ~= '' and tonumber(owner.epoch) and
|
|
tonumber(owner.epoch) > 0 and tonumber(owner.expiresAtMs) and tonumber(owner.expiresAtMs) > now_ms then
|
|
table.insert(candidates, {
|
|
upstreamId = upstream_id,
|
|
policyRevision = policy.revision,
|
|
proxyId = proxy_id,
|
|
workerId = owner.workerId,
|
|
assignmentEpoch = tonumber(owner.epoch),
|
|
})
|
|
if #candidates == limit then
|
|
break
|
|
end
|
|
end
|
|
end
|
|
return finish({status = 'ok', count = 0, disabledDrainCandidates = candidates})
|
|
end
|
|
if operation == 'inventory' then
|
|
local count = tonumber(redis.call('HGET', inventory_key, upstream_id) or '0')
|
|
if count < 0 then
|
|
count = 0
|
|
redis.call('HSET', inventory_key, upstream_id, 0)
|
|
end
|
|
return finish({status = 'ok', count = count})
|
|
end
|
|
return finish({status = 'invalid', count = 0})
|