Skip to content

Commit 3d0bbed

Browse files
author
Pablo Polvorin
committed
Ability to attach arbitrary data to cosocket instances
Usage: sock:binddata(strkey, value) sock:getbounddata(strkey) Impl: A table is created and a reference to it stored in the ngx_http_lua_socket_tcp_upstream_t structure (on demand, on the first call to binddata() on that socket). binddata() just put the data into that table, and getbounddata retrieve it. Table is dereferenced when the socket is freed
1 parent bda5b97 commit 3d0bbed

File tree

3 files changed

+269
-0
lines changed

3 files changed

+269
-0
lines changed

‎src/ngx_http_lua_socket_tcp.c

Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -93,6 +93,8 @@ static int ngx_http_lua_socket_cleanup_compiled_pattern(lua_State *L);
9393
static int ngx_http_lua_req_socket(lua_State *L);
9494
static void ngx_http_lua_req_socket_rev_handler(ngx_http_request_t *r);
9595
static int ngx_http_lua_socket_tcp_getreusedtimes(lua_State *L);
96+
static int ngx_http_lua_socket_tcp_binddata(lua_State *L);
97+
static int ngx_http_lua_socket_tcp_getbounddata(lua_State *L);
9698
static int ngx_http_lua_socket_tcp_setkeepalive(lua_State *L);
9799
static ngx_int_t ngx_http_lua_get_keepalive_peer(ngx_http_request_t *r,
98100
lua_State *L, int key_index,
@@ -300,6 +302,12 @@ ngx_http_lua_inject_socket_tcp_api(ngx_log_t *log, lua_State *L)
300302
lua_pushcfunction(L, ngx_http_lua_socket_tcp_setkeepalive);
301303
lua_setfield(L, -2, "setkeepalive");
302304

305+
lua_pushcfunction(L, ngx_http_lua_socket_tcp_binddata);
306+
lua_setfield(L, -2, "binddata");
307+
308+
lua_pushcfunction(L, ngx_http_lua_socket_tcp_getbounddata);
309+
lua_setfield(L, -2, "getbounddata");
310+
303311
lua_pushvalue(L, -1);
304312
lua_setfield(L, -2, "__index");
305313
lua_rawset(L, LUA_REGISTRYINDEX);
@@ -610,6 +618,8 @@ ngx_http_lua_socket_tcp_connect(lua_State *L)
610618

611619
/* rc == NGX_DECLINED */
612620

621+
u->bind_tbl_ref = LUA_NOREF;
622+
613623
ngx_memzero(&url, sizeof(ngx_url_t));
614624

615625
url.url.len = host.len;
@@ -4332,6 +4342,68 @@ ngx_http_lua_socket_tcp_getreusedtimes(lua_State *L)
43324342
return 1;
43334343
}
43344344

4345+
static int
4346+
ngx_http_lua_socket_tcp_getbounddata(lua_State *L)
4347+
{
4348+
ngx_http_lua_socket_tcp_upstream_t *u;
4349+
4350+
if (lua_gettop(L) != 2) {
4351+
return luaL_error(L, "expecting 2 argument "
4352+
"(including the object), but got %d", lua_gettop(L));
4353+
}
4354+
4355+
luaL_checktype(L, 1, LUA_TTABLE);
4356+
luaL_checktype(L, 2, LUA_TSTRING);
4357+
4358+
lua_rawgeti(L, 1, SOCKET_CTX_INDEX);
4359+
u = lua_touserdata(L, -1);
4360+
lua_rawgeti(L, LUA_REGISTRYINDEX, u->bind_tbl_ref);
4361+
switch (lua_type(L, -1)) {
4362+
case LUA_TTABLE:
4363+
lua_pushvalue(L,2);
4364+
lua_gettable(L,-2);
4365+
break;
4366+
case LUA_TNIL:
4367+
break;
4368+
default:
4369+
return luaL_error(L, "no bind table found");
4370+
}
4371+
return 1;
4372+
}
4373+
4374+
static int
4375+
ngx_http_lua_socket_tcp_binddata(lua_State *L)
4376+
{
4377+
ngx_http_lua_socket_tcp_upstream_t *u;
4378+
4379+
if (lua_gettop(L) != 3) {
4380+
return luaL_error(L, "expecting 3 argument "
4381+
"(including the object), but got %d", lua_gettop(L));
4382+
}
4383+
4384+
luaL_checktype(L, 1, LUA_TTABLE);
4385+
luaL_checktype(L, 2, LUA_TSTRING);
4386+
4387+
lua_rawgeti(L, 1, SOCKET_CTX_INDEX);
4388+
u = lua_touserdata(L, -1);
4389+
4390+
if (u->bind_tbl_ref == LUA_NOREF) {
4391+
lua_createtable(L, 0, 1);
4392+
lua_pushvalue(L,2); //key
4393+
lua_pushvalue(L,3); //value
4394+
lua_rawset(L,-3);
4395+
u->bind_tbl_ref = luaL_ref(L, LUA_REGISTRYINDEX);
4396+
lua_pushinteger(L, 1);
4397+
return 1;
4398+
}
4399+
4400+
lua_rawgeti(L, LUA_REGISTRYINDEX, u->bind_tbl_ref);
4401+
lua_pushvalue(L,2); //key
4402+
lua_pushvalue(L,3); //value
4403+
lua_rawset(L,-3);
4404+
lua_pushinteger(L, 1);
4405+
return 1;
4406+
}
43354407

