150 lines
5.3 KiB
Lua
150 lines
5.3 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 health_due_key = KEYS[10]
|
|
local health_queued_key = KEYS[11]
|
|
local health_leases_key = KEYS[12]
|
|
local health_tasks_key = KEYS[13]
|
|
local health_task_expiry_key = KEYS[14]
|
|
local health_ref_task_key = KEYS[15]
|
|
|
|
local now_ms = tonumber(ARGV[1])
|
|
local cleanup_limit = tonumber(ARGV[2])
|
|
local upstream_ids = cjson.decode(ARGV[3])
|
|
|
|
local function is_managed(state)
|
|
return state == 'FETCHED' or state == 'CHECKING' or state == 'AVAILABLE' or
|
|
state == 'SUSPECT' or state == 'DRAINING'
|
|
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_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 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_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)
|
|
if decoded and type(record) == 'table' then
|
|
if is_managed(record.state) then
|
|
decrement_inventory(record.sourceUpstream)
|
|
end
|
|
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)
|
|
redis.call('HDEL', owners_key, proxy_id)
|
|
redis.call('ZREM', owner_expiry_key, proxy_id)
|
|
remove_health_task(proxy_id)
|
|
end
|
|
|
|
local expired = redis.call('ZRANGEBYSCORE', expiry_key, '-inf', now_ms, 'LIMIT', 0, cleanup_limit)
|
|
for _, proxy_id in ipairs(expired) do
|
|
remove_proxy(proxy_id)
|
|
end
|
|
|
|
local oldest = redis.call('ZRANGE', expiry_key, 0, 0, 'WITHSCORES')
|
|
if #oldest == 2 and tonumber(oldest[2]) <= now_ms then
|
|
return cjson.encode({status = 'unavailable', inventories = cjson.decode('[]')})
|
|
end
|
|
|
|
local inventories = cjson.decode('[]')
|
|
local states = {'FETCHED', 'CHECKING', 'AVAILABLE', 'SUSPECT', 'DRAINING', 'UNHEALTHY', 'EXTRACTED'}
|
|
for _, upstream_id in ipairs(upstream_ids) do
|
|
if type(upstream_id) ~= 'string' or upstream_id == '' then
|
|
return cjson.encode({status = 'invalid', inventories = cjson.decode('[]')})
|
|
end
|
|
local counts = {}
|
|
for _, state in ipairs(states) do
|
|
local count = tonumber(redis.call('HGET', state_inventory_key, state_field(upstream_id, state)) or '0')
|
|
if count < 0 then
|
|
return cjson.encode({status = 'unavailable', inventories = cjson.decode('[]')})
|
|
end
|
|
counts[string.lower(state)] = count
|
|
end
|
|
counts.upstreamId = upstream_id
|
|
inventories[#inventories + 1] = counts
|
|
end
|
|
|
|
return cjson.encode({status = 'ok', inventories = inventories})
|