diff --git a/kong/globalpatches.lua b/kong/globalpatches.lua index c3782f0c8a0f..1a65d4479ee2 100644 --- a/kong/globalpatches.lua +++ b/kong/globalpatches.lua @@ -409,6 +409,7 @@ return function(options) local seeded = {} local randomseed = math.randomseed + _G.math.native_randomseed = randomseed _G.math.randomseed = function() local pid = ngx.worker.pid() local id diff --git a/kong/resty/dns_client/utils.lua b/kong/resty/dns_client/utils.lua index a9347ade39c0..42df70edc994 100644 --- a/kong/resty/dns_client/utils.lua +++ b/kong/resty/dns_client/utils.lua @@ -3,6 +3,7 @@ local utils = require("kong.resty.dns.utils") +local math_random = math.random local table_insert = table.insert local DEFAULT_HOSTS_FILE = "/etc/hosts" @@ -70,14 +71,7 @@ function _M.get_rr_ans(answers) end --- based on the Nginx's SWRR algorithm -local function swrr_init(answers) - for _, answer in ipairs(answers) do - answer.cw = 0 -- current weight - end -end - - +-- based on the Nginx's SWRR algorithm and lua-resty-balancer local function swrr_next(answers) local total = 0 local best = nil -- best answer in answers[] @@ -97,6 +91,17 @@ local function swrr_next(answers) end +local function swrr_init(answers) + for _, answer in ipairs(answers) do + answer.cw = 0 -- current weight + end + -- random start + for _ = 1, math_random(#answers) do + swrr_next(answers) + end +end + + -- gather all records with the lowest priority into one array (answers.l) -- and return it local function filter_lowest_priority_answers(answers) @@ -112,7 +117,7 @@ local function filter_lowest_priority_answers(answers) end end - answer.l = l + answers.l = l return l end diff --git a/spec/01-unit/30-new-dns-client/01-utils_spec.lua b/spec/01-unit/30-new-dns-client/01-utils_spec.lua index 52d13b67f768..74438a18c086 100644 --- a/spec/01-unit/30-new-dns-client/01-utils_spec.lua +++ b/spec/01-unit/30-new-dns-client/01-utils_spec.lua @@ -1,17 +1,4 @@ local utils = require "kong.resty.dns_client.utils" -local splitlines = require("pl.stringx").splitlines -local writefile = require("pl.utils").writefile -local tempfilename = require("pl.path").tmpname - -local sleep -if ngx then - gettime = ngx.now -- luacheck: ignore - sleep = ngx.sleep -else - local socket = require("socket") - gettime = socket.gettime -- luacheck: ignore - sleep = socket.sleep -end describe("[utils]", function () @@ -92,4 +79,105 @@ describe("[utils]", function () end) end) + describe("round robin getion", function () + + local function get_and_count(answers, n, get_ans) + local count = {} + for _ = 1, n do + local answer = get_ans(answers) + count[answer.target] = (count[answer.target] or 0) + 1 + end + return count + end + + it("rr", function () + local answers = { + { target = "1" }, -- 25% + { target = "2" }, -- 25% + { target = "3" }, -- 25% + { target = "4" }, -- 25% + } + local count = get_and_count(answers, 100, utils.get_rr_ans) + assert.same(count, { ["1"] = 25, ["2"] = 25, ["3"] = 25, ["4"] = 25 }) + end) + + it("swrr", function () + -- simple one + local answers = { + { target = "w5-p10-a", weight = 5, priority = 10, }, -- hit 100% + } + local count = get_and_count(answers, 20, utils.get_wrr_ans) + assert.same(count, { ["w5-p10-a"] = 20 }) + + -- only get the lowest priority + local answers = { + { target = "w5-p10-a", weight = 5, priority = 10, }, -- hit 50% + { target = "w5-p20", weight = 5, priority = 20, }, -- hit 0% + { target = "w5-p10-b", weight = 5, priority = 10, }, -- hit 50% + { target = "w0-p10", weight = 0, priority = 10, }, -- hit 0% + } + local count = get_and_count(answers, 20, utils.get_wrr_ans) + assert.same(count, { ["w5-p10-a"] = 10, ["w5-p10-b"] = 10 }) + + -- weight: 6, 3, 1 + local answers = { + { target = "w6", weight = 6, priority = 10, }, -- hit 60% + { target = "w3", weight = 3, priority = 10, }, -- hit 30% + { target = "w1", weight = 1, priority = 10, }, -- hit 10% + } + local count = get_and_count(answers, 100 * 1000, utils.get_wrr_ans) + assert.same(count, { ["w6"] = 60000, ["w3"] = 30000, ["w1"] = 10000 }) + + -- random start + _G.math.native_randomseed(9975098) -- math.randomseed() ignores @seed + local answers1 = { + { target = "1", weight = 1, priority = 10, }, + { target = "2", weight = 1, priority = 10, }, + { target = "3", weight = 1, priority = 10, }, + { target = "4", weight = 1, priority = 10, }, + } + local answers2 = { + { target = "1", weight = 1, priority = 10, }, + { target = "2", weight = 1, priority = 10, }, + { target = "3", weight = 1, priority = 10, }, + { target = "4", weight = 1, priority = 10, }, + } + + local a1 = utils.get_wrr_ans(answers1) + local a2 = utils.get_wrr_ans(answers2) + assert.not_equal(a1.target, a2.target) + + -- weight 0 + local answers = { + { target = "w0", weight = 0, priority = 10, }, + { target = "w1", weight = 1, priority = 10, }, -- hit 100% + { target = "w2", weight = 0, priority = 10, }, + { target = "w3", weight = 0, priority = 10, }, + } + local count = get_and_count(answers, 100, utils.get_wrr_ans) + assert.same(count, { ["w1"] = 100 }) + + -- weight 0 and lowest priority + local answers = { + { target = "w0-a", weight = 0, priority = 0, }, -- hit 100% + { target = "w1", weight = 1, priority = 10, }, + { target = "w0-b", weight = 0, priority = 0, }, + { target = "w0-c", weight = 0, priority = 0, }, + } + local count = get_and_count(answers, 100, utils.get_wrr_ans) + assert.same(count, { ["w0-a"] = 100 }) + + -- all weights are 0 + local answers = { + { target = "1", weight = 0, priority = 10, }, + { target = "2", weight = 0, priority = 10, }, + { target = "3", weight = 0, priority = 10, }, + { target = "4", weight = 0, priority = 10, }, + } + local count = get_and_count(answers, 100, utils.get_wrr_ans) + assert.same(count, { ["1"] = 100 }) + + end) + end) + end)