random start in swrr_init() and add test cases

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

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
 

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)