Skip to content

Commit fdf752d

Browse files
agentzhagentzh
committed
feature: added undocumented API for attaching arbitrary key-value pair data to a cosocket TCP connecton, which can survive connection pool reuse.
This API may change in the future without notice.
1 parent 7367a23 commit fdf752d

File tree

2 files changed

+267
-0
lines changed

2 files changed

+267
-0
lines changed

‎src/ngx_http_lua_socket_tcp.c‎

Lines changed: 244 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -221,6 +221,9 @@ static char ngx_http_lua_ssl_session_metatable_key;
221221
#endif
222222

223223

224+
#define ngx_http_lua_tcp_socket_metatable_literal_key "__tcp_cosocket_mt"
225+
226+
224227
void
225228
ngx_http_lua_inject_socket_tcp_api(ngx_log_t *log, lua_State *L)
226229
{
@@ -355,6 +358,12 @@ ngx_http_lua_inject_socket_tcp_api(ngx_log_t *log, lua_State *L)
355358
lua_pushvalue(L, -1);
356359
lua_setfield(L, -2, "__index");
357360
lua_rawset(L, LUA_REGISTRYINDEX);
361+
362+
lua_pushliteral(L, ngx_http_lua_tcp_socket_metatable_literal_key);
363+
lua_pushlightuserdata(L, ngx_http_lua_lightudata_mask(
364+
tcp_socket_metatable_key));
365+
lua_rawget(L, LUA_REGISTRYINDEX);
366+
lua_rawset(L, LUA_REGISTRYINDEX);
358367
/* }}} */
359368

360369
/* {{{upstream userdata metatable */
@@ -5382,6 +5391,8 @@ ngx_http_lua_socket_tcp_setkeepalive(lua_State *L)
53825391
item->socklen = pc->socklen;
53835392
ngx_memcpy(&item->sockaddr, pc->sockaddr, pc->socklen);
53845393
item->reused = u->reused;
5394+
item->udata_queue = u->udata_queue;
5395+
u->udata_queue = NULL;
53855396

53865397
if (c->read->ready) {
53875398
rc = ngx_http_lua_socket_keepalive_close_handler(c->read);
@@ -5457,6 +5468,8 @@ ngx_http_lua_get_keepalive_peer(ngx_http_request_t *r,
54575468
pc->cached = 1;
54585469

54595470
u->reused = item->reused + 1;
5471+
u->udata_queue = item->udata_queue;
5472+
item->udata_queue = NULL;
54605473

54615474
#if 1
54625475
u->write_event_handler = ngx_http_lua_socket_dummy_handler;
@@ -6151,4 +6164,235 @@ ngx_http_lua_cleanup_conn_pools(lua_State *L)
61516164
lua_pop(L, 1);
61526165
}
61536166

6167+
6168+
int
6169+
ngx_http_lua_ffi_socket_tcp_init_udata_queue(
6170+
ngx_http_lua_socket_tcp_upstream_t *u, int capacity, char **err_msg)
6171+
{
6172+
int i, max_size;
6173+
ngx_pool_t *pool;
6174+
ngx_http_lua_socket_udata_queue_t *udata_queue;
6175+
ngx_http_lua_socket_node_t *node;
6176+
6177+
pool = u->peer.connection->pool;
6178+
6179+
if (u->udata_queue == NULL) {
6180+
max_size = capacity;
6181+
if (max_size == 0) {
6182+
max_size = 4;
6183+
}
6184+
6185+
udata_queue = ngx_palloc(pool,
6186+
sizeof(ngx_http_lua_socket_udata_queue_t) +
6187+
sizeof(ngx_http_lua_socket_node_t) * max_size);
6188+
6189+
if (udata_queue == NULL) {
6190+
*err_msg = "no memory";
6191+
return NGX_ERROR;
6192+
}
6193+
6194+
udata_queue->pool = pool;
6195+
udata_queue->capacity = capacity;
6196+
udata_queue->len = 0;
6197+
ngx_queue_init(&udata_queue->queue);
6198+
ngx_queue_init(&udata_queue->free);
6199+
6200+
node = (ngx_http_lua_socket_node_t *) (udata_queue + 1);
6201+
6202+
for (i = 0; i < max_size; i++) {
6203+
ngx_queue_insert_head(&udata_queue->free, &node->queue);
6204+
node++;
6205+
}
6206+
6207+
u->udata_queue = udata_queue;
6208+
6209+
ngx_log_debug3(NGX_LOG_DEBUG_HTTP, u->request->connection->log, 0,
6210+
"init udata_queue %uD, cosocket %p udata %p",
6211+
capacity, u, udata_queue);
6212+
}
6213+
6214+
return NGX_OK;
6215+
}
6216+
6217+
6218+
int
6219+
ngx_http_lua_ffi_socket_tcp_count_udata(ngx_http_lua_socket_tcp_upstream_t *u)
6220+
{
6221+
/* return NGX_ERROR (-1) for missing udata_queue to
6222+
* distinguish it from empty udata_queue */
6223+
if (u->udata_queue == NULL) {
6224+
return NGX_ERROR;
6225+
}
6226+
6227+
return u->udata_queue->len;
6228+
}
6229+
6230+
6231+
int
6232+
ngx_http_lua_ffi_socket_tcp_add_udata(ngx_http_lua_socket_tcp_upstream_t *u,
6233+
uint64_t key, uint64_t value, uint64_t *evicted_key,
6234+
uint64_t *evicted_value, char **err_msg)
6235+
{
6236+
int evicted = 0;
6237+
ngx_pool_t *pool;
6238+
ngx_http_lua_socket_node_t *node = NULL;
6239+
ngx_queue_t *q, *uqueue;
6240+
6241+
pool = u->peer.connection->pool;
6242+
6243+
if (u->udata_queue == NULL) {
6244+
*err_msg = "no udata queue";
6245+
return NGX_ERROR;
6246+
}
6247+
6248+
uqueue = &u->udata_queue->queue;
6249+
6250+
for (q = ngx_queue_head(uqueue);
6251+
q != ngx_queue_sentinel(uqueue);
6252+
q = ngx_queue_next(q))
6253+
{
6254+
node = ngx_queue_data(q, ngx_http_lua_socket_node_t, queue);
6255+
6256+
if (node->key == key) {
6257+
/* key exists */
6258+
ngx_log_debug3(NGX_LOG_DEBUG_HTTP, u->request->connection->log, 0,
6259+
"found %uD, cosocket %p udata %p",
6260+
key, u, u->udata_queue);
6261+
ngx_queue_remove(q);
6262+
node->value = value;
6263+
6264+
break;
6265+
}
6266+
}
6267+
6268+
if (q == ngx_queue_sentinel(uqueue)) {
6269+
6270+
if (u->udata_queue->capacity
6271+
&& u->udata_queue->capacity == u->udata_queue->len)
6272+
{
6273+
/* evict key */
6274+
q = ngx_queue_last(uqueue);
6275+
node = ngx_queue_data(q, ngx_http_lua_socket_node_t, queue);
6276+
ngx_queue_remove(q);
6277+
ngx_log_debug4(NGX_LOG_DEBUG_HTTP, u->request->connection->log, 0,
6278+
"evict %uD for %uD, cosocket %p udata %p",
6279+
node->key, key, u, u->udata_queue);
6280+
*evicted_key = node->key;
6281+
*evicted_value = node->value;
6282+
evicted = 1;
6283+
6284+
} else {
6285+
/* insert key */
6286+
ngx_log_debug3(NGX_LOG_DEBUG_HTTP, u->request->connection->log, 0,
6287+
"insert %uD, cosocket %p udata %p",
6288+
key, u, u->udata_queue);
6289+
6290+
if (!ngx_queue_empty(&u->udata_queue->free)) {
6291+
q = ngx_queue_head(&u->udata_queue->free);
6292+
node = ngx_queue_data(q, ngx_http_lua_socket_node_t, queue);
6293+
ngx_queue_remove(q);
6294+
ngx_log_debug3(NGX_LOG_DEBUG_HTTP, u->request->connection->log,
6295+
0, "reuse free node %p, cosocket %p udata %p",
6296+
node, u, u->udata_queue);
6297+
6298+
} else {
6299+
node = ngx_palloc(pool, sizeof(ngx_http_lua_socket_node_t));
6300+
if (node == NULL) {
6301+
goto nomem;
6302+
}
6303+
6304+
ngx_log_debug3(NGX_LOG_DEBUG_HTTP, u->request->connection->log,
6305+
0, "allocate new node %p, cosocket %p udata %p",
6306+
node, u, u->udata_queue);
6307+
}
6308+
6309+
u->udata_queue->len++;
6310+
}
6311+
6312+
node->key = key;
6313+
node->value = value;
6314+
}
6315+
6316+
ngx_queue_insert_head(uqueue, &node->queue);
6317+
return evicted ? NGX_DONE : NGX_OK;
6318+
6319+
nomem:
6320+
6321+
*err_msg = "no memory";
6322+
return NGX_ERROR;
6323+
}
6324+
6325+
6326+
int
6327+
ngx_http_lua_ffi_socket_tcp_get_udata(ngx_http_lua_socket_tcp_upstream_t *u,
6328+
uint64_t key, uint64_t *value, char **err_msg)
6329+
{
6330+
ngx_http_lua_socket_node_t *node;
6331+
ngx_queue_t *q, *uqueue;
6332+
6333+
if (u->udata_queue == NULL) {
6334+
*err_msg = "no udata queue";
6335+
return NGX_ERROR;
6336+
}
6337+
6338+
uqueue = &u->udata_queue->queue;
6339+
6340+
for (q = ngx_queue_head(uqueue);
6341+
q != ngx_queue_sentinel(uqueue);
6342+
q = ngx_queue_next(q))
6343+
{
6344+
node = ngx_queue_data(q, ngx_http_lua_socket_node_t, queue);
6345+
6346+
if (node->key == key) {
6347+
ngx_log_debug3(NGX_LOG_DEBUG_HTTP, u->request->connection->log, 0,
6348+
"found %uD, cosocket %p udata %p",
6349+
key, u, u->udata_queue);
6350+
ngx_queue_remove(q);
6351+
ngx_queue_insert_head(uqueue, &node->queue);
6352+
*value = node->value;
6353+
return NGX_OK;
6354+
}
6355+
}
6356+
6357+
*err_msg = "not found";
6358+
return NGX_ERROR;
6359+
}
6360+
6361+
6362+
int
6363+
ngx_http_lua_ffi_socket_tcp_del_udata(ngx_http_lua_socket_tcp_upstream_t *u,
6364+
uint64_t key, char **err_msg)
6365+
{
6366+
ngx_http_lua_socket_node_t *node;
6367+
ngx_queue_t *q, *uqueue;
6368+
6369+
if (u->udata_queue == NULL) {
6370+
*err_msg = "no udata queue";
6371+
return NGX_ERROR;
6372+
}
6373+
6374+
uqueue = &u->udata_queue->queue;
6375+
6376+
for (q = ngx_queue_head(uqueue);
6377+
q != ngx_queue_sentinel(uqueue);
6378+
q = ngx_queue_next(q))
6379+
{
6380+
node = ngx_queue_data(q, ngx_http_lua_socket_node_t, queue);
6381+
6382+
if (node->key == key) {
6383+
ngx_log_debug3(NGX_LOG_DEBUG_HTTP, u->request->connection->log, 0,
6384+
"delete %uD, cosocket %p udata %p",
6385+
key, u, u->udata_queue);
6386+
ngx_queue_remove(q);
6387+
ngx_queue_insert_head(&u->udata_queue->free, &node->queue);
6388+
u->udata_queue->len--;
6389+
return NGX_OK;
6390+
}
6391+
}
6392+
6393+
*err_msg = "not found";
6394+
return NGX_ERROR;
6395+
}
6396+
6397+
61546398
/* vi:set ft=c ts=4 sw=4 et fdm=marker: */

‎src/ngx_http_lua_socket_tcp.h‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,10 @@ typedef struct ngx_http_lua_socket_tcp_upstream_s
2626
ngx_http_lua_socket_tcp_upstream_t;
2727

2828

29+
typedef struct ngx_http_lua_socket_udata_queue_s
30+
ngx_http_lua_socket_udata_queue_t;
31+
32+
2933
typedef
3034
int (*ngx_http_lua_socket_tcp_retval_handler)(ngx_http_request_t *r,
3135
ngx_http_lua_socket_tcp_upstream_t *u, lua_State *L);
@@ -79,6 +83,8 @@ struct ngx_http_lua_socket_tcp_upstream_s {
7983
ngx_http_lua_socket_tcp_upstream_handler_pt read_event_handler;
8084
ngx_http_lua_socket_tcp_upstream_handler_pt write_event_handler;
8185

86+
ngx_http_lua_socket_udata_queue_t *udata_queue;
87+
8288
ngx_http_lua_socket_pool_t *socket_pool;
8389

8490
ngx_http_lua_loc_conf_t *conf;
@@ -166,9 +172,26 @@ typedef struct {
166172

167173
ngx_uint_t reused;
168174

175+
ngx_http_lua_socket_udata_queue_t *udata_queue;
169176
} ngx_http_lua_socket_pool_item_t;
170177

171178

179+
struct ngx_http_lua_socket_udata_queue_s {
180+
ngx_pool_t *pool;
181+
ngx_queue_t queue;
182+
ngx_queue_t free;
183+
int len;
184+
int capacity;
185+
};
186+
187+
188+
typedef struct {
189+
ngx_queue_t queue;
190+
uint64_t key;
191+
uint64_t value;
192+
} ngx_http_lua_socket_node_t;
193+
194+
172195
void ngx_http_lua_inject_socket_tcp_api(ngx_log_t *log, lua_State *L);
173196
void ngx_http_lua_inject_req_socket_api(lua_State *L);
174197
void ngx_http_lua_cleanup_conn_pools(lua_State *L);

0 commit comments

Comments
 (0)