feat: control the proxy_upstream in lua side

lualib/resty/kong/upstream.lua

@@ -0,0 +1,82 @@
+-- Copyright 2019-2022 Kong Inc.
+
+-- Licensed under the Apache License, Version 2.0 (the "License");
+-- you may not use this file except in compliance with the License.
+-- You may obtain a copy of the License at
+
+--    http://www.apache.org/licenses/LICENSE-2.0
+
+-- Unless required by applicable law or agreed to in writing, software
+-- distributed under the License is distributed on an "AS IS" BASIS,
+-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+-- See the License for the specific language governing permissions and
+-- limitations under the License.
+
+local _M = {}
+
+local ffi = require("ffi")
+local base = require("resty.core.base")
+base.allows_subsystem("http")
+
+ffi.cdef([[
+int
+ngx_http_lua_ffi_set_upstream_next(ngx_http_request_t *r, uint next_upstream, char **err);
+]])
+
+local type = type
+local error = error
+local tostring = tostring
+local C = ffi.C
+local get_request = base.get_request
+local ffi_str = ffi.string
+
+local NGX_OK = ngx.OK
+
+local next_upstream_table = {
+	error = 0x00000002,
+	timeout = 0x00000004,
+	invalid_header = 0x00000008,
+	http_500 = 0x00000010,
+	http_502 = 0x00000020,
+	http_503 = 0x00000040,
+	http_504 = 0x00000080,
+	http_403 = 0x00000100,
+	http_404 = 0x00000200,
+	http_429 = 0x00000400,
+	updating = 0x00000800,
+	off = 0x00001000,
+}
+
+function _M.set_upstream_next(...)
+	local nargs = select("#", ...)
+	if nargs == 0 then
+		error("no argument")
+	end
+
+	local r = get_request()
+	if not r then
+		error("no request found")
+	end
+
+	local next_upstream_table = { ... }
+	local next_upstream = 0
+	for i = 1, nargs do
+		local v = next_upstream_table[i]
+		if type(v) ~= "number" then
+			error("argument #" .. i .. " is not a valid argument")
+		end
+
+		next_upstream = bit.bor(next_upstream, v)
+	end
+
+	local err = ffi.new("char *[1]")
+	local rc = C.ngx_http_lua_ffi_set_upstream_next(r, next_upstream, err)
+
+	if rc ~= NGX_OK then
+		error("failed to set upstream next: " .. tostring(ffi_str(err[0])))
+	end
+
+	return true
+end
+
+return _M

src/ngx_http_lua_kong_common.h

@@ -28,6 +28,7 @@ typedef struct {
     ngx_lua_kong_ssl_ctx_t   ssl_ctx;
     ngx_str_t                grpc_authority;
     ngx_http_log_handler_pt  orig_log_handler;
+    ngx_uint_t               next_upstream;
 } ngx_http_lua_kong_ctx_t;
 
 

src/ngx_http_lua_kong_module.c

@@ -158,4 +158,75 @@ ngx_http_lua_kong_merge_loc_conf(ngx_conf_t *cf, void *parent, void *child)
     return NGX_CONF_OK;
 }
 
