feature: implemented the ngx.balancer Lua module to support dynamic nginx stream upstream balancers written in Lua.

.travis.yml

@@ -42,6 +42,7 @@ install:
   - git clone https://github.com/openresty/nginx-devel-utils.git
   - git clone https://github.com/simpl/ngx_devel_kit.git ../ndk-nginx-module
   - git clone https://github.com/openresty/lua-nginx-module.git ../lua-nginx-module
+  - git clone -b balancer https://github.com/doujiang24/stream-lua-nginx-module.git ../stream-lua-nginx-module
   - git clone https://github.com/openresty/no-pool-nginx.git ../no-pool-nginx
   - git clone https://github.com/openresty/echo-nginx-module.git ../echo-nginx-module
   - git clone https://github.com/openresty/lua-resty-lrucache.git
@@ -68,7 +69,7 @@ script:
   - export LD_LIBRARY_PATH=$PWD/mockeagain:$LD_LIBRARY_PATH
   - export TEST_NGINX_RESOLVER=8.8.4.4
   - export NGX_BUILD_CC=$CC
-  - ngx-build $NGINX_VERSION --with-ipv6 --with-http_realip_module --with-http_ssl_module --with-cc-opt="-I$OPENSSL_INC" --with-ld-opt="-L$OPENSSL_LIB -Wl,-rpath,$OPENSSL_LIB" --add-module=../ndk-nginx-module --add-module=../echo-nginx-module --add-module=../headers-more-nginx-module --add-module=../lua-nginx-module --with-debug > build.log 2>&1 || (cat build.log && exit 1)
+  - ngx-build $NGINX_VERSION --with-ipv6 --with-http_realip_module --with-http_ssl_module --with-cc-opt="-I$OPENSSL_INC" --with-ld-opt="-L$OPENSSL_LIB -Wl,-rpath,$OPENSSL_LIB" --add-module=../ndk-nginx-module --add-module=../echo-nginx-module --add-module=../headers-more-nginx-module --add-module=../lua-nginx-module --add-module=../stream-lua-nginx-module --with-stream --with-stream_ssl_module --with-debug > build.log 2>&1 || (cat build.log && exit 1)
   - nginx -V
   - ldd `which nginx`|grep -E 'luajit|ssl|pcre'
   - prove -r t

Makefile

@@ -14,10 +14,12 @@ install: all
 	$(INSTALL) -d $(DESTDIR)$(LUA_LIB_DIR)/resty/core/
 	$(INSTALL) -d $(DESTDIR)$(LUA_LIB_DIR)/ngx/
 	$(INSTALL) -d $(DESTDIR)$(LUA_LIB_DIR)/ngx/ssl
