Converted sticky session balancers into separate classes.

rootfs/etc/nginx/lua/balancer.lua

@@ -6,7 +6,8 @@ local configuration = require("configuration")
 local round_robin = require("balancer.round_robin")
 local chash = require("balancer.chash")
 local chashsubset = require("balancer.chashsubset")
-local sticky = require("balancer.sticky")
+local sticky_balanced = require("balancer.sticky_balanced")
+local sticky_persistent = require("balancer.sticky_persistent")
 local ewma = require("balancer.ewma")
 
 -- measured in seconds
@@ -19,7 +20,8 @@ local IMPLEMENTATIONS = {
   round_robin = round_robin,
   chash = chash,
   chashsubset = chashsubset,
-  sticky = sticky,
+  sticky_balanced = sticky_balanced,
+  sticky_persistent = sticky_persistent,
   ewma = ewma,
 }
 
@@ -30,7 +32,11 @@ local function get_implementation(backend)
   local name = backend["load-balance"] or DEFAULT_LB_ALG
 
   if backend["sessionAffinityConfig"] and backend["sessionAffinityConfig"]["name"] == "cookie" then
-    name = "sticky"
+    if backend["sessionAffinityConfig"]["mode"] == 'persistent' then
+      name = "sticky_persistent"
+    else
+      name = "sticky_balanced"
+    end
   elseif backend["upstreamHashByConfig"] and backend["upstreamHashByConfig"]["upstream-hash-by"] then
     if backend["upstreamHashByConfig"]["upstream-hash-by-subset"] then
       name = "chashsubset"

rootfs/etc/nginx/lua/balancer/sticky.lua

@@ -1,5 +1,3 @@
-local affinity_balanced = require("affinity.balanced")
-local affinity_persistent = require("affinity.persistent")
 local balancer_resty = require("balancer.resty")
 local util = require("util")
 local ck = require("resty.cookie")
@@ -10,48 +8,24 @@ local string_format = string.format
 local ngx_log = ngx.log
 local INFO = ngx.INFO
 
-local _M = balancer_resty:new({ name = "sticky" })
+local _M = balancer_resty:new()
 local DEFAULT_COOKIE_NAME = "route"
 
-
 function _M.cookie_name(self)
   return self.cookie_session_affinity.name or DEFAULT_COOKIE_NAME
 end
 
-local function init_affinity_mode(self, backend)
+function _M.new(self)
-  local mode = backend["sessionAffinityConfig"]["mode"] or 'balanced'
-
-  -- set default mode to 'balanced' for backwards compatibility
-  if mode == nil or mode == '' then
-    mode = 'balanced'
-  end
-
-  self.affinity_mode = mode
-
-  if mode == 'persistent' then
-    return affinity_persistent:new(self, backend)
-  end
-
-  -- default is 'balanced' for backwards compatibility
-  if mode ~= 'balanced' then
-    ngx.log(ngx.WARN, string.format("Invalid affinity mode '%s'! Using 'balanced' as a default.", mode))
-  end
-
-  return affinity_balanced:new(self, backend)
-end
-
-function _M.new(self, backend)
   local o = {
-    instance = nil,
+    alternative_backends = nil,
-    affinity_mode = nil,
+    cookie_session_affinity = nil,
-    traffic_shaping_policy = backend.trafficShapingPolicy,
+    traffic_shaping_policy = nil
-    alternative_backends = backend.alternativeBackends,
-    cookie_session_affinity = backend["sessionAffinityConfig"]["cookieSessionAffinity"]
   }
+
   setmetatable(o, self)
   self.__index = self
-  
+
-  return init_affinity_mode(o, backend)
+  return o
 end
 
 function _M.get_cookie(self)
@@ -112,34 +86,8 @@ local function get_failed_upstreams()
   return indexed_upstream_addrs
 end
 
