wehub-resource-sync
164 lines
3.7 KiB
Lua
164 lines
3.7 KiB
Lua
|
|
local balancers = require "kong.runloop.balancer.balancers"
|
|
|
|
local random = math.random
|
|
|
|
local MAX_WHEEL_SIZE = 2^32
|
|
|
|
local roundrobin_algorithm = {}
|
|
roundrobin_algorithm.__index = roundrobin_algorithm
|
|
|
|
-- calculate the greater common divisor, used to find the smallest wheel
|
|
-- possible
|
|
local function gcd(a, b)
|
|
if b == 0 then
|
|
return a
|
|
end
|
|
|
|
return gcd(b, a % b)
|
|
end
|
|
|
|
|
|
local function wheel_shuffle(wheel)
|
|
for i = #wheel, 2, -1 do
|
|
local j = random(i)
|
|
wheel[i], wheel[j] = wheel[j], wheel[i]
|
|
end
|
|
return wheel
|
|
end
|
|
|
|
|
|
function roundrobin_algorithm:afterHostUpdate()
|
|
local total_points = 0
|
|
local total_weight = 0
|
|
local divisor = 0
|
|
|
|
local targets = self.balancer.targets or {}
|
|
|
|
-- calculate the gcd to find the proportional weight of each address
|
|
for _, target in ipairs(targets) do
|
|
for _, address in ipairs(target.addresses) do
|
|
local address_weight = address.weight
|
|
divisor = gcd(divisor, address_weight)
|
|
total_weight = total_weight + address_weight
|
|
end
|
|
end
|
|
|
|
self.balancer.totalWeight = total_weight
|
|
if total_weight == 0 then
|
|
ngx.log(ngx.DEBUG, "trying to set a round-robin balancer with no addresses")
|
|
return
|
|
end
|
|
|
|
if divisor > 0 then
|
|
total_points = total_weight / divisor
|
|
end
|
|
|
|
-- add all addresses to the wheel
|
|
local new_wheel = {}
|
|
local idx = 1
|
|
|
|
for _, target in ipairs(targets) do
|
|
for _, address in ipairs(target.addresses) do
|
|
local address_points = address.weight / divisor
|
|
for _ = 1, address_points do
|
|
new_wheel[idx] = address
|
|
idx = idx + 1
|
|
end
|
|
end
|
|
end
|
|
|
|
-- store the shuffled wheel
|
|
self.wheel = wheel_shuffle(new_wheel)
|
|
self.wheelSize = total_points
|
|
end
|
|
|
|
|
|
function roundrobin_algorithm:getPeer(cacheOnly, handle, hashValue)
|
|
if handle then
|
|
-- existing handle, so it's a retry
|
|
handle.retryCount = handle.retryCount + 1
|
|
|
|
else
|
|
-- no handle, so this is a first try
|
|
handle = {} -- self:getHandle() -- no GC specific handler needed
|
|
handle.retryCount = 0
|
|
end
|
|
|
|
local starting_pointer = self.pointer
|
|
local address
|
|
local ip, port, hostname
|
|
|
|
repeat
|
|
self.pointer = self.pointer + 1
|
|
|
|
if self.pointer > self.wheelSize then
|
|
self.pointer = 1
|
|
end
|
|
|
|
address = self.wheel[self.pointer]
|
|
if not address or not address.available or address.disabled then
|
|
goto continue
|
|
end
|
|
|
|
-- address ~= nil and address.available and not address.disabled
|
|
|
|
ip, port, hostname = balancers.getAddressPeer(address, cacheOnly)
|
|
|
|
if ip then
|
|
-- success, update handle
|
|
handle.address = address
|
|
return ip, port, hostname, handle
|
|
end
|
|
|
|
-- no ip, port is an error hint
|
|
|
|
if port == balancers.errors.ERR_DNS_UPDATED then
|
|
-- if healthy we just need to try again
|
|
if self.balancer.healthy then
|
|
goto continue
|
|
end
|
|
|
|
-- not healthy
|
|
return nil, balancers.errors.ERR_BALANCER_UNHEALTHY
|
|
end
|
|
|
|
if port == balancers.errors.ERR_ADDRESS_UNAVAILABLE then
|
|
ngx.log(ngx.DEBUG, "found address but it was unavailable. ",
|
|
" trying next one.")
|
|
goto continue
|
|
end
|
|
|
|
-- an unknown error occurred
|
|
do return nil, port end
|
|
|
|
::continue::
|
|
|
|
until self.pointer == starting_pointer
|
|
|
|
return nil, balancers.errors.ERR_NO_PEERS_AVAILABLE
|
|
end
|
|
|
|
|
|
function roundrobin_algorithm.new(opts)
|
|
assert(type(opts) == "table", "Expected an options table, but got: " .. type(opts))
|
|
|
|
local balancer = opts.balancer
|
|
|
|
local self = setmetatable({
|
|
health_threshold = balancer.health_threshold,
|
|
balancer = balancer,
|
|
|
|
pointer = 1,
|
|
wheelSize = 0,
|
|
maxWheelSize = balancer.maxWheelSize or balancer.wheelSize or MAX_WHEEL_SIZE,
|
|
wheel = {},
|
|
}, roundrobin_algorithm)
|
|
|
|
self:afterHostUpdate()
|
|
|
|
return self
|
|
end
|
|
|
|
return roundrobin_algorithm
|