+	$(INSTALL) -d $(DESTDIR)$(LUA_LIB_DIR)/ngx/balancer
 	$(INSTALL) lib/resty/*.lua $(DESTDIR)$(LUA_LIB_DIR)/resty/
 	$(INSTALL) lib/resty/core/*.lua $(DESTDIR)$(LUA_LIB_DIR)/resty/core/
 	$(INSTALL) lib/ngx/*.lua $(DESTDIR)$(LUA_LIB_DIR)/ngx/
 	$(INSTALL) lib/ngx/ssl/*.lua $(DESTDIR)$(LUA_LIB_DIR)/ngx/ssl/
+	$(INSTALL) lib/ngx/balancer/*.lua $(DESTDIR)$(LUA_LIB_DIR)/ngx/balancer/
 
 test: all
 	PATH=$(OPENRESTY_PREFIX)/nginx/sbin:$$PATH prove -I../test-nginx/lib -r t

lib/ngx/balancer.lua

@@ -1,6 +1,11 @@
 -- Copyright (C) Yichun Zhang (agentzh)
 
 
+if ngx.config.subsystem == "stream" then
+    return require "ngx.balancer.stream"
+end
+
+
 local ffi = require "ffi"
 local base = require "resty.core.base"
 

lib/ngx/balancer/stream.lua

@@ -0,0 +1,79 @@
+-- Copyright (C) Yichun Zhang (agentzh)
+
+
+local ffi = require "ffi"
+local base = require "resty.core.base"
+
+
+local C = ffi.C
+local ffi_str = ffi.string
+local errmsg = base.get_errmsg_ptr()
+local FFI_OK = base.FFI_OK
+local FFI_ERROR = base.FFI_ERROR
+local int_out = ffi.new("int[1]")
+local getfenv = getfenv
+local error = error
+local type = type
+local tonumber = tonumber
+
+
+ffi.cdef[[
+int ngx_stream_lua_ffi_balancer_set_current_peer(ngx_stream_session_t *s,
+    const unsigned char *addr, size_t addr_len, int port, char **err);
+
+int ngx_stream_lua_ffi_balancer_set_more_tries(ngx_stream_session_t *s,
+    int count, char **err);
+
+int ngx_stream_lua_ffi_balancer_get_last_failure(ngx_stream_session_t *s,
+    int *status, char **err);
+
+int ngx_stream_lua_ffi_balancer_set_timeouts(ngx_stream_session_t *s,
+    long connect_timeout, long send_timeout,
+    long read_timeout, char **err);
+]]
+
+
+local _M = { version = base.version }
+
+
+function _M.set_current_peer(addr, port)
+    local s = getfenv(0).__ngx_sess
+    if not s then
+        return error("no request found")
+    end
+
+    if not port then
+        port = 0
+    elseif type(port) ~= "number" then
+        port = tonumber(port)
+    end
+
+    local rc = C.ngx_stream_lua_ffi_balancer_set_current_peer(s, addr, #addr,
+                                                              port, errmsg)
+    if rc == FFI_OK then
+        return true
+    end
+
+    return nil, ffi_str(errmsg[0])
+end
+
+
+function _M.set_more_tries(count)
+    local s = getfenv(0).__ngx_sess
+    if not s then
+        return error("no request found")
+    end
+
+    local rc = C.ngx_stream_lua_ffi_balancer_set_more_tries(s, count, errmsg)
+    if rc == FFI_OK then
+        if errmsg[0] == nil then
+            return true
+        end
+        return true, ffi_str(errmsg[0])  -- return the warning
+    end
+
+    return nil, ffi_str(errmsg[0])
+end
+
+
+return _M

lib/resty/core/base.lua

@@ -79,6 +79,14 @@ if not pcall(ffi.typeof, "ngx_http_request_t") then
 end
 
 
+if not pcall(ffi.typeof, "ngx_stream_session_t") then
+    ffi.cdef[[
+        struct ngx_stream_session_s;
+        typedef struct ngx_stream_session_s  ngx_stream_session_t;
+    ]]
+end
+
+
 if not pcall(ffi.typeof, "ngx_http_lua_ffi_str_t") then
     ffi.cdef[[
         typedef struct {

t/balancer-stream.t

@@ -0,0 +1,359 @@
+# vim:set ft= ts=4 sw=4 et fdm=marker:
+
+use Test::Nginx::Socket::Lua::Stream 'no_plan';
+use Cwd qw(cwd);
+
+#worker_connections(1014);
+#master_on();
+#workers(2);
+#log_level('warn');
+
+repeat_each(1);
+
+#plan tests => repeat_each() * (blocks() * 4 + 5);
+
+$ENV{TEST_NGINX_CWD} = cwd();
+
+#worker_connections(1024);
+#no_diff();
+no_long_string();
+run_tests();
+
+__DATA__
+
+=== TEST 1: set current peer (separate addr and port)
+--- stream_config
+    lua_package_path "$TEST_NGINX_CWD/lib/?.lua;;";
+
+    upstream backend {
+        server 0.0.0.1:80;
+        balancer_by_lua_block {
+            print("hello from balancer by lua!")
+            local b = require "ngx.balancer"
+            assert(b.set_current_peer("127.0.0.3", 12345))
+        }
+    }
+--- stream_server_config
+        proxy_pass backend;
+--- error_log eval
+[
+'[lua] balancer_by_lua:2: hello from balancer by lua! while connecting to upstream,',
+qr{connect\(\) failed .*?, upstream: "127\.0\.0\.3:12345"},
+]
+--- no_error_log
+[warn]
+
+
+
+=== TEST 2: set current peer & next upstream (3 tries)
+--- stream_config
+    lua_package_path "$TEST_NGINX_CWD/lib/?.lua;;";
+
+    proxy_next_upstream_tries 10;
+
+    upstream backend {
+        server 0.0.0.1:80;
+        balancer_by_lua_block {
+            print("hello from balancer by lua!")
+            local b = require "ngx.balancer"
+            if not ngx.ctx.tries then
+                ngx.ctx.tries = 0
+            end
+
+            if ngx.ctx.tries < 2 then
+                local ok, err = b.set_more_tries(1)
+                if not ok then
+                    return error("failed to set more tries: ", err)
+                elseif err then
+                    ngx.log(ngx.WARN, "set more tries: ", err)
+                end
+            end
+            ngx.ctx.tries = ngx.ctx.tries + 1
+            assert(b.set_current_peer("127.0.0.3", 12345))
+        }
+    }
+--- stream_server_config
+        proxy_pass backend;
+--- grep_error_log eval: qr{connect\(\) failed .*, upstream: ".*?"}
+--- grep_error_log_out eval
+qr#^(?:connect\(\) failed .*?, upstream: "127.0.0.3:12345"\n){3}$#
+--- no_error_log
+[warn]
+
+
+
+=== TEST 3: set current peer & next upstream (no retries)
+--- stream_config
+    lua_package_path "$TEST_NGINX_CWD/lib/?.lua;;";
+
+    upstream backend {
+        server 0.0.0.1:80;
+        balancer_by_lua_block {
+            print("hello from balancer by lua!")
+            local b = require "ngx.balancer"
+            if not ngx.ctx.tries then
+                ngx.ctx.tries = 0
+            end
+
+            ngx.ctx.tries = ngx.ctx.tries + 1
+            assert(b.set_current_peer("127.0.0.3", 12345))
+        }
+    }
+--- stream_server_config
+        proxy_pass backend;
+--- grep_error_log eval: qr{connect\(\) failed .*, upstream: ".*?"}
+--- grep_error_log_out eval
+qr#^(?:connect\(\) failed .*?, upstream: "127.0.0.3:12345"\n){1}$#
+--- no_error_log
+[warn]
+
+
+
+=== TEST 4: set current peer & next upstream (3 tries exceeding the limit)
+--- SKIP
+--- TODO
+--- stream_config
+    lua_package_path "$TEST_NGINX_CWD/lib/?.lua;;";
+
+    proxy_next_upstream_tries 2;
+
+    upstream backend {
+        server 0.0.0.1:80;
+        balancer_by_lua_block {
+            local b = require "ngx.balancer"
+
+            if not ngx.ctx.tries then
+                ngx.ctx.tries = 0
+            end
+
+            if ngx.ctx.tries < 2 then
+                local ok, err = b.set_more_tries(1)
+                if not ok then
+                    return error("failed to set more tries: ", err)
+                elseif err then
+                    ngx.log(ngx.WARN, "set more tries: ", err)
+                end
+            end
+            ngx.ctx.tries = ngx.ctx.tries + 1
+            assert(b.set_current_peer("127.0.0.3", 12345))
+        }
+    }
+--- stream_server_config
+        proxy_pass backend;
+--- grep_error_log eval: qr{connect\(\) failed .*, upstream: ".*?"}
+--- grep_error_log_out eval
+qr#^(?:connect\(\) failed .*?, upstream: "127.0.0.3:12345"\n){2}$#
+--- error_log
+set more tries: reduced tries due to limit
+
+
+
+=== TEST 5: set current peer (port embedded in addr)
+--- stream_config
+    lua_package_path "$TEST_NGINX_CWD/lib/?.lua;;";
+
+    upstream backend {
+        server 0.0.0.1:80;
+        balancer_by_lua_block {
+            print("hello from balancer by lua!")
+            local b = require "ngx.balancer"
+            assert(b.set_current_peer("127.0.0.3:12345"))
+        }
+    }
+--- stream_server_config
+        proxy_pass backend;
+--- error_log eval
+[
+'[lua] balancer_by_lua:2: hello from balancer by lua! while connecting to upstream,',
+qr{connect\(\) failed .*?, upstream: "127\.0\.0\.3:12345"},
+]
+--- no_error_log
+[warn]
+
+
+
+=== TEST 6: set_current_peer called in a wrong context
+--- wait: 0.2
+--- stream_config
+    lua_package_path "$TEST_NGINX_CWD/lib/?.lua;;";
+
+    upstream backend {
+        server 127.0.0.1:$TEST_NGINX_SERVER_PORT;
+        balancer_by_lua_block {
+            print("hello from balancer by lua!")
+        }
+    }
+
+--- stream_server_config
+        proxy_pass backend;
+
+        content_by_lua_block {
+            local balancer = require "ngx.balancer"
+            local ok, err = balancer.set_current_peer("127.0.0.1", 1234)
+            if not ok then
+                ngx.log(ngx.ERR, "failed to call: ", err)
+                return
+            end
+            ngx.log(ngx.ALERT, "unexpected success")
+        }
+--- error_log eval
+qr/\[error\] .*? content_by_lua.*? failed to call: API disabled in the current context/
+--- no_error_log
+[alert]
+
+
+
+=== TEST 7: set_more_tries called in a wrong context
+--- wait: 0.2
+--- stream_config
+    lua_package_path "$TEST_NGINX_CWD/lib/?.lua;;";
+
+    upstream backend {
+        server 127.0.0.1:$TEST_NGINX_SERVER_PORT;
+        balancer_by_lua_block {
+            print("hello from balancer by lua!")
+        }
+    }
+
+--- stream_server_config
+        content_by_lua_block {
+            local balancer = require "ngx.balancer"
+            local ok, err = balancer.set_more_tries(1)
+            if not ok then
+                ngx.log(ngx.ERR, "failed to call: ", err)
+                return
+            end
+            ngx.log(ngx.ALERT, "unexpected success")
+        }
+--- error_log eval
+qr/\[error\] .*? content_by_lua.*? failed to call: API disabled in the current context/
+--- no_error_log
+[alert]
+
+
+
+=== TEST 8: https => http
+--- stream_config
+    lua_package_path "$TEST_NGINX_CWD/lib/?.lua;;";
+
+    upstream backend {
+        server 0.0.0.1:80;
+        balancer_by_lua_block {
+            local b = require "ngx.balancer"
+            print("hello from balancer by lua!")
+            assert(b.set_current_peer("127.0.0.1", 1234))
+        }
+    }
+
+    server {
+        listen 1234;
+
+        content_by_lua_block {
+            ngx.say("ok")
+        }
+    }
+
+    server {
+        listen 1235 ssl;
+        ssl_certificate ../../cert/test.crt;
+        ssl_certificate_key ../../cert/test.key;
+
+        proxy_pass backend;
+    }
+--- stream_server_config
+        content_by_lua_block {
+            local sock, err = ngx.socket.tcp()
+            assert(sock, err)
+
+            local ok, err = sock:connect("127.0.0.1", 1235)
+            if not ok then
+                ngx.say("connect to stream server error: ", err)
+                return
+            end
+
+            local sess, err = sock:sslhandshake(nil, "test.com", false)
+            if not sess then
+                ngx.say("failed to do SSL handshake: ", err)
+                return
+            end
+
+            local data, err = sock:receive("*a")
+            if not data then
+                sock:close()
+                ngx.say("receive stream response error: ", err)
+                return
+            end
+            ngx.print(data)
+        }
+--- no_error_log
+[alert]
+[error]
+--- stream_response
+ok
+--- grep_error_log eval: qr{hello from balancer by lua!}
+--- grep_error_log_out
+hello from balancer by lua!
+
+
+
+=== TEST 9: http => https
+--- stream_config
+    lua_package_path "$TEST_NGINX_CWD/lib/?.lua;;";
+
+    upstream backend {
+        server 0.0.0.1:80;
+        balancer_by_lua_block {
+            local b = require "ngx.balancer"
+            print("hello from balancer by lua!")
+            assert(b.set_current_peer("127.0.0.1", 1234))
+        }
+    }
+
+    server {
+        listen 1234 ssl;
+        ssl_certificate ../../cert/test.crt;
+        ssl_certificate_key ../../cert/test.key;
+
+        content_by_lua_block {
+            ngx.say("ok")
+        }
+    }
+
+    server {
+        listen 1235;
+
+        proxy_pass backend;
+    }
+--- stream_server_config
+        content_by_lua_block {
+            local sock, err = ngx.socket.tcp()
+            assert(sock, err)
+
+            local ok, err = sock:connect("127.0.0.1", 1235)
+            if not ok then
+                ngx.say("connect to stream server error: ", err)
+                return
+            end
+
+            local sess, err = sock:sslhandshake(nil, "test.com", false)
+            if not sess then
+                ngx.say("failed to do SSL handshake: ", err)
+                return
+            end
+
+            local data, err = sock:receive("*a")
+            if not data then
+                sock:close()
+                ngx.say("receive stream response error: ", err)
+                return
+            end
+            ngx.print(data)
+        }
+--- no_error_log
+[alert]
+[error]
+--- stream_response
+ok
+--- grep_error_log eval: qr{hello from balancer by lua!}
+--- grep_error_log_out
+hello from balancer by lua!