---- get_routing_key gets the current routing key from the cookie
--- @treturn string, string The routing key and an error message if an error occured.
-function _M.get_routing_key(self)
-  -- interface method to get the routing key from the cookie
-  -- has to be overridden by an affinity mode
-  ngx.log(ngx.ERR, "[BUG] Failed to get routing key as no implementation has been provided!")
-  return nil, nil
-end
-
---- set_routing_key sets the current routing key on the cookie
--- @tparam string key The routing key to set on the cookie.
-function _M.set_routing_key(self, key)
-  -- interface method to set the routing key on the cookie
-  -- has to be overridden by an affinity mode
-  ngx.log(ngx.ERR, "[BUG] Failed to set routing key as no implementation has been provided!")
-end
-
---- pick_new_upstream picks a new upstream while ignoring the given failed upstreams.
--- @tparam {[string]=boolean} A table of upstreams to ignore where the key is the endpoint and the value a boolean.
--- @treturn string, string The endpoint and its key.
-function _M.pick_new_upstream(self, failed_upstreams)
-  -- interface method to get a new upstream
-  -- has to be overridden by an affinity mode
-  ngx.log(ngx.ERR, "[BUG] Failed to pick new upstream as no implementation has been provided!")
-  return nil, nil
-end
-
 local function should_set_cookie(self)
+
   if self.cookie_session_affinity.locations and ngx.var.host then
     local locs = self.cookie_session_affinity.locations[ngx.var.host]
     if locs == nil then
@@ -193,19 +141,11 @@ function _M.balance(self)
 end
 
 function _M.sync(self, backend)
-  local changed = false
-
-  -- check and reinit affinity mode before syncing the balancer which will reinit the nodes
-  if self.affinity_mode ~= backend.sessionAffinityConfig.mode then
-    changed = true
-    init_affinity_mode(self, backend)
-  end
-
   -- reload balancer nodes
   balancer_resty.sync(self, backend)
 
   -- Reload the balancer if any of the annotations have changed.
