feature: allow plugins to handler balance phase. (#299)

*  feature: allow plugins to handler balance phase.

lua/apisix.lua

@@ -75,6 +75,29 @@ local function run_plugin(phase, plugins, api_ctx)
         return api_ctx
     end
 
+    if phase == "balancer" then
+        local balancer_name = api_ctx.balancer_name
+        local balancer_plugin = api_ctx.balancer_plugin
+        if balancer_name and balancer_plugin then
+            local phase_fun = balancer_plugin[phase]
+            phase_fun(balancer_plugin, api_ctx)
+            return api_ctx
+        end
+
+        for i = 1, #plugins, 2 do
+            local phase_fun = plugins[i][phase]
+            if phase_fun and
+               (not balancer_name or balancer_name == plugins[i].name) then
+                phase_fun(plugins[i + 1], api_ctx)
+                if api_ctx.balancer_name == plugins[i].name then
+                    api_ctx.balancer_plugin = plugins[i]
+                    return api_ctx
+                end
+            end
+        end
+        return api_ctx
+    end
+
     if phase ~= "log" then
         for i = 1, #plugins, 2 do
             local phase_fun = plugins[i][phase]
@@ -214,6 +237,19 @@ function _M.http_balancer_phase()
         return core.response.exit(500)
     end
 
+    -- first time
+    if not api_ctx.balancer_name then
+        run_plugin("balancer", nil, api_ctx)
+        if api_ctx.balancer_name then
+            return
+        end
+    end
+
+    if api_ctx.balancer_name and api_ctx.balancer_name ~= "default" then
+        return run_plugin("balancer", nil, api_ctx)
+    end
+
+    api_ctx.balancer_name = "default"
     load_balancer(api_ctx.matched_route, api_ctx)
 end
 

lua/apisix/http/balancer.lua

@@ -223,10 +223,8 @@ local function pick_server(route, ctx)
     end
 
     local ip, port, err = parse_addr(server)
-    if upstream.checker then
-        ctx.balancer_ip = ip
-        ctx.balancer_port = port
-    end
+    ctx.balancer_ip = ip
+    ctx.balancer_port = port
 
     return ip, port, err
 end

lua/apisix/plugins/example-plugin.lua

@@ -1,5 +1,5 @@
 local core = require("apisix.core")
-
+local balancer = require("ngx.balancer")
 
 -- You can follow this document to write schema:
 -- https://github.com/Tencent/rapidjson/blob/master/bin/draft-04/schema
@@ -10,8 +10,10 @@ local schema = {
         i = {type = "number", minimum = 0},
         s = {type = "string"},
         t = {type = "array", minItems = 1},
+        ip = {type = "string"},
+        port = {type = "integer"},
     },
-    required = {"i"}
+    required = {"i"},
 }
 
 
@@ -48,4 +50,23 @@ function _M.access(conf, ctx)
 end
 
 
+function _M.balancer(conf, ctx)
+    core.log.warn("plugin access phase, conf: ", core.json.encode(conf))
+
+    if not conf.ip then
+        return
+    end
+
+    -- NOTE: update `ctx.balancer_name` is important, APISIX will skip other
+    -- balancer handler.
+    ctx.balancer_name = plugin_name
+
+    local ok, err = balancer.set_current_peer(conf.ip, conf.port)
+    if not ok then
+        core.log.error("failed to set server peer: ", err)
+        return core.response.exit(502)
+    end
+end
+
+
 return _M

t/plugin/example.t

@@ -4,6 +4,7 @@ repeat_each(2);
 no_long_string();
 no_root_location();
 no_shuffle();
+
 run_tests;
 
 __DATA__
@@ -197,3 +198,47 @@ GET /t
 plugin [example-plugin] config: {"i":1,"s":"s","t":[1,2]}
 --- no_error_log
 [error]
+
+
+
+=== TEST 8: set route(id: 1)
+--- config
+    location /t {
+        content_by_lua_block {
+            local t = require("lib.test_admin").test
+            local code, body = t('/apisix/admin/routes/1',
+                 ngx.HTTP_PUT,
+                 [[{
+                        "plugins": {
+                            "example-plugin": {
+                                "i": 11,
+                                "ip": "127.0.0.1",
+                                "port": 1981
+                            }
+                        },
+                        "uri": "/server_port"
+                }]]
+                )
+
+            if code >= 300 then
+                ngx.status = code
+            end
+            ngx.say(body)
+        }
+    }
+--- request
+GET /t
+--- response_body
+passed
+--- no_error_log
+[error]
+
+
+
+=== TEST 9: hit route
+--- request
+GET /server_port
+--- response_body_like eval
+qr/1981/
+--- no_error_log
+[error]