Merge pull request #4048 from fedunineyu/change-upstream-on-error-with-sticky-session

Change upstream on error when sticky session balancer is used

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

@@ -3,6 +3,8 @@ local resty_chash = require("resty.chash")
 local util = require("util")
 local ck = require("resty.cookie")
 local math = require("math")
+local ngx_balancer = require("ngx.balancer")
+local split = require("util.split")
 
 local string_format = string.format
 local ngx_log = ngx.log
@@ -11,6 +13,12 @@ local INFO = ngx.INFO
 local _M = balancer_resty:new({ factory = resty_chash, name = "sticky" })
 local DEFAULT_COOKIE_NAME = "route"
 
+-- 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.
+-- For the worst case (2 upstreams; 20 iterations) it would be ~10**(-6)
+-- which is much better then ~10**(-3) for 10 iterations.
+local MAX_UPSTREAM_CHECKS_COUNT = 20
+
 function _M.cookie_name(self)
   return self.cookie_session_affinity.name or DEFAULT_COOKIE_NAME
 end
@@ -63,6 +71,10 @@ local function set_cookie(self, value)
   end
 end
 
+function _M.get_last_failure()
+  return ngx_balancer.get_last_failure()
+end
+
 function _M.balance(self)
   local cookie, err = ck:new()
   if not cookie then
@@ -70,24 +82,75 @@ function _M.balance(self)
     return
   end
 
-  local key = cookie:get(self:cookie_name())
-  if not key then
-    key = string.format("%s.%s.%s", ngx.now(), ngx.worker.pid(), math.random(999999))
+  -- upstream_from_cookie: upstream which is pointed by sticky cookie
+  local upstream_from_cookie = nil
 
-    if self.cookie_session_affinity.locations then
-      local locs = self.cookie_session_affinity.locations[ngx.var.host]
-      if locs ~= nil then
-        for _, path in pairs(locs) do
-          if ngx.var.location_path == path then
-            set_cookie(self, key)
-            break
-          end
-        end
-      end
+  local key = cookie:get(self:cookie_name())
+  if key then
+    upstream_from_cookie = self.instance:find(key)
+  end
+
+  -- get state of the previous attempt
+  local state_name = self.get_last_failure()
+
+  if upstream_from_cookie ~= nil then
+    -- use previous upstream if this is the first attempt or previous attempt succeeded
+    -- or ingress is configured to ignore previous request result
+    if state_name == nil or self.cookie_session_affinity.changeonfailure == "false" then
+      return upstream_from_cookie
     end
   end
 
-  return self.instance:find(key)
+  -- failed_upstream: upstream which failed during previous attempt
+  local failed_upstream = nil
+
+  -- If previous attempt failed recent upstream can be obtained from ngx.var.upstream_addr.
+  -- Do nothing if ingress is configured to ignore previous request result.
+  if state_name ~= nil and self.cookie_session_affinity.changeonfailure == "true" then
+    local upstream_addr = ngx.var.upstream_addr
+    failed_upstream = split.get_last_value(upstream_addr)
+
+    if failed_upstream == nil then
+      ngx.log(ngx.ERR, string.format("failed to get failed_upstream from upstream_addr (%s)", upstream_addr))
+    end
+  end
+
+  -- new_upstream: upstream which is pointed by new key
+  local new_upstream = nil
+
+  -- generate new upstream key if sticky cookie not set or previous attempt failed
+  for _ = 1, MAX_UPSTREAM_CHECKS_COUNT do
+    key = string.format("%s.%s.%s", ngx.now(), ngx.worker.pid(), math.random(999999))
+
+    new_upstream = self.instance:find(key)
+
+    if failed_upstream ~= new_upstream then
+      -- set cookie only when we get NOT THE SAME upstream
+      if upstream_from_cookie ~= new_upstream then
+        if self.cookie_session_affinity.locations then
+          local locs = self.cookie_session_affinity.locations[ngx.var.host]
+          if locs ~= nil then
+            for _, path in pairs(locs) do
+              if ngx.var.location_path == path then
+                set_cookie(self, key)
+                break
+              end
+            end
+          end
+        end
+
+      end
+
+      -- new upstream was obtained; return it to the balancer
+      do return new_upstream end
+    end
+
+    -- generated key points to the failed upstream; try another key
+  end
+
+  ngx.log(ngx.WARN, string.format("failed to get new upstream; using upstream %s", new_upstream))
+
+  return new_upstream
 end
 
 function _M.sync(self, backend)

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

@@ -164,7 +164,7 @@ describe("Sticky", function()
           s = spy.on(cookie_instance, "set")
           return cookie_instance, false
         end
-        local sticky_balancer_instance = sticky:new(get_test_backend())        
+        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)
@@ -193,4 +193,73 @@ describe("Sticky", function()
       end)
     end)
   end)
+
+  local function get_several_test_backends(changeOnFailure)
+    return {
+      name = "access-router-production-web-80",
+      endpoints = {
+        { address = "10.184.7.40", port = "8080", maxFails = 0, failTimeout = 0 },
+        { address = "10.184.7.41", port = "8080", maxFails = 0, failTimeout = 0 },
+      },
+      sessionAffinityConfig = {
+        name = "cookie",
+        cookieSessionAffinity = { name = "test_name", hash = "sha1", changeonfailure = changeOnFailure }
+      },
+    }
+  end
+
+  describe("balance() after error", function()
+    local mocked_cookie_new = cookie.new
+
+    before_each(function()
+      package.loaded["balancer.sticky"] = nil
+      sticky = require("balancer.sticky")
+    end)
+
+    after_each(function()
+      cookie.new = mocked_cookie_new
+    end)
+
+    context("when request to upstream fails", function()
+      it("changes upstream when changeOnFailure option is true", function()
+        -- create sticky cookie
+        cookie.new = function(self)
+          local return_obj = {
+            set = function(v) return false, nil end,
+            get = function(k) return "" end,
+          }
+          return return_obj, false
+        end
+
+        local options = {'false', 'true'}
+
+        for _, option in ipairs(options) do
+          local sticky_balancer_instance = sticky:new(get_several_test_backends(option))
+
+          local old_upstream = sticky_balancer_instance:balance()
+          for _ = 1, 100 do
+            -- make sure upstream doesn't change on subsequent calls of balance()
+            assert.equal(old_upstream, sticky_balancer_instance:balance())
+          end
+
+          -- simulate request failure
+          sticky_balancer_instance.get_last_failure = function()
+            return "failed"
+          end
+          _G.ngx.var = { upstream_addr = old_upstream }
+
+          for _ = 1, 100 do
+            local new_upstream = sticky_balancer_instance:balance()
+            if option == 'false' then
+              -- upstream should be the same inspite of error, if changeOnFailure option is false
+              assert.equal(new_upstream, old_upstream)
+            else
+              -- upstream should change after error, if changeOnFailure option is true
+              assert.not_equal(new_upstream, old_upstream)
+            end
+          end
+        end
+      end)
+    end)
+  end)
 end)

rootfs/etc/nginx/lua/util/split.lua

@@ -16,6 +16,12 @@ function _M.get_first_value(var)
   return t[1]
 end
 
+function _M.get_last_value(var)
+  local t = _M.split_upstream_var(var) or {}
+  if #t == 0 then return nil end
+  return t[#t]
+end
+
 -- http://nginx.org/en/docs/http/ngx_http_upstream_module.html#example
 -- CAVEAT: nginx is giving out : instead of , so the docs are wrong
 -- 127.0.0.1:26157 : 127.0.0.1:26157 , ngx.var.upstream_addr