-  changed = changed or not util.deep_compare(
+  local changed = not util.deep_compare(
     self.cookie_session_affinity,
     backend.sessionAffinityConfig.cookieSessionAffinity
   )
@@ -216,6 +156,8 @@ function _M.sync(self, backend)
 
   ngx_log(INFO, string_format("[%s] nodes have changed for backend %s", self.name, backend.name))
 
+  self.traffic_shaping_policy = backend.trafficShapingPolicy
+  self.alternative_backends = backend.alternativeBackends
   self.cookie_session_affinity = backend.sessionAffinityConfig.cookieSessionAffinity
 end
 

rootfs/etc/nginx/lua/balancer/sticky_balanced.lua

@@ -4,13 +4,12 @@
 -- will lose their session, where c is the current number of pods and n is the new number of 
 -- pods.
 --
--- This class extends/implements the abstract class balancer.sticky.
+local balancer_sticky = require("balancer.sticky")
---
 local math = require("math")
 local resty_chash = require("resty.chash")
 local util = require("util")
 
-local _M = {}
+local _M = balancer_sticky:new()
 
 -- Consider the situation of N upstreams one of which is failing.
 -- Then the probability to obtain failing upstream after M iterations would be close to (1/N)**M.
@@ -18,18 +17,33 @@ local _M = {}
 -- which is much better then ~10**(-3) for 10 iterations.
 local MAX_UPSTREAM_CHECKS_COUNT = 20
 
-local function get_routing_key(self)
+function _M.new(self, backend)
+  local nodes = util.get_nodes(backend.endpoints)
+
+  local o = {
+    name = "sticky_balanced",
+    instance = resty_chash:new(nodes)
+  }
+
+  setmetatable(o, self)
+  self.__index = self
+  
+  balancer_sticky.sync(o, backend)
+
+  return o
+end
+
+function _M.get_routing_key(self)
   return self:get_cookie(), nil
 end
 
-local function set_routing_key(self, key)
+function _M.set_routing_key(self, key)
 	self:set_cookie(key)
 end
 
-local function pick_new_upstream(self, failed_upstreams)
+function _M.pick_new_upstream(self, failed_upstreams)
   for i = 1, MAX_UPSTREAM_CHECKS_COUNT do
     local key = string.format("%s.%s.%s", ngx.now() + i, ngx.worker.pid(), math.random(999999))
-
     local new_upstream = self.instance:find(key)
 
     if not failed_upstreams[new_upstream] then
@@ -40,18 +54,4 @@ local function pick_new_upstream(self, failed_upstreams)
   return nil, nil
 end
 
-function _M.new(self, sticky_balancer, backend)
-  local o = sticky_balancer or {}
-  
-  local nodes = util.get_nodes(backend.endpoints)
-
-  -- override sticky.balancer methods
-  o.instance = resty_chash:new(nodes)
-  o.get_routing_key = get_routing_key
-  o.set_routing_key = set_routing_key
-  o.pick_new_upstream = pick_new_upstream
-
-  return sticky_balancer
-end
-
 return _M

rootfs/etc/nginx/lua/balancer/sticky_persistent.lua

@@ -3,12 +3,30 @@
 -- The drawback of this mode is that when scaling up a deployment, sessions will not
 -- be rebalanced.
 --
+local balancer_sticky = require("balancer.sticky")
 local util = require("util")
 local util_nodemap = require("util.nodemap")
 
-local _M = {}
+local _M = balancer_sticky:new()
 
-local function get_routing_key(self)
+function _M.new(self, backend)
+  local nodes = util.get_nodes(backend.endpoints)
+  local hash_salt = backend["name"]
+
+  local o = {
+    name = "sticky_persistent",
+    instance = util_nodemap:new(nodes, hash_salt)
+  }
+
+  setmetatable(o, self)
+  self.__index = self
+  
+  balancer_sticky.sync(o, backend)
+
+  return o
+end
+
+function _M.get_routing_key(self)
   local cookie_value = self:get_cookie()
 
   if cookie_value then
@@ -26,28 +44,13 @@ local function get_routing_key(self)
   return nil, nil
 end
 
-local function set_routing_key(self, key)
+function _M.set_routing_key(self, key)
   local value = string.format("%s.%s.%s", ngx.now(), ngx.worker.pid(), key)
   self:set_cookie(value);
 end
 
-local function pick_new_upstream(self, failed_upstreams)
+function _M.pick_new_upstream(self, failed_upstreams)
   return self.instance:random_except(failed_upstreams)
 end
 
-function _M.new(self, sticky_balancer, backend)
-  local o = sticky_balancer or {}
-  
-  local nodes = util.get_nodes(backend.endpoints)
-  local hash_salt = backend["name"]
-
-  -- override sticky.balancer methods
-  o.instance = util_nodemap:new(nodes, hash_salt)
-  o.get_routing_key = get_routing_key
-  o.set_routing_key = set_routing_key
-  o.pick_new_upstream = pick_new_upstream
-
-  return sticky_balancer
-end
-
 return _M

rootfs/etc/nginx/lua/test/balancer/sticky_test.lua

@@ -1,4 +1,5 @@
-local sticky = require("balancer.sticky")
+local sticky_balanced = require("balancer.sticky_balanced")
+local sticky_persistent = require("balancer.sticky_persistent")
 local cookie = require("resty.cookie")
 local util = require("util")
 
@@ -57,21 +58,27 @@ describe("Sticky", function()
 
   describe("new(backend)", function()
     context("when backend specifies cookie name", function()
-      it("returns an instance containing the corresponding cookie name", function()
+      local function test(sticky)
         local sticky_balancer_instance = sticky:new(test_backend)
         local test_backend_cookie_name = test_backend.sessionAffinityConfig.cookieSessionAffinity.name
         assert.equal(sticky_balancer_instance:cookie_name(), test_backend_cookie_name)
-      end)
+      end
+
+      it("returns an instance containing the corresponding cookie name", function() test(sticky_balanced) end)
+      it("returns an instance containing the corresponding cookie name", function() test(sticky_persistent) end)
     end)
 
     context("when backend does not specify cookie name", function()
-      it("returns an instance with 'route' as cookie name", function()
+      local function test(sticky)
         local temp_backend = util.deepcopy(test_backend)
         temp_backend.sessionAffinityConfig.cookieSessionAffinity.name = nil
         local sticky_balancer_instance = sticky:new(temp_backend)
         local default_cookie_name = "route"
         assert.equal(sticky_balancer_instance:cookie_name(), default_cookie_name)
-      end)
+      end
+
+      it("returns an instance with 'route' as cookie name", function() test(sticky_balanced) end)
+      it("returns an instance with 'route' as cookie name", function() test(sticky_persistent) end)
     end)
   end)
 
@@ -79,8 +86,10 @@ describe("Sticky", function()
     local mocked_cookie_new = cookie.new
 
     before_each(function()
-      package.loaded["balancer.sticky"] = nil
+      package.loaded["balancer.sticky_balanced"] = nil
-      sticky = require("balancer.sticky")
+      package.loaded["balancer.sticky_persistent"] = nil
+      sticky_balanced = require("balancer.sticky_balanced")
+      sticky_persistent = require("balancer.sticky_persistent")
     end)
 
     after_each(function()
@@ -88,13 +97,17 @@ describe("Sticky", function()
     end)
 
     context("when client doesn't have a cookie set and location is in cookie_locations", function()
-      it("picks an endpoint for the client", function()
+
+      local function test_pick_endpoint(sticky)
         local sticky_balancer_instance = sticky:new(test_backend)
         local peer = sticky_balancer_instance:balance()
-        assert.equal(peer, test_backend_endpoint)
+        assert.equal(test_backend_endpoint, peer)
-      end)
+      end
 
-      it("sets a cookie on the client", function()
+      it("picks an endpoint for the client", function() test_pick_endpoint(sticky_balanced) end)
+      it("picks an endpoint for the client", function() test_pick_endpoint(sticky_persistent) end)
+
+      local function test_set_cookie(sticky)
         local s = {}
         cookie.new = function(self)
           local cookie_instance = {
@@ -117,9 +130,12 @@ describe("Sticky", function()
         local sticky_balancer_instance = sticky:new(b)
         assert.has_no.errors(function() sticky_balancer_instance:balance() end)
         assert.spy(s).was_called()
-      end)
+      end
 
-      it("sets a secure cookie on the client when being in ssl mode", function()
+      it("sets a cookie on the client", function() test_set_cookie(sticky_balanced) end)
+      it("sets a cookie on the client", function() test_set_cookie(sticky_persistent) end)
+
+      local function test_set_ssl_cookie(sticky)
         ngx.var.https = "on"
         local s = {}
         cookie.new = function(self)
@@ -143,10 +159,18 @@ describe("Sticky", function()
         local sticky_balancer_instance = sticky:new(b)
         assert.has_no.errors(function() sticky_balancer_instance:balance() end)
         assert.spy(s).was_called()
+      end
+
+      it("sets a secure cookie on the client when being in ssl mode", function()
+          test_set_ssl_cookie(sticky_balanced)
+      end)
+      it("sets a secure cookie on the client when being in ssl mode", function()
+        test_set_ssl_cookie(sticky_persistent)
       end)
     end)
 
-    context("when client doesn't have a cookie set and cookie_locations contains a matching wildcard location", function()
+    context("when client doesn't have a cookie set and cookie_locations contains a matching wildcard location",
+    function()
       before_each(function ()
         ngx.var.host = "dev.test.com"
       end)
@@ -154,7 +178,7 @@ describe("Sticky", function()
         ngx.var.host = "test.com"
       end)
 
-      it("sets a cookie on the client", function()
+      local function test(sticky)
         local s = {}
         cookie.new = function(self)
           local cookie_instance = {
@@ -178,17 +202,24 @@ describe("Sticky", function()
         local sticky_balancer_instance = sticky:new(b)
         assert.has_no.errors(function() sticky_balancer_instance:balance() end)
         assert.spy(s).was_called()
-      end)
+      end 
+
+      it("sets a cookie on the client", function() test(sticky_balanced) end)
+      it("sets a cookie on the client", function() test(sticky_persistent) end)
     end)
 
     context("when client doesn't have a cookie set and location not in cookie_locations", function()
-      it("picks an endpoint for the client", function()
+
+      local function test_pick_endpoint(sticky)
         local sticky_balancer_instance = sticky:new(test_backend)
         local peer = sticky_balancer_instance:balance()
         assert.equal(peer, test_backend_endpoint)
-      end)
+      end
 
-      it("does not set a cookie on the client", function()
+      it("picks an endpoint for the client", function() test_pick_endpoint(sticky_balanced) end)
+      it("picks an endpoint for the client", function() test_pick_endpoint(sticky_persistent) end)
+
+      local function test_no_cookie(sticky)
         local s = {}
         cookie.new = function(self)
           local cookie_instance = {
@@ -207,11 +238,15 @@ describe("Sticky", function()
         local sticky_balancer_instance = sticky:new(get_test_backend())
         assert.has_no.errors(function() sticky_balancer_instance:balance() end)
         assert.spy(s).was_not_called()
-      end)
+      end
+
+      it("does not set a cookie on the client", function() test_no_cookie(sticky_balanced) end)
+      it("does not set a cookie on the client", function() test_no_cookie(sticky_persistent) end)
     end)
 
     context("when client has a cookie set", function()
-      it("does not set a cookie", function()
+
+      local function test_no_cookie(sticky)
         local s = {}
         cookie.new = function(self)
           local return_obj = {
@@ -224,17 +259,23 @@ describe("Sticky", function()
         local sticky_balancer_instance = sticky:new(test_backend)
         assert.has_no.errors(function() sticky_balancer_instance:balance() end)
         assert.spy(s).was_not_called()
-      end)
+      end
 
-      it("returns the correct endpoint for the client", function()
+      it("does not set a cookie", test_no_cookie(sticky_balanced))
+      it("does not set a cookie", test_no_cookie(sticky_persistent))
+
+      local function test_correct_endpoint(sticky)
         local sticky_balancer_instance = sticky:new(test_backend)
         local peer = sticky_balancer_instance:balance()
         assert.equal(peer, test_backend_endpoint)
-      end)
+      end
+
+      it("returns the correct endpoint for the client", function() test_correct_endpoint(sticky_balanced) end)
+      it("returns the correct endpoint for the client", function() test_correct_endpoint(sticky_persistent) end)
     end)
   end)
 
-  local function get_several_test_backends(option)
+  local function get_several_test_backends(change_on_failure)
     return {
       name = "access-router-production-web-80",
       endpoints = {
@@ -243,11 +284,10 @@ describe("Sticky", function()
       },
       sessionAffinityConfig = {
         name = "cookie",
-        mode = option["mode"],
         cookieSessionAffinity = {
           name = "test_name",
           hash = "sha1",
-          change_on_failure = option["change_on_failure"],
+          change_on_failure = change_on_failure,
           locations = { ['test.com'] = {'/'} }
         }
       },
@@ -258,52 +298,58 @@ describe("Sticky", function()
     local mocked_cookie_new = cookie.new
 
     before_each(function()
-      package.loaded["balancer.sticky"] = nil
+      package.loaded["balancer.sticky_balanced"] = nil
-      sticky = require("balancer.sticky")
+      package.loaded["balancer.sticky_persistent"] = nil
+      sticky_balanced = require("balancer.sticky_balanced")
+      sticky_persistent = require("balancer.sticky_persistent")
+      mock_ngx({ var = { location_path = "/", host = "test.com" } })
     end)
 
     after_each(function()
-      cookie.new = mocked_cookie_new
+      reset_ngx()
     end)
 
     context("when request to upstream fails", function()
-      it("changes upstream when change_on_failure option is true", function()
-        local options = {
-          {["change_on_failure"] = false, ["mode"] = nil},
-          {["change_on_failure"] = false, ["mode"] = 'balanced'},
-          {["change_on_failure"] = false, ["mode"] = 'persistent'},
-          {["change_on_failure"] = true, ["mode"] = nil},
-          {["change_on_failure"] = true, ["mode"] = 'balanced'},
-          {["change_on_failure"] = true, ["mode"] = 'persistent'}
-        }
 
-        for _, option in ipairs(options) do
+      local function test(sticky, change_on_failure)
-          local sticky_balancer_instance = sticky:new(get_several_test_backends(option))
+        local sticky_balancer_instance = sticky:new(get_several_test_backends(change_on_failure))
 
-          local old_upstream = sticky_balancer_instance:balance()
+        local old_upstream = sticky_balancer_instance:balance()
-          assert.is.Not.Nil(old_upstream)
+        assert.is.Not.Nil(old_upstream)
-          for _ = 1, 100 do
+        for _ = 1, 100 do
-            -- make sure upstream doesn't change on subsequent calls of balance()
+          -- make sure upstream doesn't change on subsequent calls of balance()
-            assert.equal(old_upstream, sticky_balancer_instance:balance())
+          assert.equal(old_upstream, sticky_balancer_instance:balance())
-          end
+        end
 
-          -- simulate request failure
+        -- simulate request failure
-          sticky_balancer_instance.get_last_failure = function()
+        sticky_balancer_instance.get_last_failure = function()
-            return "failed"
+          return "failed"
-          end
+        end
-          _G.ngx.var.upstream_addr = old_upstream
+        _G.ngx.var.upstream_addr = old_upstream
 
-          for _ = 1, 100 do
+        for _ = 1, 100 do
-            local new_upstream = sticky_balancer_instance:balance()
+          local new_upstream = sticky_balancer_instance:balance()
-            if option["change_on_failure"] == false then
+          if change_on_failure == false then
-              -- upstream should be the same inspite of error, if change_on_failure option is false
+            -- upstream should be the same inspite of error, if change_on_failure option is false
-              assert.equal(new_upstream, old_upstream)
+            assert.equal(new_upstream, old_upstream)
-            else
+          else
-              -- upstream should change after error, if change_on_failure option is true
+            -- upstream should change after error, if change_on_failure option is true
-              assert.not_equal(new_upstream, old_upstream)
+            assert.not_equal(new_upstream, old_upstream)
-            end
           end
         end
+      end
+
+      it("changes upstream when change_on_failure option is true", function()
+        test(sticky_balanced, 'balanced', true)
+      end)
+      it("changes upstream when change_on_failure option is true", function()
+        test(sticky_balanced, 'balanced', false)
+      end)
+      it("changes upstream when change_on_failure option is true", function()
+        test(sticky_persistent, 'balanced', true)
+      end)
+      it("changes upstream when change_on_failure option is true", function()
+        test(sticky_persistent, 'balanced', false)
       end)
     end)
   end)

rootfs/etc/nginx/lua/test/balancer_test.lua

@@ -23,9 +23,9 @@ local function reset_expected_implementations()
     ["access-router-production-web-80"] = package.loaded["balancer.round_robin"],
     ["my-dummy-app-1"] = package.loaded["balancer.round_robin"],
     ["my-dummy-app-2"] = package.loaded["balancer.chash"],
-    ["my-dummy-app-3"] = package.loaded["balancer.sticky"],
+    ["my-dummy-app-3"] = package.loaded["balancer.sticky_persistent"],
     ["my-dummy-app-4"] = package.loaded["balancer.ewma"],
-    ["my-dummy-app-5"] = package.loaded["balancer.sticky"],
+    ["my-dummy-app-5"] = package.loaded["balancer.sticky_balanced"]
   }
 end
 
@@ -55,7 +55,7 @@ local function reset_backends()
     },
     {
       name = "my-dummy-app-3", ["load-balance"] = "ewma",
-      sessionAffinityConfig = { name = "cookie", cookieSessionAffinity = { name = "route" } }
+      sessionAffinityConfig = { name = "cookie", mode = 'persistent', cookieSessionAffinity = { name = "route" } }
     },
     { name = "my-dummy-app-4", ["load-balance"] = "ewma", },
     {