Skip to content

Commit 1d8d157

Browse files
agentzhagentzh
committed
feature: added an optional "optional table" to tcpsock:connect() which accepts a "pool" option to allow the user specify a custom pool name intead of the automatically generated one based on the host-port pair or the socket file path. thanks Brian Akins for the patches.
1 parent 06abbea commit 1d8d157

File tree

2 files changed

+276
-21
lines changed

2 files changed

+276
-21
lines changed

‎src/ngx_http_lua_socket_tcp.c

Lines changed: 52 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -253,13 +253,16 @@ ngx_http_lua_socket_tcp_connect(lua_State *L)
253253
ngx_http_lua_loc_conf_t *llcf;
254254
ngx_peer_connection_t *pc;
255255
int timeout;
256+
unsigned custom_pool;
257+
int key_index;
258+
const char *msg;
256259

257260
ngx_http_lua_socket_tcp_upstream_t *u;
258261

259262
n = lua_gettop(L);
260-
if (n != 2 && n != 3) {
261-
return luaL_error(L, "ngx.socket connect: expecting 2 or 3 arguments "
262-
"(including the object), but seen %d", n);
263+
if (n != 2 && n != 3 && n != 4) {
264+
return luaL_error(L, "ngx.socket connect: expecting 2, 3, or 4 "
265+
"arguments (including the object), but seen %d", n);
263266
}
264267

265268
lua_pushlightuserdata(L, &ngx_http_lua_request_key);
@@ -294,6 +297,41 @@ ngx_http_lua_socket_tcp_connect(lua_State *L)
294297
ngx_memcpy(host.data, p, len);
295298
host.data[len] = '\0';
296299

300+
key_index = 2;
301+
custom_pool = 0;
302+
303+
if (lua_type(L, n) == LUA_TTABLE) {
304+
305+
/* found the last optional option table */
306+
307+
lua_getfield(L, n, "pool");
308+
309+
switch (lua_type(L, -1)) {
310+
case LUA_TSTRING:
311+
case LUA_TNUMBER:
312+
custom_pool = 1;
313+
314+
lua_pushvalue(L, -1);
315+
lua_rawseti(L, 1, SOCKET_KEY_INDEX);
316+
317+
key_index = n + 1;
318+
319+
break;
320+
321+
case LUA_TNIL:
322+
/* ignore it */
323+
break;
324+
325+
default:
326+
msg = lua_pushfstring(L, "bad \"pool\" option type: %s",
327+
luaL_typename(L, -1));
328+
luaL_argerror(L, n, msg);
329+
break;
330+
}
331+
332+
n--;
333+
}
334+
297335
if (n == 3) {
298336
port = luaL_checkinteger(L, 3);
299337

@@ -303,20 +341,24 @@ ngx_http_lua_socket_tcp_connect(lua_State *L)
303341
return 2;
304342
}
305343

306-
lua_pushliteral(L, ":");
307-
lua_insert(L, 3);
308-
lua_concat(L, 3);
344+
if (!custom_pool) {
345+
lua_pushliteral(L, ":");
346+
lua_insert(L, 3);
347+
lua_concat(L, 3);
348+
}
309349

310350
dd("socket key: %s", lua_tostring(L, -1));
311351

312352
} else { /* n == 2 */
313353
port = 0;
314354
}
315355

316-
/* the key's index is 2 */
356+
if (!custom_pool) {
357+
/* the key's index is 2 */
317358

318-
lua_pushvalue(L, -1);
319-
lua_rawseti(L, 1, SOCKET_KEY_INDEX);
359+
lua_pushvalue(L, 2);
360+
lua_rawseti(L, 1, SOCKET_KEY_INDEX);
361+
}
320362

321363
lua_rawgeti(L, 1, SOCKET_CTX_INDEX);
322364
u = lua_touserdata(L, -1);
@@ -390,7 +432,7 @@ ngx_http_lua_socket_tcp_connect(lua_State *L)
390432

391433
r->connection->single_connection = 0;
392434

393-
rc = ngx_http_lua_get_keepalive_peer(r, L, 2, u);
435+
rc = ngx_http_lua_get_keepalive_peer(r, L, key_index, u);
394436

