Adding support for traffic shaping for more than one alternative backends

rootfs/etc/nginx/lua/balancer.lua

@@ -191,18 +191,8 @@ local function sync_backends()
   backends_last_synced_at = raw_backends_last_synced_at
 end
 
-local function route_to_alternative_balancer(balancer)
-  if balancer.is_affinitized(balancer) then
-    -- If request is already affinitized to a primary balancer, keep the primary balancer.
-    return false
-  end
-
-  if not balancer.alternative_backends then
-    return false
-  end
-
-  -- TODO: support traffic shaping for n > 1 alternative backends
-  local backend_name = balancer.alternative_backends[1]
+local function is_alternative_backend_suitable(backend_name)
+  -- check whether an alternative backend is suitable for the request context
   if not backend_name then
     ngx.log(ngx.ERR, "empty alternative backend")
     return false
@@ -234,7 +224,7 @@ local function route_to_alternative_balancer(balancer)
   local header = ngx.var["http_" .. target_header]
   if header then
     if traffic_shaping_policy.headerValue
-	   and #traffic_shaping_policy.headerValue > 0 then
+           and #traffic_shaping_policy.headerValue > 0 then
       if traffic_shaping_policy.headerValue == header then
         return true
       end
@@ -276,6 +266,28 @@ local function route_to_alternative_balancer(balancer)
   return false
 end
 
+local function get_alternative_or_original_balancer(balancer)
+  -- checks whether an alternative backend should be used
+  -- the first backend matching the request context is returned
+  -- if no suitable alterntative backends are found, the original is returned
+  if balancer.is_affinitized(balancer) then
+    -- If request is already affinitized to a primary balancer, keep the primary balancer.
+    return balancer
+  end
+
+  if not balancer.alternative_backends then
+    return balancer
+  end
+
+  for _, backend_name in ipairs(balancer.alternative_backends) do
+    if is_alternative_backend_suitable(backend_name) then
+      return balancers[backend_name]
+    end
+  end
+
+  return balancer
+end
+
 local function get_balancer_by_upstream_name(upstream_name)
   return balancers[upstream_name]
 end
@@ -292,16 +304,9 @@ local function get_balancer()
     return nil
   end
 
-  if route_to_alternative_balancer(balancer) then
-    local alternative_backend_name = balancer.alternative_backends[1]
-    ngx.var.proxy_alternative_upstream_name = alternative_backend_name
-
-    balancer = balancers[alternative_backend_name]
-  end
-
-  ngx.ctx.balancer = balancer
+  ngx.ctx.balancer = get_alternative_or_original_balancer(balancer)
 
-  return balancer
+  return ngx.ctx.balancer
 end
 
 function _M.init_worker()
@@ -376,7 +381,7 @@ end
 setmetatable(_M, {__index = {
   get_implementation = get_implementation,
   sync_backend = sync_backend,
-  route_to_alternative_balancer = route_to_alternative_balancer,
+  get_alternative_or_original_balancer = get_alternative_or_original_balancer,
   get_balancer = get_balancer,
   get_balancer_by_upstream_name = get_balancer_by_upstream_name,
 }})

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

@@ -150,14 +150,17 @@ describe("Balancer", function()
     end)
   end)
 
