373 lines
14 KiB
Lua
373 lines
14 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 idempotency_key = KEYS[11]
|
|
|
|
local now_ms = tonumber(ARGV[1])
|
|
local requested = tonumber(ARGV[2])
|
|
local fulfillment = ARGV[3]
|
|
local reserve = tonumber(ARGV[4])
|
|
local min_remaining_ttl_ms = tonumber(ARGV[5])
|
|
local max_health_age_ms = tonumber(ARGV[6])
|
|
local max_candidate_scan = tonumber(ARGV[7])
|
|
local cleanup_limit = tonumber(ARGV[8])
|
|
local idempotency_ttl_ms = tonumber(ARGV[9])
|
|
local operation_ttl_ms = tonumber(ARGV[10])
|
|
local request_digest = ARGV[11]
|
|
local has_idempotency = tonumber(ARGV[12]) == 1
|
|
local filters = cjson.decode(ARGV[13])
|
|
|
|
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 committed = redis.call('GET', operation_key)
|
|
if committed then
|
|
local reply = cjson.decode(committed)
|
|
if reply.requestDigest == request_digest then
|
|
return committed
|
|
end
|
|
return cjson.encode({status = 'conflict', requestDigest = request_digest})
|
|
end
|
|
|
|
if has_idempotency then
|
|
local replay = redis.call('GET', idempotency_key)
|
|
if replay then
|
|
local replay_record = cjson.decode(replay)
|
|
if tonumber(replay_record.expiresAtMs) <= now_ms then
|
|
redis.call('DEL', idempotency_key)
|
|
elseif replay_record.requestDigest ~= request_digest then
|
|
return finish({status = 'conflict', requestDigest = request_digest})
|
|
else
|
|
return finish({status = 'ok', requestDigest = request_digest, record = replay}, tonumber(replay_record.expiresAtMs))
|
|
end
|
|
end
|
|
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 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(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_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' 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)
|
|
redis.call('HDEL', owners_key, proxy_id)
|
|
redis.call('ZREM', owner_expiry_key, 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 function to_set(values)
|
|
local result = {}
|
|
for _, value in ipairs(values or {}) do
|
|
result[value] = true
|
|
end
|
|
return result
|
|
end
|
|
|
|
local protocol_filter = to_set(filters.protocols)
|
|
local region_filter = to_set(filters.regions)
|
|
local carrier_filter = to_set(filters.carriers)
|
|
local upstream_filter = to_set(filters.upstreams)
|
|
|
|
local function matches(filter_values, filter_set, value)
|
|
return #filter_values == 0 or filter_set[value] == true
|
|
end
|
|
|
|
local function valid_scheme(value)
|
|
return value == 'http' or value == 'https' or value == 'socks5'
|
|
end
|
|
|
|
local function valid_optional_string(value)
|
|
return value == nil or type(value) == 'string'
|
|
end
|
|
|
|
local function valid_integer(value)
|
|
return type(value) == 'number' and value == math.floor(value)
|
|
end
|
|
|
|
local function valid_proxy_state(value)
|
|
return value == 'FETCHED' or value == 'CHECKING' or value == 'AVAILABLE' or
|
|
value == 'SUSPECT' or value == 'DRAINING' or value == 'UNHEALTHY' or
|
|
value == 'EXTRACTED' or value == 'EXPIRED' or value == 'REMOVED'
|
|
end
|
|
|
|
local function valid_proxy_record(proxy_id, record)
|
|
if type(record) ~= 'table' or record.version ~= 1 or record.id ~= proxy_id or
|
|
not valid_scheme(record.scheme) or type(record.host) ~= 'string' or record.host == '' or
|
|
type(record.sourceUpstream) ~= 'string' or record.sourceUpstream == '' or
|
|
not valid_optional_string(record.username) or not valid_optional_string(record.password) or
|
|
not valid_optional_string(record.credentialVersion) or
|
|
not valid_optional_string(record.ownerWorkerId) or not valid_proxy_state(record.state) or
|
|
(record.tags ~= nil and type(record.tags) ~= 'table') or
|
|
(record.indexKeys ~= nil and type(record.indexKeys) ~= 'table') then
|
|
return false
|
|
end
|
|
if not valid_integer(record.port) or record.port <= 0 or record.port > 65535 or
|
|
not valid_integer(record.createdAtMs) or record.createdAtMs <= 0 or
|
|
not valid_integer(record.expiresAtMs) or record.expiresAtMs <= 0 or
|
|
not valid_integer(record.usableUntilMs) or record.usableUntilMs <= 0 or
|
|
record.usableUntilMs > record.expiresAtMs or
|
|
not valid_integer(record.latencyNs) or record.latencyNs < 0 or
|
|
not valid_integer(record.maxConcurrency) or record.maxConcurrency < 0 or
|
|
(record.lastCheckedAtMs ~= nil and
|
|
(not valid_integer(record.lastCheckedAtMs) or record.lastCheckedAtMs < 0)) or
|
|
(record.lastSuccessAtMs ~= nil and
|
|
(not valid_integer(record.lastSuccessAtMs) or record.lastSuccessAtMs < 0)) then
|
|
return false
|
|
end
|
|
for _, index_key in ipairs(record.indexKeys or {}) do
|
|
if type(index_key) ~= 'string' or index_key == '' or
|
|
not string.find(index_key, '{activity}', 1, true) then
|
|
return false
|
|
end
|
|
end
|
|
for tag_key, tag_value in pairs(record.tags or {}) do
|
|
if type(tag_key) ~= 'string' or type(tag_value) ~= 'string' then
|
|
return false
|
|
end
|
|
end
|
|
return true
|
|
end
|
|
|
|
local function empty_result_record(expires_at_ms)
|
|
return '{"version":1,"requestDigest":' .. cjson.encode(request_digest) ..
|
|
',"expiresAtMs":' .. tostring(expires_at_ms) ..
|
|
',"result":{"requested":' .. tostring(requested) .. ',"returned":0,"items":[]}}'
|
|
end
|
|
|
|
if requested == 0 then
|
|
local expires_at_ms = now_ms + idempotency_ttl_ms
|
|
local record = empty_result_record(expires_at_ms)
|
|
if has_idempotency then
|
|
redis.call('SET', idempotency_key, record)
|
|
redis.call('PEXPIREAT', idempotency_key, expires_at_ms)
|
|
end
|
|
return finish({status = 'ok', requestDigest = request_digest, record = record}, expires_at_ms)
|
|
end
|
|
|
|
local driver_key = available_key
|
|
local driver_size = redis.call('ZCARD', available_key)
|
|
for index = 12, #KEYS do
|
|
local size = redis.call('ZCARD', KEYS[index])
|
|
if size < driver_size then
|
|
driver_key = KEYS[index]
|
|
driver_size = size
|
|
end
|
|
end
|
|
|
|
local candidate_ids = redis.call('ZREVRANGE', driver_key, 0, max_candidate_scan - 1)
|
|
local matches_found = {}
|
|
local required_matches = requested + reserve
|
|
local scanned = 0
|
|
|
|
for _, proxy_id in ipairs(candidate_ids) do
|
|
scanned = scanned + 1
|
|
local raw = redis.call('HGET', records_key, proxy_id)
|
|
if not raw then
|
|
redis.call('ZREM', driver_key, proxy_id)
|
|
redis.call('ZREM', available_key, proxy_id)
|
|
else
|
|
local decoded, record = pcall(cjson.decode, raw)
|
|
if not decoded or not valid_proxy_record(proxy_id, record) then
|
|
remove_available(proxy_id, decoded and record or nil)
|
|
elseif tonumber(record.expiresAtMs) <= now_ms then
|
|
remove_proxy(proxy_id)
|
|
else
|
|
local owned = (record.ownerWorkerId and record.ownerWorkerId ~= '') or redis.call('HEXISTS', owners_key, proxy_id) == 1
|
|
local usable = record.state == 'AVAILABLE' and not owned and tonumber(record.usableUntilMs) > now_ms
|
|
if not usable then
|
|
remove_available(proxy_id, record)
|
|
else
|
|
local tags = record.tags or {}
|
|
local health_fresh = max_health_age_ms == 0 or
|
|
(record.lastCheckedAtMs and now_ms - tonumber(record.lastCheckedAtMs) <= max_health_age_ms)
|
|
local ttl_eligible = tonumber(record.expiresAtMs) - now_ms >= min_remaining_ttl_ms
|
|
if health_fresh and ttl_eligible and
|
|
matches(filters.protocols, protocol_filter, record.scheme) and
|
|
matches(filters.regions, region_filter, tags.region or '') and
|
|
matches(filters.carriers, carrier_filter, tags.carrier or '') and
|
|
matches(filters.upstreams, upstream_filter, record.sourceUpstream) then
|
|
matches_found[#matches_found + 1] = {id = proxy_id, record = record}
|
|
if #matches_found >= required_matches then
|
|
break
|
|
end
|
|
end
|
|
end
|
|
end
|
|
end
|
|
end
|
|
|
|
if #matches_found < required_matches and driver_size > scanned then
|
|
return finish({status = 'unavailable', requestDigest = request_digest})
|
|
end
|
|
|
|
local available_count = #matches_found - reserve
|
|
if available_count < 0 then
|
|
available_count = 0
|
|
end
|
|
if fulfillment == 'allOrNothing' and available_count < requested then
|
|
return finish({status = 'insufficient', requestDigest = request_digest})
|
|
end
|
|
|
|
local selected_count = requested
|
|
if selected_count > available_count then
|
|
selected_count = available_count
|
|
end
|
|
local result_items = cjson.decode('[]')
|
|
local earliest_expiry_ms = nil
|
|
|
|
for index = 1, selected_count do
|
|
local selected = matches_found[index]
|
|
local record = selected.record
|
|
remove_available(selected.id, record)
|
|
if is_managed(record.state) then
|
|
decrement_inventory(record.sourceUpstream)
|
|
end
|
|
decrement_state(record.sourceUpstream, record.state)
|
|
record.state = 'EXTRACTED'
|
|
increment_state(record.sourceUpstream, record.state)
|
|
local encoded = cjson.encode(record)
|
|
redis.call('HSET', records_key, selected.id, encoded)
|
|
|
|
local tags = record.tags or {}
|
|
local item = {
|
|
id = record.id,
|
|
protocol = record.scheme,
|
|
host = record.host,
|
|
port = record.port,
|
|
username = record.username,
|
|
password = record.password,
|
|
region = tags.region,
|
|
carrier = tags.carrier,
|
|
upstream = record.sourceUpstream,
|
|
ownerWorkerId = record.ownerWorkerId,
|
|
state = 'EXTRACTED',
|
|
expiresAtMs = record.expiresAtMs,
|
|
lastCheckedAtMs = record.lastCheckedAtMs,
|
|
}
|
|
result_items[#result_items + 1] = item
|
|
local hard_expiry_ms = tonumber(record.expiresAtMs)
|
|
if not earliest_expiry_ms or hard_expiry_ms < earliest_expiry_ms then
|
|
earliest_expiry_ms = hard_expiry_ms
|
|
end
|
|
end
|
|
|
|
local result_expiry_ms = now_ms + idempotency_ttl_ms
|
|
if earliest_expiry_ms and earliest_expiry_ms < result_expiry_ms then
|
|
result_expiry_ms = earliest_expiry_ms
|
|
end
|
|
local result_value = {
|
|
requested = requested,
|
|
returned = selected_count,
|
|
items = result_items,
|
|
}
|
|
if selected_count > 0 then
|
|
result_value.extractedAtMs = now_ms
|
|
end
|
|
local result_record
|
|
if selected_count == 0 then
|
|
result_record = empty_result_record(result_expiry_ms)
|
|
else
|
|
result_record = cjson.encode({
|
|
version = 1,
|
|
requestDigest = request_digest,
|
|
expiresAtMs = result_expiry_ms,
|
|
result = result_value,
|
|
})
|
|
end
|
|
|
|
if has_idempotency then
|
|
redis.call('SET', idempotency_key, result_record)
|
|
redis.call('PEXPIREAT', idempotency_key, result_expiry_ms)
|
|
end
|
|
return finish({status = 'ok', requestDigest = request_digest, record = result_record}, result_expiry_ms)
|