Socket pool

src/ngx_http_lua_socket_tcp.c

@@ -260,12 +260,13 @@ ngx_http_lua_socket_tcp_connect(lua_State *L)
     ngx_http_lua_loc_conf_t     *llcf;
     ngx_peer_connection_t       *pc;
     int                          timeout;
+    const char *                 pool;
 
     ngx_http_lua_socket_tcp_upstream_t      *u;
 
     n = lua_gettop(L);
-    if (n != 2 && n != 3) {
-        return luaL_error(L, "ngx.socket connect: expecting 2 or 3 arguments "
+    if (n != 2 && n != 3 && n != 4) {
+        return luaL_error(L, "ngx.socket connect: expecting 2, 3, or 4 arguments "
                           "(including the object), but seen %d", n);
     }
 
@@ -301,25 +302,69 @@ ngx_http_lua_socket_tcp_connect(lua_State *L)
     ngx_memcpy(host.data, p, len);
     host.data[len] = '\0';
 
-    if (n == 3) {
-        port = luaL_checkinteger(L, 3);
+    if (n >= 3) {
+        switch (lua_type(L, 3)) {
+           case LUA_TNIL:
+               /* no port passed in */
+               port = 0;
+               break;
+
+            case LUA_TNUMBER:
+            case LUA_TSTRING:
+                port = lua_tonumber(L, 3);
+                if (port < 0 || port > 65536) {
+                    lua_pushnil(L);
+                    lua_pushfstring(L, "bad port number: %d", port);
+                    return 2;
+                }
+                break;
 
-        if (port < 0 || port > 65536) {
-            lua_pushnil(L);
-            lua_pushfstring(L, "bad port number: %d", port);
-            return 2;
+            default:
+               lua_pushnil(L);
+               lua_pushfstring(L, "bad port option. Expected number, got %s", luaL_typename(L, 3));
+               return 2;
         }
 
+    } else { /* n == 2 */
+        port = 0;
+    }
+
+    pool = NULL;
+    if(n == 4) {
+        luaL_checktype(L, 4, LUA_TTABLE);
+
+        lua_getfield(L, 4, "pool");
+
+        switch (lua_type(L, -1)) {
+
+            case LUA_TNIL:
+                /* remove table and nil from stack */
+                lua_pop(L, 2);
+                break;
+
+            case LUA_TSTRING:
+                /* clean up the stack */
+                lua_replace(L, 2);
+                lua_pop(L, n - 2);
+                pool = lua_tostring(L, -1);
+                break;
+
+            default:
+                return luaL_error(L, "bad \"pool\" option value type: %s",
+                                  luaL_typename(L, -1));
+
+        }
+    }
+
+    if(!pool && n == 3) {
         lua_pushliteral(L, ":");
         lua_insert(L, 3);
         lua_concat(L, 3);
-
-        dd("socket key: %s", lua_tostring(L, -1));
-
-    } else { /* n == 2 */
-        port = 0;
     }
 
+    /* key is on top of stack. it may be host, host:port, or a pool name*/
+    dd("socket key: %s", lua_tostring(L, -1));
+
     /* the key's index is 2 */
 
     lua_pushvalue(L, -1);

t/086-socket-keepalive-pool.t

@@ -0,0 +1,155 @@
+# vim:set ft= ts=4 sw=4 et fdm=marker:
+
+use lib 'lib';
+use Test::Nginx::Socket;
+
+repeat_each(2);
+
+plan tests => blocks() * 2 * repeat_each();
+
+our $HtmlDir = html_dir;
+
+$ENV{TEST_NGINX_CLIENT_PORT} ||= server_port();
+$ENV{TEST_NGINX_MEMCACHED_PORT} ||= 11211;
+$ENV{TEST_NGINX_HTML_DIR} = $HtmlDir;
+#$ENV{TEST_NGINX_REDIS_PORT} ||= 6379;
+
+$ENV{LUA_PATH} ||=
+    '/usr/local/openresty-debug/lualib/?.lua;/usr/local/openresty/lualib/?.lua;;';
+
+no_long_string();
+#no_diff();
+#log_level 'warn';
+
+no_shuffle();
+
+run_tests();
+
+__DATA__
+
+=== TEST 1: sanity
+--- config
+    location /t {
+        set $port $TEST_NGINX_MEMCACHED_PORT;
+        content_by_lua '
+            local port = ngx.var.port
+            for i=1,2 do
+                local sock = ngx.socket.tcp()
+                local ok, err = sock:connect("127.0.0.1", port, { pool = "test" })
+                if not ok then
+                    ngx.say("failed to connect: ", err)
+                    return
+                end
+                ngx.say("connected: ", ok, ", reused: ", sock:getreusedtimes())
+                local ok, err = sock:setkeepalive()
+                if not ok then
+                    ngx.say("failed to set reusable: ", err)
+                end
+            end
+        ';
+    }
+--- request
+GET /t
+--- response_body_like
+^connected: 1, reused: \d+
+connected: 1, reused: [1-9]\d*
+
+=== TEST 2: two pools
+--- config
+    location /t {
+        set $port $TEST_NGINX_MEMCACHED_PORT;
+        content_by_lua '
+            local port = ngx.var.port
+            function socktest(port, pool)
+                local sock = ngx.socket.tcp()
+                local ok, err = sock:connect("127.0.0.1", port, { pool = pool })
+                if not ok then
+                    ngx.say("failed to connect: ", err)
+                    return
+                end
+                ngx.say("connected: ", ok, ", reused: ", sock:getreusedtimes())
+                local ok, err = sock:setkeepalive()
+                if not ok then
+                    ngx.say("failed to set reusable: ", err)
+                end
+            end
+            for i=1,10 do
+                socktest(port, "one")
+                socktest(port, "two")
+            end
+        ';
+    }
+--- request
+GET /t
+--- response_body_like
+(connected: 1, reused: \d+\s+)*
+
+=== TEST 3: unix socket
+--- http_config
+    server {
+        listen unix:$TEST_NGINX_HTML_DIR/nginx.sock;
+        default_type 'text/plain';
+    }
+--- config
+    location /t {
+        set $port $TEST_NGINX_MEMCACHED_PORT;
+        content_by_lua '
+            local port = ngx.var.port
+            for i=1,2 do
+                local sock = ngx.socket.tcp()
+                local ok, err = sock:connect("unix:$TEST_NGINX_HTML_DIR/nginx.sock", nil, { pool = "nginx.sock" })
+                if not ok then
+                    ngx.say("failed to connect: ", err)
+                    return
+                end
+                ngx.say("connected: ", ok, ", reused: ", sock:getreusedtimes())
+                local ok, err = sock:setkeepalive()
+                if not ok then
+                    ngx.say("failed to set reusable: ", err)
+                end
+            end
+        ';
+    }
+--- request
+GET /t
+--- response_body_like
+^connected: 1, reused: \d+
+connected: 1, reused: [1-9]\d*
+
+
+=== TEST 2: unix socket - two pools
+--- http_config
+    server {
+        listen unix:$TEST_NGINX_HTML_DIR/nginx.sock;
+        default_type 'text/plain';
+    }
+--- config
+    location /t {
+        set $port $TEST_NGINX_MEMCACHED_PORT;
+        content_by_lua '
+            local port = ngx.var.port
+            function socktest(port, pool)
+                local sock = ngx.socket.tcp()
+                local ok, err = sock:connect("unix:$TEST_NGINX_HTML_DIR/nginx.sock", nil, { pool = pool })
+                if not ok then
+                    ngx.say("failed to connect: ", err)
+                    return
+                end
+                ngx.say("connected: ", ok, ", reused: ", sock:getreusedtimes())
+                local ok, err = sock:setkeepalive()
+                if not ok then
+                    ngx.say("failed to set reusable: ", err)
+                end
+            end
+            for i=1,10 do
+                socktest(port, "one")
+                socktest(port, "two")
+            end
+        ';
+    }
+--- request
+GET /t
+--- response_body_like
+(connected: 1, reused: \d+\s+)*
+
+