+#define NGX_HTTP_UPSTREAM_FT_ERROR           0x00000002
+#define NGX_HTTP_UPSTREAM_FT_TIMEOUT         0x00000004
+#define NGX_HTTP_UPSTREAM_FT_INVALID_HEADER  0x00000008
+#define NGX_HTTP_UPSTREAM_FT_HTTP_500        0x00000010
+#define NGX_HTTP_UPSTREAM_FT_HTTP_502        0x00000020
+#define NGX_HTTP_UPSTREAM_FT_HTTP_503        0x00000040
+#define NGX_HTTP_UPSTREAM_FT_HTTP_504        0x00000080
+#define NGX_HTTP_UPSTREAM_FT_HTTP_403        0x00000100
+#define NGX_HTTP_UPSTREAM_FT_HTTP_404        0x00000200
+#define NGX_HTTP_UPSTREAM_FT_HTTP_429        0x00000400
+#define NGX_HTTP_UPSTREAM_FT_UPDATING        0x00000800
+#define NGX_HTTP_UPSTREAM_FT_BUSY_LOCK       0x00001000
+#define NGX_HTTP_UPSTREAM_FT_MAX_WAITING     0x00002000
+#define NGX_HTTP_UPSTREAM_FT_NON_IDEMPOTENT  0x00004000
+#define NGX_HTTP_UPSTREAM_FT_NOLIVE          0x40000000
+#define NGX_HTTP_UPSTREAM_FT_OFF             0x80000000
+
+#define NGX_HTTP_UPSTREAM_FT_STATUS          (NGX_HTTP_UPSTREAM_FT_HTTP_500  \
+                                             |NGX_HTTP_UPSTREAM_FT_HTTP_502  \
+                                             |NGX_HTTP_UPSTREAM_FT_HTTP_503  \
+                                             |NGX_HTTP_UPSTREAM_FT_HTTP_504  \
+                                             |NGX_HTTP_UPSTREAM_FT_HTTP_403  \
+                                             |NGX_HTTP_UPSTREAM_FT_HTTP_404  \
+                                             |NGX_HTTP_UPSTREAM_FT_HTTP_429)
+
+static ngx_conf_bitmask_t  ngx_http_proxy_next_upstream_masks[] = {
+    { ngx_string("error"), NGX_HTTP_UPSTREAM_FT_ERROR },
+    { ngx_string("timeout"), NGX_HTTP_UPSTREAM_FT_TIMEOUT },
+    { ngx_string("invalid_header"), NGX_HTTP_UPSTREAM_FT_INVALID_HEADER },
+    { ngx_string("non_idempotent"), NGX_HTTP_UPSTREAM_FT_NON_IDEMPOTENT },
+    { ngx_string("http_500"), NGX_HTTP_UPSTREAM_FT_HTTP_500 },
+    { ngx_string("http_502"), NGX_HTTP_UPSTREAM_FT_HTTP_502 },
+    { ngx_string("http_503"), NGX_HTTP_UPSTREAM_FT_HTTP_503 },
+    { ngx_string("http_504"), NGX_HTTP_UPSTREAM_FT_HTTP_504 },
+    { ngx_string("http_403"), NGX_HTTP_UPSTREAM_FT_HTTP_403 },
+    { ngx_string("http_404"), NGX_HTTP_UPSTREAM_FT_HTTP_404 },
+    { ngx_string("http_429"), NGX_HTTP_UPSTREAM_FT_HTTP_429 },
+    { ngx_string("updating"), NGX_HTTP_UPSTREAM_FT_UPDATING },
+    { ngx_string("off"), NGX_HTTP_UPSTREAM_FT_OFF },
+    { ngx_null_string, 0 }
+};
+
+ngx_flag_t
+ngx_http_lua_kong_get_next_upstream_mask(ngx_http_request_t *r,
+    ngx_flag_t upstream_next)
+{
+    ngx_http_lua_kong_ctx_t      *ctx;
+
+    ctx = ngx_http_lua_kong_get_module_ctx(r);
+    if (ctx == NULL) {
+        return upstream_next;
+    }
+
+    if ctx->next_upstream ~= 0 {
+        return ctx->next_upstream;
+    }
+    return upstream_next;
+}
+
+int
+ngx_http_lua_ffi_set_upstream_next(ngx_http_request_t *r, ngx_uint_t next_upstream, char **err)
+{
+    ngx_http_lua_kong_ctx_t      *ctx;
+
+    ctx = ngx_http_lua_kong_get_module_ctx(r);
+    if (ctx == NULL) {
+        return NGX_ERROR;
+    }
 
+    ctx->next_upstream = next_upstream;
+    return NGX_OK;
+}