-  describe("route_to_alternative_balancer()", function()
-    local backend, _primaryBalancer
+  describe("get_alternative_or_original_balancer()", function()
+    local backend1, backend2, _primaryBalancer
 
     before_each(function()
-      backend = backends[1]
+      backend1 = backends[1]
+      backend2 = backends[1]
+      backend2.name = "second-router-production-web-80"
       _primaryBalancer = {
         alternative_backends = {
-          backend.name,
+          backend1.name,
+          backend2.name,
         }
       }
       mock_ngx({ var = { request_uri = "/" } })
@@ -172,50 +175,72 @@ describe("Balancer", function()
         end
       end)
 
-      it("returns false when no trafficShapingPolicy is set", function()
-        balancer.sync_backend(backend)
-        assert.equal(false, balancer.route_to_alternative_balancer(_primaryBalancer))
+      it("returns _primaryBalancer when no trafficShapingPolicy is set", function()
+        balancer.sync_backend(backend1)
+        assert.equal(_primaryBalancer, balancer.get_alternative_or_original_balancer(_primaryBalancer))
       end)
 
-      it("returns false when no alternative backends is set", function()
-        backend.trafficShapingPolicy.weight = 100
-        balancer.sync_backend(backend)
+      it("returns _primaryBalancer when no alternative backends is set", function()
+        backend1.trafficShapingPolicy.weight = 100
+        backend2.trafficShapingPolicy.weight = 100
+        balancer.sync_backend(backend1)
+        balancer.sync_backend(backend2)
         _primaryBalancer.alternative_backends = nil
-        assert.equal(false, balancer.route_to_alternative_balancer(_primaryBalancer))
+        assert.equal(_primaryBalancer, balancer.get_alternative_or_original_balancer(_primaryBalancer))
       end)
 
-      it("returns false when alternative backends name does not match", function()
-        backend.trafficShapingPolicy.weight = 100
-        balancer.sync_backend(backend)
-        _primaryBalancer.alternative_backends[1] = "nonExistingBackend"
-        assert.equal(false, balancer.route_to_alternative_balancer(_primaryBalancer))
+      it("returns _primaryBalancer when alternative backends name does not match", function()
+        backend1.trafficShapingPolicy.weight = 100
+        backend2.trafficShapingPolicy.weight = 100
+        balancer.sync_backend(backend1)
+        balancer.sync_backend(backend2)
+        _primaryBalancer.alternative_backends = {"nonExistingBackend", "secondNonExistingBackend"}
+        assert.equal(_primaryBalancer, balancer.get_alternative_or_original_balancer(_primaryBalancer))
       end)
 
       describe("canary by weight", function()
-        it("returns true when weight is 100", function()
-          backend.trafficShapingPolicy.weight = 100
-          balancer.sync_backend(backend)
-          assert.equal(true, balancer.route_to_alternative_balancer(_primaryBalancer))
+        it("returns first backend when weight is 100", function()
+          backend1.trafficShapingPolicy.weight = 100
+          balancer.sync_backend(backend1)
+          assert.equal(backend1, balancer.get_alternative_or_original_balancer(_primaryBalancer))
+        end)
+
+        it("returns second backend when weight is 100", function()
+          backend2.trafficShapingPolicy.weight = 100
+          balancer.sync_backend(backend2)
+          assert.equal(backend2, balancer.get_alternative_or_original_balancer(_primaryBalancer))
+        end)
+
+        it("returns _primaryBalancer when weight is 0", function()
+          backend1.trafficShapingPolicy.weight = 0
+          backend2.trafficShapingPolicy.weight = 0
+          balancer.sync_backend(backend1)
+          balancer.sync_backend(backend2)
+          assert.equal(_primaryBalancer, balancer.get_alternative_or_original_balancer(_primaryBalancer))
         end)
 
-        it("returns false when weight is 0", function()
-          backend.trafficShapingPolicy.weight = 0
-          balancer.sync_backend(backend)
-          assert.equal(false, balancer.route_to_alternative_balancer(_primaryBalancer))
+        it("returns first backend when weight is 1000 and weight total is 1000", function()
+          backend1.trafficShapingPolicy.weight = 1000
+          backend1.trafficShapingPolicy.weightTotal = 1000
+          balancer.sync_backend(backend1)
+          assert.equal(backend1, balancer.get_alternative_or_original_balancer(_primaryBalancer))
         end)
 
-        it("returns true when weight is 1000 and weight total is 1000", function()
-          backend.trafficShapingPolicy.weight = 1000
-          backend.trafficShapingPolicy.weightTotal = 1000
-          balancer.sync_backend(backend)
-          assert.equal(true, balancer.route_to_alternative_balancer(_primaryBalancer))
+        it("returns second backend when weight is 1000 and weight total is 1000", function()
+          backend2.trafficShapingPolicy.weight = 1000
+          backend2.trafficShapingPolicy.weightTotal = 1000
+          balancer.sync_backend(backend2)
+          assert.equal(backend2, balancer.get_alternative_or_original_balancer(_primaryBalancer))
         end)
 
-        it("returns false when weight is 0 and weight total is 1000", function()
-          backend.trafficShapingPolicy.weight = 1000
-          backend.trafficShapingPolicy.weightTotal = 1000
-          balancer.sync_backend(backend)
-          assert.equal(true, balancer.route_to_alternative_balancer(_primaryBalancer))
+        it("returns _primaryBalancer when weight is 0 and weight total is 1000", function()
+          backend1.trafficShapingPolicy.weight = 0
+          backend1.trafficShapingPolicy.weightTotal = 1000
+          backend2.trafficShapingPolicy.weight = 0
+          backend2.trafficShapingPolicy.weightTotal = 1000
+          balancer.sync_backend(backend1)
+          balancer.sync_backend(backend2)
+          assert.equal(_primaryBalancer, balancer.get_alternative_or_original_balancer(_primaryBalancer))
         end)
       end)
 
@@ -252,10 +277,11 @@ describe("Balancer", function()
               ["cookie_" .. test_pattern.request_cookie_name] = test_pattern.request_cookie_value,
               request_uri = "/"
             }})
-            backend.trafficShapingPolicy.cookie = "canaryCookie"
-            balancer.sync_backend(backend)
+            backend2.trafficShapingPolicy.cookie = "canaryCookie"
+            balancer.sync_backend(backend2)
+            local _expected_result = test_pattern.expected_result and backend2 or _primaryBalancer
             assert.message("\nTest data pattern: " .. test_pattern.case_title)
-              .equal(test_pattern.expected_result, balancer.route_to_alternative_balancer(_primaryBalancer))
+              .equal(_expected_result, balancer.get_alternative_or_original_balancer(_primaryBalancer))
             reset_ngx()
           end
         end)
@@ -329,11 +355,12 @@ describe("Balancer", function()
               ["http_" .. test_pattern.request_header_name] = test_pattern.request_header_value,
               request_uri = "/"
             }})
-            backend.trafficShapingPolicy.header = test_pattern.header_name
-            backend.trafficShapingPolicy.headerValue = test_pattern.header_value
-            balancer.sync_backend(backend)
+            backend2.trafficShapingPolicy.header = test_pattern.header_name
+            backend2.trafficShapingPolicy.headerValue = test_pattern.header_value
+            balancer.sync_backend(backend2)
+            local _expected_result = test_pattern.expected_result and backend2 or _primaryBalancer
             assert.message("\nTest data pattern: " .. test_pattern.case_title)
-              .equal(test_pattern.expected_result, balancer.route_to_alternative_balancer(_primaryBalancer))
+              .equal(_expected_result, balancer.get_alternative_or_original_balancer(_primaryBalancer))
             reset_ngx()
           end
         end)
@@ -345,41 +372,68 @@ describe("Balancer", function()
     describe("affinitized", function()
 
       before_each(function()
-        mock_ngx({ var = { request_uri = "/", proxy_upstream_name = backend.name } })
-        balancer.sync_backend(backend)
+        mock_ngx({ var = { request_uri = "/", proxy_upstream_name = backend2.name } })
+        balancer.sync_backend(backend2)
       end)
 
-      it("returns false if request is affinitized to primary backend", function()
+      it("returns _primaryBalancer if request is affinitized to primary backend", function()
         _primaryBalancer.is_affinitized = function (_)
           return true
         end
 
-        local alternativeBalancer = balancer.get_balancer_by_upstream_name(backend.name)
+        local alternativeBalancer1 = balancer.get_balancer_by_upstream_name(backend1.name)
+        local alternativeBalancer2 = balancer.get_balancer_by_upstream_name(backend2.name)
+
+        local primarySpy = spy.on(_primaryBalancer, "is_affinitized")
+        local alternativeSpy1 = spy.on(alternativeBalancer1, "is_affinitized")
+        local alternativeSpy2 = spy.on(alternativeBalancer2, "is_affinitized")
+
+        assert.equal(_primaryBalancer , balancer.get_alternative_or_original_balancer(_primaryBalancer))
+        assert.spy(_primaryBalancer.is_affinitized).was_called()
+        assert.spy(alternativeBalancer1.is_affinitized).was_not_called()
+        assert.spy(alternativeBalancer2.is_affinitized).was_not_called()
+      end)
+
+      it("returns first alternative balancer if request is affinitized to alternative backend", function()
+        _primaryBalancer.is_affinitized = function (_)
+          return false
+        end
+
+        local alternativeBalancer1 = balancer.get_balancer_by_upstream_name(backend1.name)
+        local alternativeBalancer2 = balancer.get_balancer_by_upstream_name(backend2.name)
+        alternativeBalancer1.is_affinitized = function (_)
+          return true
+        end
 
         local primarySpy = spy.on(_primaryBalancer, "is_affinitized")
-        local alternativeSpy = spy.on(alternativeBalancer, "is_affinitized")
+        local alternativeSpy1 = spy.on(alternativeBalancer1, "is_affinitized")
+        local alternativeSpy2 = spy.on(alternativeBalancer2, "is_affinitized")
 
-        assert.is_false(balancer.route_to_alternative_balancer(_primaryBalancer))
+        assert.equal(alternativeBalancer1 , balancer.get_alternative_or_original_balancer(_primaryBalancer))
         assert.spy(_primaryBalancer.is_affinitized).was_called()
-        assert.spy(alternativeBalancer.is_affinitized).was_not_called()
+        assert.spy(alternativeBalancer1.is_affinitized).was_called()
+        assert.spy(alternativeBalancer2.is_affinitized).was_not_called()
       end)
 
-      it("returns true if request is affinitized to alternative backend", function()
+      it("returns second alternative balancer if request is affinitized to alternative backend", function()
         _primaryBalancer.is_affinitized = function (_)
           return false
         end
 
-        local alternativeBalancer = balancer.get_balancer_by_upstream_name(backend.name)
-        alternativeBalancer.is_affinitized = function (_)
+        local alternativeBalancer1 = balancer.get_balancer_by_upstream_name(backend1.name)
+        local alternativeBalancer2 = balancer.get_balancer_by_upstream_name(backend2.name)
+        alternativeBalancer2.is_affinitized = function (_)
           return true
         end
 
         local primarySpy = spy.on(_primaryBalancer, "is_affinitized")
-        local alternativeSpy = spy.on(alternativeBalancer, "is_affinitized")
+        local alternativeSpy1 = spy.on(alternativeBalancer1, "is_affinitized")
+        local alternativeSpy2 = spy.on(alternativeBalancer2, "is_affinitized")
 
-        assert.is_true(balancer.route_to_alternative_balancer(_primaryBalancer))
+        assert.equal(alternativeBalancer2 , balancer.get_alternative_or_original_balancer(_primaryBalancer))
         assert.spy(_primaryBalancer.is_affinitized).was_called()
-        assert.spy(alternativeBalancer.is_affinitized).was_called()
+        assert.spy(alternativeBalancer1.is_affinitized).was_called()
+        assert.spy(alternativeBalancer2.is_affinitized).was_called()
       end)
 
     end)