395437
if (rc == NGX_OK) {
396438
lua_pushinteger(L, 1);

‎t/068-socket-keepalive.t

Lines changed: 224 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -3,9 +3,9 @@
33
use lib 'lib';
44
use Test::Nginx::Socket;
55

6-
repeat_each(2);
6+
#repeat_each(2);
77

8-
plan tests => repeat_each() * (blocks() * 5 + 8);
8+
plan tests => repeat_each() * (blocks() * 5 + 14);
99

1010
our $HtmlDir = html_dir;
1111

@@ -80,11 +80,11 @@ function go(port)
8080
end
8181
--- request
8282
GET /t
83-
--- response_body_like
84-
^connected: 1, reused: \d+
83+
--- response_body
84+
connected: 1, reused: 0
8585
request sent: 11
8686
received: OK
87-
connected: 1, reused: [1-9]\d*
87+
connected: 1, reused: 1
8888
request sent: 11
8989
received: OK
9090
--- no_error_log eval
@@ -153,11 +153,11 @@ function go(port, keepalive)
153153
sock:close()
154154
end
155155
end
156-
--- response_body_like
157-
^connected: 1, reused: \d+
156+
--- response_body
157+
connected: 1, reused: 0
158158
request sent: 11
159159
received: OK
160-
connected: 1, reused: [1-9]\d*
160+
connected: 1, reused: 1
161161
request sent: 11
162162
received: OK
163163
--- no_error_log
@@ -844,11 +844,11 @@ function go(path, port)
844844
ngx.say("failed to set reusable: ", err)
845845
end
846846
end
847-
--- response_body_like
848-
^connected: 1, reused: \d+
847+
--- response_body
848+
connected: 1, reused: 0
849849
request sent: 61
850850
received response of 119 bytes
851-
connected: 1, reused: [1-9]\d*
851+
connected: 1, reused: 1
852852
request sent: 61
853853
received response of 119 bytes
854854
--- no_error_log eval
@@ -964,3 +964,216 @@ Not found, dear...
964964
--- no_error_log
965965
[error]
966966

967+
968+
969+
=== TEST 14: custom pools (different pool for the same host:port) - tcp
970+
--- http_config eval
971+
"lua_package_path '$::HtmlDir/?.lua;./?.lua';"
972+
--- config
973+
location /t {
974+
set $port $TEST_NGINX_MEMCACHED_PORT;
975+
content_by_lua '
976+
local test = require "test"
977+
local port = ngx.var.port
978+
test.go(port, "A")
979+
test.go(port, "B")
980+
';
981+
}
982+
--- user_files
983+
>>> test.lua
984+
module("test", package.seeall)
985+
986+
function go(port, pool)
987+
local sock = ngx.socket.tcp()
988+
local ok, err = sock:connect("127.0.0.1", port, {pool = pool})
989+
if not ok then
990+
ngx.say("failed to connect: ", err)
991+
return
992+
end
993+
994+
ngx.say("connected: ", ok, ", reused: ", sock:getreusedtimes())
995+
996+
local ok, err = sock:setkeepalive()
997+
if not ok then
998+
ngx.say("failed to set reusable: ", err)
999+
end
1000+
end
1001+
--- request
1002+
GET /t
1003+
--- response_body
1004+
connected: 1, reused: 0
1005+
connected: 1, reused: 0
1006+
--- no_error_log eval
1007+
["[error]",
1008+
"lua tcp socket keepalive: free connection pool for ",
1009+
"lua tcp socket get keepalive peer: using connection"
1010+
]
1011+
--- error_log
1012+
lua tcp socket keepalive create connection pool for key "A"
1013+
lua tcp socket keepalive create connection pool for key "B"
1014+
1015+
1016+
1017+
=== TEST 15: custom pools (same pool for different host:port) - tcp
1018+
--- http_config eval
1019+
"lua_package_path '$::HtmlDir/?.lua;./?.lua';"
1020+
--- config
1021+
location /t {
1022+
set $port $TEST_NGINX_MEMCACHED_PORT;
1023+
content_by_lua '
1024+
local test = require "test"
1025+
local port = ngx.var.port
1026+
test.go($TEST_NGINX_MEMCACHED_PORT, "foo")
1027+
test.go($TEST_NGINX_CLIENT_PORT, "foo")
1028+
';
1029+
}
1030+
--- user_files
1031+
>>> test.lua
1032+
module("test", package.seeall)
1033+
1034+
function go(port, pool)
1035+
local sock = ngx.socket.tcp()
1036+
local ok, err = sock:connect("127.0.0.1", port, {pool = pool})
1037+
if not ok then
1038+
ngx.say("failed to connect: ", err)
1039+
return
1040+
end
1041+
1042+
ngx.say("connected: ", ok, ", reused: ", sock:getreusedtimes())
1043+
1044+
local ok, err = sock:setkeepalive()
1045+
if not ok then
1046+
ngx.say("failed to set reusable: ", err)
1047+
end
1048+
end
1049+
--- request
1050+
GET /t
1051+
--- response_body
1052+
connected: 1, reused: 0
1053+
connected: 1, reused: 1
1054+
--- no_error_log eval
1055+
["[error]",
1056+
"lua tcp socket keepalive: free connection pool for ",
1057+
]
1058+
--- error_log
1059+
lua tcp socket keepalive create connection pool for key "foo"
1060+
lua tcp socket get keepalive peer: using connection
1061+
1062+
1063+
1064+
=== TEST 16: custom pools (different pool for the same host:port) - unix
1065+
--- http_config eval
1066+
"
1067+
lua_package_path '$::HtmlDir/?.lua;./?.lua';
1068+
server {
1069+
listen unix:$::HtmlDir/nginx.sock;
1070+
default_type 'text/plain';
1071+
1072+
server_tokens off;
1073+
location /foo {
1074+
echo foo;
1075+
more_clear_headers Date;
1076+
}
1077+
}
1078+
"
1079+
--- config
1080+
location /t {
1081+
set $port $TEST_NGINX_MEMCACHED_PORT;
1082+
content_by_lua '
1083+
local test = require "test"
1084+
local path = "$TEST_NGINX_HTML_DIR/nginx.sock";
1085+
test.go(path, "A")
1086+
test.go(path, "B")
1087+
';
1088+
}
1089+
--- user_files
1090+
>>> test.lua
1091+
module("test", package.seeall)
1092+
1093+
function go(path, pool)
1094+
local sock = ngx.socket.tcp()
1095+
local ok, err = sock:connect("unix:" .. path, {pool = pool})
1096+
if not ok then
1097+
ngx.say("failed to connect: ", err)
1098+
return
1099+
end
1100+
1101+
ngx.say("connected: ", ok, ", reused: ", sock:getreusedtimes())
1102+
1103+
local ok, err = sock:setkeepalive()
1104+
if not ok then
1105+
ngx.say("failed to set reusable: ", err)
1106+
end
1107+
end
1108+
--- request
1109+
GET /t
1110+
--- response_body
1111+
connected: 1, reused: 0
1112+
connected: 1, reused: 0
1113+
--- no_error_log eval
1114+
["[error]",
1115+
"lua tcp socket keepalive: free connection pool for ",
1116+
"lua tcp socket get keepalive peer: using connection"
1117+
]
1118+
--- error_log
1119+
lua tcp socket keepalive create connection pool for key "A"
1120+
lua tcp socket keepalive create connection pool for key "B"
1121+
1122+
1123+
=== TEST 16: custom pools (same pool for the same host:port) - unix
1124+
--- http_config eval
1125+
"
1126+
lua_package_path '$::HtmlDir/?.lua;./?.lua';
1127+
server {
1128+
listen unix:$::HtmlDir/nginx.sock;
1129+
default_type 'text/plain';
1130+
1131+
server_tokens off;
1132+
location /foo {
1133+
echo foo;
1134+
more_clear_headers Date;
1135+
}
1136+
}
1137+
"
1138+
--- config
1139+
location /t {
1140+
set $port $TEST_NGINX_MEMCACHED_PORT;
1141+
content_by_lua '
1142+
local test = require "test"
1143+
local path = "$TEST_NGINX_HTML_DIR/nginx.sock";
1144+
test.go(path, "A")
1145+
test.go(path, "A")
1146+
';
1147+
}
1148+
--- user_files
1149+
>>> test.lua
1150+
module("test", package.seeall)
1151+
1152+
function go(path, pool)
1153+
local sock = ngx.socket.tcp()
1154+
local ok, err = sock:connect("unix:" .. path, {pool = pool})
1155+
if not ok then
1156+
ngx.say("failed to connect: ", err)
1157+
return
1158+
end
1159+
1160+
ngx.say("connected: ", ok, ", reused: ", sock:getreusedtimes())
1161+
1162+
local ok, err = sock:setkeepalive()
1163+
if not ok then
1164+
ngx.say("failed to set reusable: ", err)
1165+
end
1166+
end
1167+
--- request
1168+
GET /t
1169+
--- response_body
1170+
connected: 1, reused: 0
1171+
connected: 1, reused: 1
1172+
--- no_error_log eval
1173+
["[error]",
1174+
"lua tcp socket keepalive: free connection pool for ",
1175+
]
1176+
--- error_log
1177+
lua tcp socket keepalive create connection pool for key "A"
1178+
lua tcp socket get keepalive peer: using connection
1179+

0 commit comments

Comments
 (0)