43364408
static int ngx_http_lua_socket_tcp_setkeepalive(lua_State *L)
43374409
{
@@ -4577,6 +4649,7 @@ static int ngx_http_lua_socket_tcp_setkeepalive(lua_State *L)
45774649
item->socklen = pc->socklen;
45784650
ngx_memcpy(&item->sockaddr, pc->sockaddr, pc->socklen);
45794651
item->reused = u->reused;
4652+
item->bind_tbl_ref = u->bind_tbl_ref;
45804653

45814654
if (c->read->ready) {
45824655
rc = ngx_http_lua_socket_keepalive_close_handler(c->read);
@@ -4667,6 +4740,7 @@ ngx_http_lua_get_keepalive_peer(ngx_http_request_t *r, lua_State *L,
46674740
pc->cached = 1;
46684741

46694742
u->reused = item->reused + 1;
4743+
u->bind_tbl_ref = item->bind_tbl_ref;
46704744

46714745
#if 1
46724746
u->write_event_handler = ngx_http_lua_socket_dummy_handler;
@@ -4843,6 +4917,7 @@ ngx_http_lua_socket_tcp_upstream_destroy(lua_State *L)
48434917
return 0;
48444918
}
48454919

4920+
luaL_unref(L, LUA_REGISTRYINDEX, u->bind_tbl_ref);
48464921
if (u->cleanup) {
48474922
ngx_http_lua_socket_tcp_cleanup(u); /* it will clear u->cleanup */
48484923
}

‎src/ngx_http_lua_socket_tcp.h

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -89,6 +89,7 @@ struct ngx_http_lua_socket_tcp_upstream_s {
8989
ngx_http_lua_co_ctx_t *write_co_ctx;
9090

9191
ngx_uint_t reused;
92+
ngx_int_t bind_tbl_ref;
9293

9394
#if (NGX_HTTP_SSL)
9495
ngx_str_t ssl_name;
@@ -142,6 +143,7 @@ typedef struct {
142143
struct sockaddr_storage sockaddr;
143144

144145
ngx_uint_t reused;
146+
ngx_int_t bind_tbl_ref;
145147

146148
} ngx_http_lua_socket_pool_item_t;
147149

‎t/130-socket-binddata.t

Lines changed: 192 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,192 @@
1+
# vim:set ft= ts=4 sw=4 et fdm=marker:
2+
3+
use lib 'lib';
4+
use Test::Nginx::Socket::Lua;
5+
6+
repeat_each(1);
7+
8+
plan tests => 4 + 6 + 2 + 2; #needed? how to calculate it when each tests is different? :(
9+
10+
11+
our $HtmlDir = html_dir;
12+
13+
$ENV{TEST_NGINX_MEMCACHED_PORT} ||= 11211;
14+
$ENV{TEST_NGINX_HTML_DIR} = $HtmlDir;
15+
$ENV{TEST_NGINX_REDIS_PORT} ||= 6379;
16+
17+
$ENV{LUA_PATH} ||=
18+
'/usr/local/openresty-debug/lualib/?.lua;/usr/local/openresty/lualib/?.lua;;';
19+
20+
no_long_string();
21+
22+
log_level 'debug';
23+
24+
no_shuffle();
25+
26+
27+
run_tests();
28+
29+
__DATA__
30+
31+
=== TEST 1: sanity
32+
--- no_check_leak
33+
--- http_config eval
34+
"lua_package_path '$::HtmlDir/?.lua;./?.lua';"
35+
--- config
36+
location /t {
37+
set $port $TEST_NGINX_MEMCACHED_PORT;
38+
content_by_lua '
39+
local port = ngx.var.port
40+
local sock = ngx.socket.tcp()
41+
local ok, err = sock:connect("127.0.0.1", port)
42+
if not ok then
43+
ngx.say("failed to connect: ", err)
44+
return
45+
end
46+
ngx.say("connected: ", ok, ", reused: ", sock:getreusedtimes())
47+
local ok, err = sock:getbounddata("test")
48+
if not ok then
49+
sock:binddata("test", "a")
50+
ngx.say("bind data not found")
51+
else
52+
ngx.say("bind data ", ok)
53+
end
54+
55+
local ok, err = sock:setkeepalive()
56+
if not ok then
57+
ngx.say("failed to set reusable: ", err)
58+
end
59+
';
60+
}
61+
--- request eval
62+
["GET /t", "GET /t"]
63+
64+
--- response_body eval
65+
[
66+
"connected: 1, reused: 0
67+
bind data not found
68+
",
69+
"connected: 1, reused: 1
70+
bind data a
71+
"]
72+
73+
74+
=== TEST 2: unbind data
75+
--- no_check_leak
76+
--- http_config eval
77+
"lua_package_path '$::HtmlDir/?.lua;./?.lua';"
78+
--- config
79+
location /t {
80+
set $port $TEST_NGINX_MEMCACHED_PORT;
81+
content_by_lua '
82+
local port = ngx.var.port
83+
local sock = ngx.socket.tcp()
84+
local ok, err = sock:connect("127.0.0.1", port)
85+
if not ok then
86+
ngx.say("failed to connect: ", err)
87+
return
88+
end
89+
ngx.say("connected: ", ok, ", reused: ", sock:getreusedtimes())
90+
local ok, err = sock:getbounddata("test")
91+
if not ok then
92+
sock:binddata("test", "a")
93+
ngx.say("bind data not found")
94+
else
95+
ngx.say("bind data ", ok)
96+
sock:binddata("test", nil)
97+
end
98+
99+
local ok, err = sock:setkeepalive()
100+
if not ok then
101+
ngx.say("failed to set reusable: ", err)
102+
end
103+
';
104+
}
105+
--- request eval
106+
["GET /t", "GET /t", "GET /t"]
107+
108+
--- response_body eval
109+
[
110+
"connected: 1, reused: 0
111+
bind data not found
112+
",
113+
"connected: 1, reused: 1
114+
bind data a
115+
",
116+
"connected: 1, reused: 2
117+
bind data not found
118+
"]
119+
120+
121+
122+
=== TEST 3: bound data is gc with socket
123+
# For TEST_NGINX_CHECK_LEAK
124+
--- http_config eval
125+
"lua_package_path '$::HtmlDir/?.lua;./?.lua';"
126+
--- config
127+
location /t {
128+
set $port $TEST_NGINX_MEMCACHED_PORT;
129+
content_by_lua '
130+
local port = ngx.var.port
131+
local sock = ngx.socket.tcp()
132+
local ok, err = sock:connect("127.0.0.1", port)
133+
if not ok then
134+
ngx.say("failed to connect: ", err)
135+
return
136+
end
137+
ngx.say("connected: ", ok, ", reused: ", sock:getreusedtimes())
138+
local ok, err = sock:getbounddata("test")
139+
if not ok then
140+
sock:binddata("test", string.rep("a",60))
141+
ngx.say("bind data not found")
142+
end
143+
local ok, err = sock:getbounddata("test")
144+
ngx.say("bind data ", ok)
145+
';
146+
}
147+
--- request
148+
GET /t
149+
150+
--- response_body
151+
connected: 1, reused: 0
152+
bind data not found
153+
bind data aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa
154+
155+
=== TEST 4: test bound data is GC when not referenced
156+
# For TEST_NGINX_CHECK_LEAK
157+
--- http_config eval
158+
"lua_package_path '$::HtmlDir/?.lua;./?.lua';"
159+
--- config
160+
location /t {
161+
set $port $TEST_NGINX_MEMCACHED_PORT;
162+
content_by_lua '
163+
local port = ngx.var.port
164+
local sock = ngx.socket.tcp()
165+
local ok, err = sock:connect("127.0.0.1", port)
166+
if not ok then
167+
ngx.say("failed to connect: ", err)
168+
return
169+
end
170+
local ok, err = sock:getbounddata("test")
171+
if not ok then
172+
sock:binddata("test", string.rep("a",60))
173+
ok, err = sock:getbounddata("test")
174+
end
175+
ngx.say("bind data ", ok)
176+
sock:binddata("test", nil)
177+
local ok, err = sock:getbounddata("test")
178+
if not ok then
179+
ngx.say("removed")
180+
end
181+
local ok, err = sock:setkeepalive()
182+
if not ok then
183+
ngx.say("failed to set reusable: ", err)
184+
end
185+
';
186+
}
187+
--- request
188+
GET /t
189+
190+
--- response_body
191+
bind data aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa
192+
removed

0 commit comments

Comments
 (0)