proxy-pool/internal/adapters/redisactivity/scripts/sweep.lua
youfak d985f211d5
Some checks are pending
ci / proto (push) Waiting to run
ci / test (ubuntu-latest) (push) Waiting to run
ci / test (windows-latest) (push) Waiting to run
ci / race (push) Waiting to run
ci / integration (push) Waiting to run
fix: clear health tasks with proxy lifecycle
2026-07-31 22:12:07 +08:00

172 lines
5.7 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 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 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)
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 == '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})