feat: enable balancer phase for plugins (#4549)

apisix/balancer.lua

@@ -286,8 +286,9 @@ do
 end
 
 
-function _M.run(route, ctx)
+function _M.run(route, ctx, plugin_funcs)
     local server, err
+    local header_changed
 
     if ctx.picked_server then
         -- use the server picked in the access phase
@@ -314,9 +315,16 @@ function _M.run(route, ctx)
             if host ~= ctx.var.upstream_host then
                 -- retried node has a different host
                 ctx.var.upstream_host = host
-                balancer.recreate_request()
+                header_changed = true
             end
         end
+
+    end
+
+    local _, run = plugin_funcs("balancer")
+    -- always recreate request as the request may be changed by plugins
+    if (run or header_changed) and balancer.recreate_request then
+        balancer.recreate_request()
     end
 
     core.log.info("proxy request to ", server.host, ":", server.port)

apisix/init.lua

@@ -556,11 +556,10 @@ local function common_phase(phase_name)
 
     if api_ctx.script_obj then
         script.run(phase_name, api_ctx)
-    else
+        return api_ctx, true
-        plugin.run_plugin(phase_name, nil, api_ctx)
     end
 
-    return api_ctx
+    return plugin.run_plugin(phase_name, nil, api_ctx)
 end
 
 
@@ -711,7 +710,7 @@ function _M.http_balancer_phase()
         return core.response.exit(500)
     end
 
-    load_balancer.run(api_ctx.matched_route, api_ctx)
+    load_balancer.run(api_ctx.matched_route, api_ctx, common_phase)
 end
 
 
@@ -921,7 +920,7 @@ function _M.stream_balancer_phase()
         return ngx_exit(1)
     end
 
-    load_balancer.run(api_ctx.matched_route, api_ctx)
+    load_balancer.run(api_ctx.matched_route, api_ctx, common_phase)
 end
 
 

apisix/plugin.lua

@@ -632,6 +632,7 @@ end
 
 
 function _M.run_plugin(phase, plugins, api_ctx)
+    local plugin_run = false
     api_ctx = api_ctx or ngx.ctx.api_ctx
     if not api_ctx then
         return
@@ -649,6 +650,7 @@ function _M.run_plugin(phase, plugins, api_ctx)
         for i = 1, #plugins, 2 do
             local phase_func = plugins[i][phase]
             if phase_func then
+                plugin_run = true
                 local code, body = phase_func(plugins[i + 1], api_ctx)
                 if code or body then
                     if is_http then
@@ -667,17 +669,18 @@ function _M.run_plugin(phase, plugins, api_ctx)
                 end
             end
         end
-        return api_ctx
+        return api_ctx, plugin_run
     end
 
     for i = 1, #plugins, 2 do
         local phase_func = plugins[i][phase]
         if phase_func then
+            plugin_run = true
             phase_func(plugins[i + 1], api_ctx)
         end
     end
 
-    return api_ctx
+    return api_ctx, plugin_run
 end
 
 

t/debug/debug-mode.t

@@ -236,7 +236,7 @@ passed
 GET /hello
 --- yaml_config eval: $::yaml_config
 --- response_headers
-Apisix-Plugins: response-rewrite, limit-conn, limit-count, response-rewrite
+Apisix-Plugins: response-rewrite, limit-conn, limit-count, response-rewrite, response-rewrite
 --- response_body
 yes
 --- error_log

t/lib/server.lua

@@ -397,6 +397,22 @@ function _M.server_error()
 end
 
 
+function _M.log_request()
+    ngx.log(ngx.WARN, "uri: ", ngx.var.uri)
+    local headers = ngx.req.get_headers()
+
+    local keys = {}
+    for k in pairs(headers) do
+        table.insert(keys, k)
+    end
+    table.sort(keys)
+
+    for _, key in ipairs(keys) do
+        ngx.log(ngx.WARN, key, ": ", headers[key])
+    end
+end
+
+
 function _M.v3_auth_authenticate()
     ngx.log(ngx.WARN, "etcd auth failed!")
 end

t/plugin/serverless.t

@@ -679,3 +679,159 @@ GET /hello
 --- error_log
 default phase: access
 match uri /hello
+
+
+
+=== TEST 23: run in the balancer phase
+--- 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": {
+                        "serverless-pre-function": {
+                            "phase": "balancer",
+                            "functions" : ["return function(conf, ctx) ngx.req.set_header('X-SERVERLESS', ctx.balancer_ip) end"]
+                        }
+                    },
+                    "upstream": {
+                        "nodes": {
+                            "127.0.0.2:1979": 100000,
+                            "127.0.0.1:1980": 1
+                        },
+                        "type": "chash",
+                        "key": "remote_addr"
+                    },
+                    "uri": "/log_request"
+                }]]
+                )
+
+            if code >= 300 then
+                ngx.status = code
+            end
+            ngx.say(body)
+        }
+    }
+--- request
+GET /t
+--- response_body
+passed
+--- no_error_log
+[error]
+
+
+
+=== TEST 24: check plugin
+--- request
+GET /log_request
+--- skip_nginx: 4: < 1.19.3
+--- grep_error_log eval
+qr/(proxy request to \S+|x-serverless: [\d.]+)/
+--- grep_error_log_out
+proxy request to 127.0.0.2:1979
+proxy request to 127.0.0.1:1980
+x-serverless: 127.0.0.1
+
+
+
+=== TEST 25: exit in the balancer phase
+--- 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": {
+                        "serverless-pre-function": {
+                            "phase": "balancer",
+                            "functions" : ["return function(conf, ctx) ngx.exit(403) end"]
+                        }
+                    },
+                    "upstream": {
+                        "nodes": {
+                            "127.0.0.2:1979": 100000,
+                            "127.0.0.1:1980": 1
+                        },
+                        "type": "chash",
+                        "key": "remote_addr"
+                    },
+                    "uri": "/log_request"
+                }]]
+                )
+
+            if code >= 300 then
+                ngx.status = code
+            end
+            ngx.say(body)
+        }
+    }
+--- request
+GET /t
+--- response_body
+passed
+--- no_error_log
+[error]
+
+
+
+=== TEST 26: check plugin
+--- request
+GET /log_request
+--- error_code: 403
+--- no_error_log
+[error]
+
+
+
+=== TEST 27: ensure balancer phase run correct time
+--- 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": {
+                        "serverless-pre-function": {
+                            "phase": "balancer",
+                            "functions" : ["return function(conf, ctx) ngx.log(ngx.WARN, 'run balancer phase with ', ctx.balancer_ip) end"]
+                        }
+                    },
+                    "upstream": {
+                        "nodes": {
+                            "127.0.0.2:1979": 100000,
+                            "127.0.0.1:1980": 1
+                        },
+                        "type": "chash",
+                        "key": "remote_addr"
+                    },
+                    "uri": "/log_request"
+                }]]
+                )
+
+            if code >= 300 then
+                ngx.status = code
+            end
+            ngx.say(body)
+        }
+    }
+--- request
+GET /t
+--- response_body
+passed
+--- no_error_log
+[error]
+
+
+
+=== TEST 28: check plugin
+--- request
+GET /log_request
+--- grep_error_log eval
+qr/(run balancer phase with [\d.]+)/
+--- grep_error_log_out
+run balancer phase with 127.0.0.2
+run balancer phase with 127.0.0.1