ability to create a new balancer using plugins

rootfs/etc/nginx/lua/balancer.lua

@@ -17,6 +17,7 @@ local tostring = tostring
 local pairs = pairs
 local math = math
 local ngx = ngx
+local type = type
 
 -- measured in seconds
 -- for an Nginx worker to pick up the new list of upstream peers
@@ -294,6 +295,27 @@ local function get_balancer()
   return balancer
 end
 
+function _M.register_implementation(name, implementation)
+  if not name or #name == 0 then
+    return false, "name is required"
+  end
+  if not implementation or type(implementation) ~= "table" then
+    return false, "implementation is required"
+  end
+  if type(implementation.new) ~= "function" or type(implementation.sync) ~= "function" or
+    type(implementation.balance) ~= "function" then
+
+    return false, "`new`, `sync` and `balance` functions must be implemented"
+  end
+  if IMPLEMENTATIONS[name] then
+    return false, "implementation with name " .. name .. " already exists"
+  end
+
+  IMPLEMENTATIONS[name] = implementation
+
+  return true, nil
+end
+
 function _M.init_worker()
   -- when worker starts, sync non ExternalName backends without delay
   sync_backends()

rootfs/etc/nginx/lua/plugins.lua

@@ -7,7 +7,9 @@ local INFO = ngx.INFO
 local ERR = ngx.ERR
 local pcall = pcall
 
-local _M = {}
+local _M = {
+  balancer_implementations = {},
+}
 local MAX_NUMBER_OF_PLUGINS = 20
 local plugins = {}
 
@@ -23,6 +25,10 @@ local function load_plugin(name)
   if (plugin.name == nil or plugin.name == '') then
     plugin.name = name
   end
+
+  if plugin["balancer_implementation"] then
+    _M.balancer_implementations[name] = plugin.balancer_implementation()
+  end
   plugins[index + 1] = plugin
 end
 

rootfs/etc/nginx/lua/plugins/README.md

@@ -14,8 +14,9 @@ By defining functions with the following names, you can run your custom Lua code
 
  - `init_worker`: useful for initializing some data per Nginx worker process
  - `rewrite`: useful for modifying request, changing headers, redirection, dropping request, doing authentication etc
+ - `balance`: when a plugin implements `balancer_implementation` function, it will be registered as a load balancer implementation. This plugin has to implement the balancer interface. Check [`hello_world`](./hello_world/main.lua) for an example implementation.
  - `header_filter`: this is called when backend response header is received, it is useful for modifying response headers
- - `body_filter`: this is called when response body is received, it is useful for logging response body 
+ - `body_filter`: this is called when response body is received, it is useful for logging response body
  - `log`: this is called when request processing is completed and a response is delivered to the client
 
 Check this [`hello_world`](https://github.com/kubernetes/ingress-nginx/tree/main/rootfs/etc/nginx/lua/plugins/hello_world) plugin as a simple example or refer to [OpenID Connect integration](https://github.com/ElvinEfendi/ingress-nginx-openidc/tree/master/rootfs/etc/nginx/lua/plugins/openidc) for more advanced usage.

rootfs/etc/nginx/lua/plugins/hello_world/main.lua

@@ -1,4 +1,5 @@
 local ngx = ngx
+local setmetatable = setmetatable
 
 local _M = {}
 
@@ -10,4 +11,33 @@ function _M.rewrite()
   end
 end
 
+function _M.balancer_implementation()
+  -- An example balancer implementation that always returns the first endpoint.
+  -- Used for demonstration and testing purposes only.
+  return {
+    name = "hello_world",
+    new = function(self, backend)
+      local o = {
+        endpoints = backend.endpoints,
+        traffic_shaping_policy = backend.trafficShapingPolicy,
+        alternative_backends = backend.alternativeBackends,
+      }
+      setmetatable(o, self)
+      self.__index = self
+      return o
+    end,
+    is_affinitized = function(_) return false end,
+    after_balance = function(_) end,
+    sync = function(self, backend)
+      self.endpoints = backend.endpoints
+      self.traffic_shaping_policy = backend.trafficShapingPolicy
+      self.alternative_backends = backend.alternativeBackends
+    end,
+    balance = function(self)
+      local endpoint = self.endpoints[1]
+      return endpoint.address .. ":" .. endpoint.port
+    end,
+  }
+end
+
 return _M

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

@@ -533,4 +533,45 @@ describe("Balancer", function()
     end)
 
   end)
