Line data Source code
1 :
2 : /*
3 : * Copyright (C) Xiaozhe Wang (chaoslawful)
4 : * Copyright (C) Yichun Zhang (agentzh)
5 : */
6 :
7 :
8 : #ifndef DDEBUG
9 : #define DDEBUG 0
10 : #endif
11 : #include "ddebug.h"
12 :
13 :
14 : #include "ngx_http_lua_coroutine.h"
15 : #include "ngx_http_lua_util.h"
16 : #include "ngx_http_lua_probe.h"
17 :
18 :
19 : /*
20 : * Design:
21 : *
22 : * In order to support using ngx.* API in Lua coroutines, we have to create
23 : * new coroutine in the main coroutine instead of the calling coroutine
24 : */
25 :
26 :
27 : static int ngx_http_lua_coroutine_create(lua_State *L);
28 : static int ngx_http_lua_coroutine_resume(lua_State *L);
29 : static int ngx_http_lua_coroutine_yield(lua_State *L);
30 : static int ngx_http_lua_coroutine_status(lua_State *L);
31 :
32 :
33 : static const ngx_str_t
34 : ngx_http_lua_co_status_names[] =
35 : {
36 : ngx_string("running"),
37 : ngx_string("suspended"),
38 : ngx_string("normal"),
39 : ngx_string("dead"),
40 : ngx_string("zombie")
41 : };
42 :
43 :
44 :
45 : static int
46 0 : ngx_http_lua_coroutine_create(lua_State *L)
47 : {
48 : ngx_http_request_t *r;
49 : ngx_http_lua_ctx_t *ctx;
50 :
51 0 : r = ngx_http_lua_get_req(L);
52 0 : if (r == NULL) {
53 0 : return luaL_error(L, "no request found");
54 : }
55 :
56 0 : ctx = ngx_http_get_module_ctx(r, ngx_http_lua_module);
57 0 : if (ctx == NULL) {
58 0 : return luaL_error(L, "no request ctx found");
59 : }
60 :
61 0 : return ngx_http_lua_coroutine_create_helper(L, r, ctx, NULL);
62 : }
63 :
64 :
65 : int
66 0 : ngx_http_lua_coroutine_create_helper(lua_State *L, ngx_http_request_t *r,
67 : ngx_http_lua_ctx_t *ctx, ngx_http_lua_co_ctx_t **pcoctx)
68 : {
69 : lua_State *vm; /* the Lua VM */
70 : lua_State *co; /* new coroutine to be created */
71 : ngx_http_lua_co_ctx_t *coctx; /* co ctx for the new coroutine */
72 :
73 0 : luaL_argcheck(L, lua_isfunction(L, 1) && !lua_iscfunction(L, 1), 1,
74 : "Lua function expected");
75 :
76 0 : ngx_http_lua_check_context(L, ctx, NGX_HTTP_LUA_CONTEXT_REWRITE
77 : | NGX_HTTP_LUA_CONTEXT_ACCESS
78 : | NGX_HTTP_LUA_CONTEXT_CONTENT
79 : | NGX_HTTP_LUA_CONTEXT_TIMER
80 : | NGX_HTTP_LUA_CONTEXT_SSL_CERT
81 : | NGX_HTTP_LUA_CONTEXT_SSL_SESS_FETCH);
82 :
83 0 : vm = ngx_http_lua_get_lua_vm(r, ctx);
84 :
85 : /* create new coroutine on root Lua state, so it always yields
86 : * to main Lua thread
87 : */
88 0 : co = lua_newthread(vm);
89 :
90 : ngx_http_lua_probe_user_coroutine_create(r, L, co);
91 :
92 0 : coctx = ngx_http_lua_get_co_ctx(co, ctx);
93 0 : if (coctx == NULL) {
94 0 : coctx = ngx_http_lua_create_co_ctx(r, ctx);
95 0 : if (coctx == NULL) {
96 0 : return luaL_error(L, "no memory");
97 : }
98 :
99 : } else {
100 0 : ngx_memzero(coctx, sizeof(ngx_http_lua_co_ctx_t));
101 0 : coctx->co_ref = LUA_NOREF;
102 : }
103 :
104 0 : coctx->co = co;
105 0 : coctx->co_status = NGX_HTTP_LUA_CO_SUSPENDED;
106 :
107 : /* make new coroutine share globals of the parent coroutine.
108 : * NOTE: globals don't have to be separated! */
109 0 : ngx_http_lua_get_globals_table(L);
110 0 : lua_xmove(L, co, 1);
111 0 : ngx_http_lua_set_globals_table(co);
112 :
113 0 : lua_xmove(vm, L, 1); /* move coroutine from main thread to L */
114 :
115 0 : lua_pushvalue(L, 1); /* copy entry function to top of L*/
116 0 : lua_xmove(L, co, 1); /* move entry function from L to co */
117 :
118 0 : if (pcoctx) {
119 0 : *pcoctx = coctx;
120 : }
121 :
122 : #ifdef NGX_LUA_USE_ASSERT
123 0 : coctx->co_top = 1;
124 : #endif
125 :
126 0 : return 1; /* return new coroutine to Lua */
127 : }
128 :
129 :
130 : static int
131 0 : ngx_http_lua_coroutine_resume(lua_State *L)
132 : {
133 : lua_State *co;
134 : ngx_http_request_t *r;
135 : ngx_http_lua_ctx_t *ctx;
136 : ngx_http_lua_co_ctx_t *coctx;
137 : ngx_http_lua_co_ctx_t *p_coctx; /* parent co ctx */
138 :
139 0 : co = lua_tothread(L, 1);
140 :
141 0 : luaL_argcheck(L, co, 1, "coroutine expected");
142 :
143 0 : r = ngx_http_lua_get_req(L);
144 0 : if (r == NULL) {
145 0 : return luaL_error(L, "no request found");
146 : }
147 :
148 0 : ctx = ngx_http_get_module_ctx(r, ngx_http_lua_module);
149 0 : if (ctx == NULL) {
150 0 : return luaL_error(L, "no request ctx found");
151 : }
152 :
153 0 : ngx_http_lua_check_context(L, ctx, NGX_HTTP_LUA_CONTEXT_REWRITE
154 : | NGX_HTTP_LUA_CONTEXT_ACCESS
155 : | NGX_HTTP_LUA_CONTEXT_CONTENT
156 : | NGX_HTTP_LUA_CONTEXT_TIMER
157 : | NGX_HTTP_LUA_CONTEXT_SSL_CERT
158 : | NGX_HTTP_LUA_CONTEXT_SSL_SESS_FETCH);
159 :
160 0 : p_coctx = ctx->cur_co_ctx;
161 0 : if (p_coctx == NULL) {
162 0 : return luaL_error(L, "no parent co ctx found");
163 : }
164 :
165 0 : coctx = ngx_http_lua_get_co_ctx(co, ctx);
166 0 : if (coctx == NULL) {
167 0 : return luaL_error(L, "no co ctx found");
168 : }
169 :
170 : ngx_http_lua_probe_user_coroutine_resume(r, L, co);
171 :
172 0 : if (coctx->co_status != NGX_HTTP_LUA_CO_SUSPENDED) {
173 : dd("coroutine resume: %d", coctx->co_status);
174 :
175 0 : lua_pushboolean(L, 0);
176 0 : lua_pushfstring(L, "cannot resume %s coroutine",
177 0 : ngx_http_lua_co_status_names[coctx->co_status].data);
178 0 : return 2;
179 : }
180 :
181 0 : p_coctx->co_status = NGX_HTTP_LUA_CO_NORMAL;
182 :
183 0 : coctx->parent_co_ctx = p_coctx;
184 :
185 : dd("set coroutine to running");
186 0 : coctx->co_status = NGX_HTTP_LUA_CO_RUNNING;
187 :
188 0 : ctx->co_op = NGX_HTTP_LUA_USER_CORO_RESUME;
189 0 : ctx->cur_co_ctx = coctx;
190 :
191 : /* yield and pass args to main thread, and resume target coroutine from
192 : * there */
193 0 : return lua_yield(L, lua_gettop(L) - 1);
194 : }
195 :
196 :
197 : static int
198 0 : ngx_http_lua_coroutine_yield(lua_State *L)
199 : {
200 : ngx_http_request_t *r;
201 : ngx_http_lua_ctx_t *ctx;
202 : ngx_http_lua_co_ctx_t *coctx;
203 :
204 0 : r = ngx_http_lua_get_req(L);
205 0 : if (r == NULL) {
206 0 : return luaL_error(L, "no request found");
207 : }
208 :
209 0 : ctx = ngx_http_get_module_ctx(r, ngx_http_lua_module);
210 0 : if (ctx == NULL) {
211 0 : return luaL_error(L, "no request ctx found");
212 : }
213 :
214 0 : ngx_http_lua_check_context(L, ctx, NGX_HTTP_LUA_CONTEXT_REWRITE
215 : | NGX_HTTP_LUA_CONTEXT_ACCESS
216 : | NGX_HTTP_LUA_CONTEXT_CONTENT
217 : | NGX_HTTP_LUA_CONTEXT_TIMER
218 : | NGX_HTTP_LUA_CONTEXT_SSL_CERT
219 : | NGX_HTTP_LUA_CONTEXT_SSL_SESS_FETCH);
220 :
221 0 : coctx = ctx->cur_co_ctx;
222 :
223 0 : coctx->co_status = NGX_HTTP_LUA_CO_SUSPENDED;
224 :
225 0 : ctx->co_op = NGX_HTTP_LUA_USER_CORO_YIELD;
226 :
227 0 : if (!coctx->is_uthread && coctx->parent_co_ctx) {
228 : dd("set coroutine to running");
229 0 : coctx->parent_co_ctx->co_status = NGX_HTTP_LUA_CO_RUNNING;
230 :
231 : ngx_http_lua_probe_user_coroutine_yield(r, coctx->parent_co_ctx->co, L);
232 :
233 : } else {
234 : ngx_http_lua_probe_user_coroutine_yield(r, NULL, L);
235 : }
236 :
237 : /* yield and pass retvals to main thread,
238 : * and resume parent coroutine there */
239 0 : return lua_yield(L, lua_gettop(L));
240 : }
241 :
242 :
243 : void
244 18 : ngx_http_lua_inject_coroutine_api(ngx_log_t *log, lua_State *L)
245 : {
246 : int rc;
247 :
248 : /* new coroutine table */
249 18 : lua_createtable(L, 0 /* narr */, 14 /* nrec */);
250 :
251 : /* get old coroutine table */
252 18 : lua_getglobal(L, "coroutine");
253 :
254 : /* set running to the old one */
255 18 : lua_getfield(L, -1, "running");
256 18 : lua_setfield(L, -3, "running");
257 :
258 18 : lua_getfield(L, -1, "create");
259 18 : lua_setfield(L, -3, "_create");
260 :
261 18 : lua_getfield(L, -1, "resume");
262 18 : lua_setfield(L, -3, "_resume");
263 :
264 18 : lua_getfield(L, -1, "yield");
265 18 : lua_setfield(L, -3, "_yield");
266 :
267 18 : lua_getfield(L, -1, "status");
268 18 : lua_setfield(L, -3, "_status");
269 :
270 : /* pop the old coroutine */
271 18 : lua_pop(L, 1);
272 :
273 18 : lua_pushcfunction(L, ngx_http_lua_coroutine_create);
274 18 : lua_setfield(L, -2, "__create");
275 :
276 18 : lua_pushcfunction(L, ngx_http_lua_coroutine_resume);
277 18 : lua_setfield(L, -2, "__resume");
278 :
279 18 : lua_pushcfunction(L, ngx_http_lua_coroutine_yield);
280 18 : lua_setfield(L, -2, "__yield");
281 :
282 18 : lua_pushcfunction(L, ngx_http_lua_coroutine_status);
283 18 : lua_setfield(L, -2, "__status");
284 :
285 18 : lua_setglobal(L, "coroutine");
286 :
287 : /* inject coroutine APIs */
288 : {
289 18 : const char buf[] =
290 : "local keys = {'create', 'yield', 'resume', 'status'}\n"
291 : "local getfenv = getfenv\n"
292 : "for _, key in ipairs(keys) do\n"
293 : "local std = coroutine['_' .. key]\n"
294 : "local ours = coroutine['__' .. key]\n"
295 : "local raw_ctx = ngx._phase_ctx\n"
296 : "coroutine[key] = function (...)\n"
297 : "local r = getfenv(0).__ngx_req\n"
298 : "if r then\n"
299 : "local ctx = raw_ctx(r)\n"
300 : /* ignore header and body filters */
301 : "if ctx ~= 0x020 and ctx ~= 0x040 then\n"
302 : "return ours(...)\n"
303 : "end\n"
304 : "end\n"
305 : "return std(...)\n"
306 : "end\n"
307 : "end\n"
308 : "local create, resume = coroutine.create, coroutine.resume\n"
309 : "coroutine.wrap = function(f)\n"
310 : "local co = create(f)\n"
311 : "return function(...) return select(2, resume(co, ...)) end\n"
312 : "end\n"
313 : "package.loaded.coroutine = coroutine";
314 :
315 : #if 0
316 : "debug.sethook(function () collectgarbage() end, 'rl', 1)"
317 : #endif
318 : ;
319 :
320 18 : rc = luaL_loadbuffer(L, buf, sizeof(buf) - 1, "=coroutine.wrap");
321 : }
322 :
323 18 : if (rc != 0) {
324 0 : ngx_log_error(NGX_LOG_ERR, log, 0,
325 : "failed to load Lua code for coroutine.wrap(): %i: %s",
326 : rc, lua_tostring(L, -1));
327 :
328 0 : lua_pop(L, 1);
329 0 : return;
330 : }
331 :
332 18 : rc = lua_pcall(L, 0, 0, 0);
333 18 : if (rc != 0) {
334 0 : ngx_log_error(NGX_LOG_ERR, log, 0,
335 : "failed to run the Lua code for coroutine.wrap(): %i: %s",
336 : rc, lua_tostring(L, -1));
337 0 : lua_pop(L, 1);
338 : }
339 : }
340 :
341 :
342 : static int
343 0 : ngx_http_lua_coroutine_status(lua_State *L)
344 : {
345 : lua_State *co; /* new coroutine to be created */
346 : ngx_http_request_t *r;
347 : ngx_http_lua_ctx_t *ctx;
348 : ngx_http_lua_co_ctx_t *coctx; /* co ctx for the new coroutine */
349 :
350 0 : co = lua_tothread(L, 1);
351 :
352 0 : luaL_argcheck(L, co, 1, "coroutine expected");
353 :
354 0 : r = ngx_http_lua_get_req(L);
355 0 : if (r == NULL) {
356 0 : return luaL_error(L, "no request found");
357 : }
358 :
359 0 : ctx = ngx_http_get_module_ctx(r, ngx_http_lua_module);
360 0 : if (ctx == NULL) {
361 0 : return luaL_error(L, "no request ctx found");
362 : }
363 :
364 0 : ngx_http_lua_check_context(L, ctx, NGX_HTTP_LUA_CONTEXT_REWRITE
365 : | NGX_HTTP_LUA_CONTEXT_ACCESS
366 : | NGX_HTTP_LUA_CONTEXT_CONTENT
367 : | NGX_HTTP_LUA_CONTEXT_TIMER
368 : | NGX_HTTP_LUA_CONTEXT_SSL_CERT
369 : | NGX_HTTP_LUA_CONTEXT_SSL_SESS_FETCH);
370 :
371 0 : coctx = ngx_http_lua_get_co_ctx(co, ctx);
372 0 : if (coctx == NULL) {
373 0 : lua_pushlstring(L, (const char *)
374 0 : ngx_http_lua_co_status_names[NGX_HTTP_LUA_CO_DEAD].data,
375 : ngx_http_lua_co_status_names[NGX_HTTP_LUA_CO_DEAD].len);
376 0 : return 1;
377 : }
378 :
379 : dd("co status: %d", coctx->co_status);
380 :
381 0 : lua_pushlstring(L, (const char *)
382 0 : ngx_http_lua_co_status_names[coctx->co_status].data,
383 0 : ngx_http_lua_co_status_names[coctx->co_status].len);
384 0 : return 1;
385 : }
386 :
387 : /* vi:set ft=c ts=4 sw=4 et fdm=marker: */
|