Exclude draining endpoints from new sticky sessions

Author: gulachek

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

@@ -3,6 +3,7 @@ local ck = require("resty.cookie")
 local ngx_balancer = require("ngx.balancer")
 local split = require("util.split")
 local same_site = require("util.same_site")
+local util = require("util")
 
 local ngx = ngx
 local pairs = pairs
@@ -24,7 +25,8 @@ function _M.new(self)
     alternative_backends = nil,
     cookie_session_affinity = nil,
     traffic_shaping_policy = nil,
-    backend_key = nil
+    backend_key = nil,
+    draining_endpoints = nil
   }
 
   setmetatable(o, self)
@@ -129,14 +131,20 @@ function _M.get_last_failure()
   return ngx_balancer.get_last_failure()
 end
 
-local function get_failed_upstreams()
+local function get_upstreams_to_exclude(self)
   local indexed_upstream_addrs = {}
-  local upstream_addrs = split.split_upstream_var(ngx.var.upstream_addr) or {}
+  local failed_upstream_addrs = split.split_upstream_var(ngx.var.upstream_addr) or {}
 
-  for _, addr in ipairs(upstream_addrs) do
+  for _, addr in ipairs(failed_upstream_addrs) do
     indexed_upstream_addrs[addr] = true
   end
 
+  if self.draining_endpoints then
+    for addr, _ in pairs(self.draining_endpoints) do
+      indexed_upstream_addrs[addr] = true
+    end
+  end
+
   return indexed_upstream_addrs
 end
 
@@ -188,7 +196,7 @@ function _M.balance(self)
 
   local new_upstream
 
-  new_upstream, key = self:pick_new_upstream(get_failed_upstreams())
+  new_upstream, key = self:pick_new_upstream(get_upstreams_to_exclude(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
@@ -202,6 +210,18 @@ function _M.sync(self, backend)
   -- reload balancer nodes
   balancer_resty.sync(self, backend)
 
+  self.draining_endpoints = nil
+  for _, endpoint in pairs(backend.endpoints) do
+    if endpoint.isDraining then
+      if not self.draining_endpoints then
+        self.draining_endpoints = {}
+      end
+
+      local endpoint_string = util.get_endpoint_string(endpoint)
+      self.draining_endpoints[endpoint_string] = true
+    end
+  end
+
   self.traffic_shaping_policy = backend.trafficShapingPolicy
   self.alternative_backends = backend.alternativeBackends
   self.cookie_session_affinity = backend.sessionAffinityConfig.cookieSessionAffinity

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

@@ -304,6 +304,68 @@ describe("Sticky", function()
       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)
+
+    -- Note below that endpoints are only marked this way when persistent
+    describe("when an endpoint is draining", function()
+      it("persists client with cookie to endpoint", function()
+        local s = {}
+        local cookie_jar = {}
+        cookie.new = function(self)
+          local cookie_instance = {
+            set = function(self, ck)
+              cookie_jar[ck.key] = ck.value
+              return true, nil
+            end,
+            get = function(self, k) return cookie_jar[k] end,
+          }
+          s = spy.on(cookie_instance, "set")
+          return cookie_instance, false
+        end
+
+        local b = get_test_backend()
+        b.sessionAffinityConfig.cookieSessionAffinity.locations = {}
+        b.sessionAffinityConfig.cookieSessionAffinity.locations["test.com"] = {"/"}
+
+        local expectedPeer = b.endpoints[1]
+
+        local sticky_balancer_instance = sticky_persistent:new(b)
+        local peer = sticky_balancer_instance:balance()
+        assert.equal(peer, test_backend_endpoint)
+
+        expectedPeer.isDraining = true
+        sticky_balancer_instance:sync(b)
+        sticky_balancer_instance.TESTING = true
+        peer = sticky_balancer_instance:balance()
+        assert.equal(peer, test_backend_endpoint)
+      end)
+
+      it("does not route client without cookie to endpoint", function()
+        local s = {}
+        local cookie_jar = {}
+        cookie.new = function(self)
+          local cookie_instance = {
+            set = function(self, ck)
+              cookie_jar[ck.key] = ck.value
+              return true, nil
+            end,
+            get = function(self, k) return cookie_jar[k] end,
+          }
+          s = spy.on(cookie_instance, "set")
+          return cookie_instance, false
+        end
+
+        local b = get_test_backend()
+        b.sessionAffinityConfig.cookieSessionAffinity.locations = {}
+        b.sessionAffinityConfig.cookieSessionAffinity.locations["test.com"] = {"/"}
+
+        local expectedPeer = b.endpoints[1]
+        expectedPeer.isDraining = true
+
+        local sticky_balancer_instance = sticky_persistent:new(b)
+        local peer = sticky_balancer_instance:balance()
+        assert.equal(peer, nil)
+      end)
+    end)
   end)
 
   local function get_several_test_backends(change_on_failure)