221 lines
7.5 KiB
Lua
221 lines
7.5 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 checked_at_ms = tonumber(ARGV[1])
|
|
local next_state = ARGV[2]
|
|
local latency_ns = tonumber(ARGV[3])
|
|
local cleanup_limit = tonumber(ARGV[4])
|
|
local operation_ttl_ms = tonumber(ARGV[5])
|
|
local proxy_id = ARGV[6]
|
|
|
|
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 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 increment_state(upstream, state)
|
|
if not is_counted(state) then
|
|
return
|
|
end
|
|
redis.call('HINCRBY', state_inventory_key, state_field(upstream, state), 1)
|
|
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 value = redis.call('HINCRBY', state_inventory_key, field, -1)
|
|
if value <= 0 then
|
|
redis.call('HDEL', state_inventory_key, field)
|
|
end
|
|
end
|
|
|
|
local function remove_available(id, record)
|
|
redis.call('ZREM', available_key, id)
|
|
for _, index_key in ipairs(record and record.indexKeys or {}) do
|
|
redis.call('ZREM', index_key, id)
|
|
end
|
|
end
|
|
|
|
local function remove_owned(id, record)
|
|
if record and type(record.ownerIndexKey) == 'string' and record.ownerIndexKey ~= '' then
|
|
redis.call('ZREM', record.ownerIndexKey, id)
|
|
end
|
|
end
|
|
|
|
local function remove_proxy(id)
|
|
local raw = redis.call('HGET', records_key, id)
|
|
local record = nil
|
|
if raw then
|
|
record = cjson.decode(raw)
|
|
remove_available(id, record)
|
|
remove_owned(id, record)
|
|
if is_managed(record.state) then
|
|
decrement_inventory(record.sourceUpstream)
|
|
end
|
|
decrement_state(record.sourceUpstream, record.state)
|
|
else
|
|
redis.call('ZREM', available_key, id)
|
|
end
|
|
local digest = redis.call('HGET', idkeys_key, id)
|
|
if digest and redis.call('HGET', unique_key, digest) == id then
|
|
redis.call('HDEL', unique_key, digest)
|
|
end
|
|
redis.call('HDEL', idkeys_key, id)
|
|
redis.call('HDEL', records_key, id)
|
|
redis.call('ZREM', expiry_key, id)
|
|
redis.call('HDEL', owners_key, id)
|
|
redis.call('ZREM', owner_expiry_key, id)
|
|
end
|
|
|
|
local function cleanup_expired()
|
|
local expired = redis.call('ZRANGEBYSCORE', expiry_key, '-inf', checked_at_ms, 'LIMIT', 0, cleanup_limit)
|
|
for _, id in ipairs(expired) do
|
|
remove_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 finish(reply, hard_expiry_ms)
|
|
local encoded = cjson.encode(reply)
|
|
redis.call('SET', operation_key, encoded, 'PX', operation_ttl_ms)
|
|
if hard_expiry_ms then
|
|
local operation_expiry_ms = redis.call('PEXPIRETIME', operation_key)
|
|
if operation_expiry_ms > hard_expiry_ms then
|
|
redis.call('PEXPIREAT', operation_key, hard_expiry_ms)
|
|
end
|
|
end
|
|
return encoded
|
|
end
|
|
|
|
local transitions = {
|
|
FETCHED = {CHECKING = true, EXPIRED = true, REMOVED = true},
|
|
CHECKING = {AVAILABLE = true, UNHEALTHY = true, EXPIRED = true, REMOVED = true},
|
|
AVAILABLE = {SUSPECT = true, DRAINING = true, EXTRACTED = true, EXPIRED = true},
|
|
SUSPECT = {AVAILABLE = true, UNHEALTHY = true, DRAINING = true, EXPIRED = true},
|
|
DRAINING = {EXPIRED = true, UNHEALTHY = true, REMOVED = true},
|
|
UNHEALTHY = {CHECKING = true, REMOVED = true, EXPIRED = true},
|
|
EXTRACTED = {EXPIRED = true, REMOVED = true},
|
|
EXPIRED = {REMOVED = true},
|
|
REMOVED = {},
|
|
}
|
|
|
|
cleanup_expired()
|
|
|
|
local raw = redis.call('HGET', records_key, proxy_id)
|
|
if not raw then
|
|
return finish({status = 'not_found'})
|
|
end
|
|
local record = cjson.decode(raw)
|
|
if type(record.ownerIndexKey) ~= 'string' or record.ownerIndexKey == '' then
|
|
return finish({status = 'invalid'})
|
|
end
|
|
if tonumber(record.expiresAtMs) <= checked_at_ms then
|
|
remove_proxy(proxy_id)
|
|
return finish({status = 'not_found'})
|
|
end
|
|
|
|
local last_checked_at_ms = tonumber(record.lastCheckedAtMs or '0')
|
|
if checked_at_ms < last_checked_at_ms then
|
|
return finish({status = 'stale'})
|
|
end
|
|
if checked_at_ms == last_checked_at_ms then
|
|
if record.state ~= next_state then
|
|
return finish({status = 'stale'})
|
|
end
|
|
return finish({status = 'ok', record = raw}, tonumber(record.expiresAtMs))
|
|
end
|
|
if record.state ~= next_state and not (transitions[record.state] and transitions[record.state][next_state]) then
|
|
return finish({status = 'invalid'})
|
|
end
|
|
|
|
local previous_state = record.state
|
|
local was_managed = is_managed(previous_state)
|
|
local will_be_managed = is_managed(next_state)
|
|
remove_available(proxy_id, record)
|
|
record.state = next_state
|
|
record.lastCheckedAtMs = checked_at_ms
|
|
record.latencyNs = latency_ns
|
|
if next_state == 'AVAILABLE' then
|
|
record.lastSuccessAtMs = checked_at_ms
|
|
end
|
|
if was_managed and not will_be_managed then
|
|
decrement_inventory(record.sourceUpstream)
|
|
elseif not was_managed and will_be_managed then
|
|
redis.call('HINCRBY', inventory_key, record.sourceUpstream, 1)
|
|
end
|
|
if previous_state ~= next_state then
|
|
decrement_state(record.sourceUpstream, previous_state)
|
|
increment_state(record.sourceUpstream, next_state)
|
|
end
|
|
|
|
local encoded = cjson.encode(record)
|
|
redis.call('HSET', records_key, proxy_id, encoded)
|
|
local owned = (record.ownerWorkerId and record.ownerWorkerId ~= '') or redis.call('HEXISTS', owners_key, proxy_id) == 1
|
|
remove_owned(proxy_id, record)
|
|
if next_state == 'AVAILABLE' and tonumber(record.usableUntilMs) > checked_at_ms then
|
|
if owned then
|
|
redis.call('ZADD', record.ownerIndexKey, record.usableUntilMs, proxy_id)
|
|
touch(record.ownerIndexKey, tonumber(record.expiresAtMs))
|
|
else
|
|
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
|
|
end
|
|
touch(records_key, tonumber(record.expiresAtMs))
|
|
touch(unique_key, tonumber(record.expiresAtMs))
|
|
touch(idkeys_key, tonumber(record.expiresAtMs))
|
|
touch(expiry_key, tonumber(record.expiresAtMs))
|
|
touch(available_key, tonumber(record.expiresAtMs))
|
|
touch(inventory_key, tonumber(record.expiresAtMs))
|
|
touch(state_inventory_key, tonumber(record.expiresAtMs))
|
|
touch(owners_key, tonumber(record.expiresAtMs))
|
|
touch(owner_expiry_key, tonumber(record.expiresAtMs))
|
|
|
|
return finish({status = 'ok', record = encoded}, tonumber(record.expiresAtMs))
|