+
+  describe("register_implementation", function ()
+    it("registers a new balancer implementation", function ()
+      local new_implementation = {
+        name = "new_implementation",
+        new = function(self, backend)
+          local o = {
+            endpoints = backend.endpoints,
+            traffic_shaping_policy = backend.trafficShapingPolicy,
+            alternative_backends = backend.alternativeBackends,
+          }
+          setmetatable(o, self)
+          self.__index = self
+          return o
+        end,
+        is_affinitized = function(_) return false end,
+        after_balance = function(_) end,
+        sync = function(self, backend)
+          self.endpoints = backend.endpoints
+          self.traffic_shaping_policy = backend.trafficShapingPolicy
+          self.alternative_backends = backend.alternativeBackends
+        end,
+        balance = function (self)
+          local endpoint = self.endpoints[1]
+          return endpoint.address .. ":" .. endpoint.port
+        end,
+      }
+
+      balancer.register_implementation("new_implementation", new_implementation)
+
+      local backend = {
+        name = "dummy",
+        ["load-balance"] = "new_implementation",
+        endpoints = {
+          { address = "1.1.1.1", port = "8080", maxFails = 0, failTimeout = 0 },
+        }
+      }
+      local implementation = balancer.get_implementation(backend)
+      assert.equal(new_implementation, implementation)
+    end)
+  end)
 end)

rootfs/etc/nginx/lua/test/plugins_test.lua

@@ -20,4 +20,19 @@ describe("plugins", function()
       assert.are.same(plugins_to_mock, called_plugins)
     end)
   end)
-end)
+
+  describe("#init", function()
+    it("marks hello_world plugin as a balancer implementation", function()
+      local plugins = require("plugins")
+      local balancer_implementations = plugins.balancer_implementations
+      assert.is_nil(balancer_implementations["hello_world"])
+      assert.has_no.errors(function()
+        plugins.init({"hello_world"})
+      end)
+      assert.is_table(balancer_implementations["hello_world"])
+      assert.is_function(balancer_implementations["hello_world"].new)
+      assert.is_function(balancer_implementations["hello_world"].sync)
+      assert.is_function(balancer_implementations["hello_world"].balance)
+    end)
+  end)
+end)

rootfs/etc/nginx/template/nginx.tmpl

@@ -122,6 +122,12 @@ http {
         end
         -- load all plugins that'll be used here
         plugins.init({ {{ range  $idx, $plugin := $cfg.Plugins }}{{ if $idx }},{{ end }}{{ $plugin | quote }}{{ end }} })
+        for name, implementation in pairs(plugins.balancer_implementations) do
+            local ok, err = balancer.register_implementation(name, implementation)
+            if not ok then
+                error("failed to register balancer implementation: " .. err)
+            end
+        end
     }
 
     init_worker_by_lua_block {
@@ -789,7 +795,7 @@ stream {
     lua_package_path "/etc/nginx/lua/?.lua;/etc/nginx/lua/vendor/?.lua;;";
 
     lua_shared_dict tcp_udp_configuration_data 5M;
-    
+
     {{ buildResolvers $cfg.Resolver $cfg.DisableIpv6DNS }}
 
     init_by_lua_block {

test/e2e/settings/plugins.go

@@ -22,6 +22,7 @@ import (
 	"strings"
 
 	"github.com/onsi/ginkgo/v2"
+	"github.com/stretchr/testify/assert"
 	"k8s.io/ingress-nginx/test/e2e/framework"
 )
 
@@ -52,4 +53,33 @@ var _ = framework.IngressNginxDescribe("plugins", func() {
 			Status(http.StatusOK).
 			Body().Contains("x-hello-world=1")
 	})
+
+	ginkgo.It("registers hello_world as a custom balancer implementation", func() {
+		f.SetNginxConfigMapData(map[string]string{
+			"plugins":      "hello_world",
+			"load-balance": "hello_world",
+		})
+
+		host := "example.com"
+		f.EnsureIngress(framework.NewSingleIngress(host, "/", host, f.Namespace, framework.EchoService, 80, nil))
+
+		f.WaitForNginxConfiguration(
+			func(server string) bool {
+				return strings.Contains(server, fmt.Sprintf("server_name %v", host)) &&
+					strings.Contains(server, `plugins.init({ "hello_world" })`) &&
+					strings.Contains(server, `local ok, err = balancer.register_implementation(name, implementation)`)
+			})
+
+		algorithm, err := f.GetLbAlgorithm(framework.EchoService, 80)
+		assert.Nil(ginkgo.GinkgoT(), err)
+		assert.Equal(ginkgo.GinkgoT(), algorithm, "hello_world")
+
+		f.HTTPTestClient().
+			GET("/").
+			WithHeader("Host", host).
+			WithHeader("User-Agent", "hello").
+			Expect().
+			Status(http.StatusOK).
+			Body().Contains("x-hello-world=1")
+	})
 })