Merge pull request #4169 from Shopify/refactor-sticky-balancer

simplify sticky balancer and fix a bug

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

@@ -75,6 +75,48 @@ function _M.get_last_failure()
   return ngx_balancer.get_last_failure()
 end
 
+local function get_failed_upstreams()
+  local indexed_upstream_addrs = {}
+  local upstream_addrs = split.split_upstream_var(ngx.var.upstream_addr) or {}
+
+  for _, addr in ipairs(upstream_addrs) do
+    indexed_upstream_addrs[addr] = true
+  end
+
+  return indexed_upstream_addrs
+end
+
+local function pick_new_upstream(self)
+  local failed_upstreams = get_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
+      return new_upstream, key
+    end
+  end
+
+  return nil, nil
+end
+
+local function should_set_cookie(self)
+  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
+          return true
+        end
+      end
+    end
+  end
+
+  return false
+end
+
 function _M.balance(self)
   local cookie, err = ck:new()
   if not cookie then
@@ -82,73 +124,26 @@ function _M.balance(self)
     return
   end
 
-  -- upstream_from_cookie: upstream which is pointed by sticky cookie
-  local upstream_from_cookie = nil
+  local upstream_from_cookie
 
   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()
+  local last_failure = self.get_last_failure()
+  local should_pick_new_upstream = last_failure ~= nil and self.cookie_session_affinity.change_on_failure or upstream_from_cookie == nil
 
-  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 not self.cookie_session_affinity.change_on_failure then
+  if not should_pick_new_upstream then
     return upstream_from_cookie
   end
-  end
-
-  -- 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.change_on_failure 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
 
+  local new_upstream, key = pick_new_upstream(self)
+  if not new_upstream then
     ngx.log(ngx.WARN, string.format("failed to get new upstream; using upstream %s", new_upstream))
+  elseif should_set_cookie(self) then
+    set_cookie(self, key)
+  end
 
   return new_upstream
 end

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

@@ -16,12 +16,6 @@ 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