feature: add the balancer.recreate_request function, which allows user to recreate request buffer in balancer phase. (#302)

This allows certain request parameters, such as headers (including
`Host` header) to be modified between balancer retries.

Author: dndx

lib/ngx/balancer.lua

@@ -38,6 +38,9 @@ if subsystem == 'http' then
     int ngx_http_lua_ffi_balancer_set_timeouts(ngx_http_request_t *r,
         long connect_timeout, long send_timeout,
         long read_timeout, char **err);
+
+    int ngx_http_lua_ffi_balancer_recreate_request(ngx_http_request_t *r,
+        char **err);
     ]]
 
     ngx_lua_ffi_balancer_set_current_peer =
@@ -207,4 +210,25 @@ function _M.set_timeouts(connect_timeout, send_timeout, read_timeout)
 end
 
 
+if subsystem == 'http' then
+    function _M.recreate_request()
+        local r = get_request()
+        if not r then
+            error("no request found")
+        end
+
+        local rc = C.ngx_http_lua_ffi_balancer_recreate_request(r, errmsg)
+        if rc == FFI_OK then
+            return true
+        end
+
+        if errmsg[0] ~= nil then
+            return nil, ffi_str(errmsg[0])
+        end
+
+        return nil, "failed to recreate the upstream request"
+    end
+end
+
+
 return _M

lib/ngx/balancer.md

@@ -9,12 +9,15 @@ Table of Contents
 * [Name](#name)
 * [Status](#status)
 * [Synopsis](#synopsis)
+    * [http subsystem](#http-subsystem)
+    * [stream subsystem](#stream-subsystem)
 * [Description](#description)
 * [Methods](#methods)
     * [set_current_peer](#set_current_peer)
     * [set_more_tries](#set_more_tries)
     * [get_last_failure](#get_last_failure)
     * [set_timeouts](#set_timeouts)
+    * [recreate_request](#recreate_request)
 * [Community](#community)
     * [English Mailing List](#english-mailing-list)
     * [Chinese Mailing List](#chinese-mailing-list)
@@ -79,6 +82,8 @@ http {
 }
 ```
 
+[Back to TOC](#table-of-contents)
+
 stream subsystem
 ----------------
 
@@ -121,6 +126,8 @@ stream {
 }
 ```
 
+[Back to TOC](#table-of-contents)
+
 Description
 ===========
 
@@ -237,6 +244,32 @@ This function was first added in the `0.1.7` version of this library.
 
 [Back to TOC](#table-of-contents)
 
+recreate_request
+----------------
+**syntax:** `ok, err = balancer.recreate_request()`
+
+**context:** *balancer_by_lua**
+
+Recreates the request buffer for sending to the upstream server. This is useful, for example
+if you want to change a request header field to the new upstream server on balancer retries.
+
+Normally this does not work because the request buffer is created once during upstream module
+initialization and won't be regenerated for subsequent retries. However you can use
+`proxy_set_header My-Header $my_header` and set the `ngx.var.my_header` variable inside the
+balancer phase. Calling `balancer.recreate_request()` after updating a header field will
+cause the request buffer to be re-generated and the `My-Header` header will thus contain
+the new value.
+
+**Warning:** because the request buffer has to be recreated and such allocation occurs on the
+request memory pool, the old buffer has to be thrown away and will only be freed after the request
+finishes. Do not call this function too often or memory leaks may be noticeable. Even so, a call
+to this function should be made **only** if you know the request buffer must be regenerated,
+instead of unconditionally in each balancer retries.
+
+This function was first added in the `0.1.20` version of this library.
+
+[Back to TOC](#table-of-contents)
+
 Community
 =========
 

t/balancer.t

@@ -9,7 +9,7 @@ use t::TestCore;
 
 repeat_each(2);
 
-plan tests => repeat_each() * (blocks() * 4 + 5);
+plan tests => repeat_each() * (blocks() * 4 + 6);
 
 $ENV{TEST_NGINX_LUA_PACKAGE_PATH} = "$t::TestCore::lua_package_path";
 
@@ -834,3 +834,51 @@ GET /t
 [lua] log_by_lua(nginx.conf:59):2: ngx.var.upstream_addr is 127.0.0.3:12345, 127.0.0.3:12346
 --- no_error_log
 [alert]
+
+
+
+=== TEST 19: recreate upstream module requests with header change
+--- http_config
+    lua_package_path "$TEST_NGINX_LUA_PACKAGE_PATH";
+
+    upstream backend {
+        server 0.0.0.1;
+
+        balancer_by_lua_block {
+            print("here")
+            local b = require "ngx.balancer"
+
+            if ngx.ctx.balancer_run then
+                assert(b.set_current_peer("127.0.0.1", tonumber(ngx.var.server_port)))
+                ngx.var.test = "second"
+                assert(b.recreate_request())
+
+            else
+                ngx.ctx.balancer_run = true
+                assert(b.set_current_peer("127.0.0.3", 12345))
+                assert(b.set_more_tries(1))
+            end
+        }
+    }
+--- config
+    location = /t {
+        proxy_next_upstream error timeout invalid_header http_500 http_502 http_503 http_504 http_403 http_404;
+        proxy_next_upstream_tries 2;
+
+        set $test "first";
+
+        proxy_set_header X-Test $test;
+        proxy_pass http://backend/upstream;
+    }
+
+    location = /upstream {
+        return 200 "value is: $http_x_test";
+    }
+--- request
+GET /t
+--- response_body: value is: second
+--- error_log
+connect() failed (111: Connection refused) while connecting to upstream, client: 127.0.0.1
+--- no_error_log
+[warn]
+[crit]