193 lines
6.6 KiB
Lua
193 lines
6.6 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 owners_key = KEYS[7]
|
|
local owner_expiry_key = KEYS[8]
|
|
local operation_key = KEYS[9]
|
|
|
|
local now_ms = tonumber(ARGV[1])
|
|
local cleanup_limit = tonumber(ARGV[2])
|
|
local max_size = tonumber(ARGV[3])
|
|
local operation_ttl_ms = tonumber(ARGV[4])
|
|
local candidates = cjson.decode(ARGV[5])
|
|
|
|
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 not upstream or upstream == '' then
|
|
return
|
|
end
|
|
local value = redis.call('HINCRBY', inventory_key, upstream, -1)
|
|
if value < 0 then
|
|
redis.call('HSET', inventory_key, upstream, 0)
|
|
end
|
|
end
|
|
|
|
local function remove_available(proxy_id, record)
|
|
redis.call('ZREM', available_key, proxy_id)
|
|
local indexes = record and record.indexKeys or {}
|
|
for _, index_key in ipairs(indexes) do
|
|
redis.call('ZREM', index_key, proxy_id)
|
|
end
|
|
end
|
|
|
|
local function remove_proxy(proxy_id)
|
|
local raw = redis.call('HGET', records_key, proxy_id)
|
|
local record = nil
|
|
if raw then
|
|
record = cjson.decode(raw)
|
|
remove_available(proxy_id, record)
|
|
if is_managed(record.state) then
|
|
decrement_inventory(record.sourceUpstream)
|
|
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)
|
|
end
|
|
|
|
local function cleanup_expired()
|
|
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
|
|
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 add_available(proxy_id, record)
|
|
if record.state ~= 'AVAILABLE' or (record.ownerWorkerId and record.ownerWorkerId ~= '') or
|
|
tonumber(record.usableUntilMs) <= now_ms then
|
|
return
|
|
end
|
|
redis.call('ZADD', available_key, record.usableUntilMs, proxy_id)
|
|
for _, index_key in ipairs(record.indexKeys or {}) do
|
|
redis.call('ZADD', index_key, record.usableUntilMs, proxy_id)
|
|
touch(index_key, tonumber(record.expiresAtMs))
|
|
end
|
|
end
|
|
|
|
local function finish(reply)
|
|
local encoded = cjson.encode(reply)
|
|
redis.call('SET', operation_key, encoded, 'PX', operation_ttl_ms)
|
|
return encoded
|
|
end
|
|
|
|
cleanup_expired()
|
|
|
|
for _, candidate in ipairs(candidates) do
|
|
local mapped = redis.call('HGET', idkeys_key, candidate.proxyId)
|
|
if mapped and mapped ~= candidate.uniqueDigest then
|
|
return finish({status = 'invalid', accepted = 0, inserted = 0, refreshed = 0, dropped = 0})
|
|
end
|
|
end
|
|
|
|
local accepted = #candidates
|
|
local inserted = 0
|
|
local refreshed = 0
|
|
local dropped = 0
|
|
local max_expiry_ms = 0
|
|
|
|
for _, candidate in ipairs(candidates) do
|
|
local incumbent_id = redis.call('HGET', unique_key, candidate.uniqueDigest)
|
|
local current_raw = incumbent_id and redis.call('HGET', records_key, incumbent_id) or nil
|
|
if incumbent_id and not current_raw then
|
|
remove_proxy(incumbent_id)
|
|
redis.call('HDEL', unique_key, candidate.uniqueDigest)
|
|
incumbent_id = nil
|
|
end
|
|
|
|
if current_raw then
|
|
local current = cjson.decode(current_raw)
|
|
if tonumber(current.expiresAtMs) <= now_ms then
|
|
remove_proxy(incumbent_id)
|
|
incumbent_id = nil
|
|
current_raw = nil
|
|
elseif current.state == 'EXTRACTED' or current.sourceUpstream ~= candidate.upstream then
|
|
refreshed = refreshed + 1
|
|
else
|
|
local incoming = cjson.decode(candidate.record)
|
|
remove_available(incumbent_id, current)
|
|
incoming.id = current.id
|
|
incoming.createdAtMs = current.createdAtMs
|
|
incoming.state = current.state
|
|
incoming.lastCheckedAtMs = current.lastCheckedAtMs
|
|
incoming.lastSuccessAtMs = current.lastSuccessAtMs
|
|
incoming.latencyNs = current.latencyNs
|
|
incoming.ownerWorkerId = current.ownerWorkerId
|
|
local encoded = cjson.encode(incoming)
|
|
redis.call('HSET', records_key, incumbent_id, encoded)
|
|
redis.call('ZADD', expiry_key, incoming.expiresAtMs, incumbent_id)
|
|
add_available(incumbent_id, incoming)
|
|
if tonumber(incoming.expiresAtMs) > max_expiry_ms then
|
|
max_expiry_ms = tonumber(incoming.expiresAtMs)
|
|
end
|
|
refreshed = refreshed + 1
|
|
end
|
|
end
|
|
|
|
if not incumbent_id then
|
|
local current_size = tonumber(redis.call('HGET', inventory_key, candidate.upstream) or '0')
|
|
if current_size >= max_size then
|
|
dropped = dropped + 1
|
|
else
|
|
local incoming = cjson.decode(candidate.record)
|
|
redis.call('HSET', records_key, candidate.proxyId, candidate.record)
|
|
redis.call('HSET', unique_key, candidate.uniqueDigest, candidate.proxyId)
|
|
redis.call('HSET', idkeys_key, candidate.proxyId, candidate.uniqueDigest)
|
|
redis.call('ZADD', expiry_key, incoming.expiresAtMs, candidate.proxyId)
|
|
if is_managed(incoming.state) then
|
|
redis.call('HINCRBY', inventory_key, candidate.upstream, 1)
|
|
end
|
|
add_available(candidate.proxyId, incoming)
|
|
if tonumber(incoming.expiresAtMs) > max_expiry_ms then
|
|
max_expiry_ms = tonumber(incoming.expiresAtMs)
|
|
end
|
|
inserted = inserted + 1
|
|
end
|
|
end
|
|
end
|
|
|
|
if max_expiry_ms > 0 then
|
|
touch(records_key, max_expiry_ms)
|
|
touch(unique_key, max_expiry_ms)
|
|
touch(idkeys_key, max_expiry_ms)
|
|
touch(expiry_key, max_expiry_ms)
|
|
touch(available_key, max_expiry_ms)
|
|
touch(inventory_key, max_expiry_ms)
|
|
touch(owners_key, max_expiry_ms)
|
|
touch(owner_expiry_key, max_expiry_ms)
|
|
end
|
|
|
|
return finish({
|
|
status = 'ok', accepted = accepted, inserted = inserted,
|
|
refreshed = refreshed, dropped = dropped,
|
|
})
|