path: root/ngx_stream_lua-0.0.16/src/ngx_stream_lua_coroutine.c

/*
 * !!! DO NOT EDIT DIRECTLY !!!
 * This file was automatically generated from the following template:
 *
 * src/subsys/ngx_subsys_lua_coroutine.c.tt2
 */


/*
 * Copyright (C) Xiaozhe Wang (chaoslawful)
 * Copyright (C) Yichun Zhang (agentzh)
 */


#ifndef DDEBUG
#define DDEBUG 0
#endif
#include "ddebug.h"


#include "ngx_stream_lua_coroutine.h"
#include "ngx_stream_lua_util.h"
#include "ngx_stream_lua_probe.h"


/*
 * Design:
 *
 * In order to support using ngx.* API in Lua coroutines, we have to create
 * new coroutine in the main coroutine instead of the calling coroutine
 */


static int ngx_stream_lua_coroutine_create(lua_State *L);
static int ngx_stream_lua_coroutine_wrap(lua_State *L);
static int ngx_stream_lua_coroutine_resume(lua_State *L);
static int ngx_stream_lua_coroutine_yield(lua_State *L);
static int ngx_stream_lua_coroutine_status(lua_State *L);


static const ngx_str_t
    ngx_stream_lua_co_status_names[] =
    {
        ngx_string("running"),
        ngx_string("suspended"),
        ngx_string("normal"),
        ngx_string("dead"),
        ngx_string("zombie")
    };



static int
ngx_stream_lua_coroutine_create(lua_State *L)
{
    ngx_stream_lua_request_t        *r;
    ngx_stream_lua_ctx_t            *ctx;

    r = ngx_stream_lua_get_req(L);
    if (r == NULL) {
        return luaL_error(L, "no request found");
    }

    ctx = ngx_stream_lua_get_module_ctx(r, ngx_stream_lua_module);
    if (ctx == NULL) {
        return luaL_error(L, "no request ctx found");
    }

    return ngx_stream_lua_coroutine_create_helper(L, r, ctx, NULL);
}


static int
ngx_stream_lua_coroutine_wrap_runner(lua_State *L)
{
    /* retrieve closure and insert it at the bottom of
     * the stack for coroutine.resume() */
    lua_pushvalue(L, lua_upvalueindex(1));
    lua_insert(L, 1);

    return ngx_stream_lua_coroutine_resume(L);
}


static int
ngx_stream_lua_coroutine_wrap(lua_State *L)
{
    ngx_stream_lua_request_t            *r;
    ngx_stream_lua_ctx_t                *ctx;
    ngx_stream_lua_co_ctx_t             *coctx = NULL;

    r = ngx_stream_lua_get_req(L);
    if (r == NULL) {
        return luaL_error(L, "no request found");
    }

    ctx = ngx_stream_lua_get_module_ctx(r, ngx_stream_lua_module);
    if (ctx == NULL) {
        return luaL_error(L, "no request ctx found");
    }

    ngx_stream_lua_coroutine_create_helper(L, r, ctx, &coctx);

    coctx->is_wrap = 1;

    lua_pushcclosure(L, ngx_stream_lua_coroutine_wrap_runner, 1);

    return 1;
}


int
ngx_stream_lua_coroutine_create_helper(lua_State *L,
    ngx_stream_lua_request_t *r, ngx_stream_lua_ctx_t *ctx,
    ngx_stream_lua_co_ctx_t **pcoctx)
{
    lua_State                     *vm;  /* the Lua VM */
    lua_State                     *co;  /* new coroutine to be created */

    /* co ctx for the new coroutine */
    ngx_stream_lua_co_ctx_t               *coctx;

    luaL_argcheck(L, lua_isfunction(L, 1) && !lua_iscfunction(L, 1), 1,
                  "Lua function expected");

    ngx_stream_lua_check_context(L, ctx, NGX_STREAM_LUA_CONTEXT_YIELDABLE);

    vm = ngx_stream_lua_get_lua_vm(r, ctx);

    /* create new coroutine on root Lua state, so it always yields
     * to main Lua thread
     */
    co = lua_newthread(vm);

    ngx_stream_lua_probe_user_coroutine_create(r, L, co);

    coctx = ngx_stream_lua_get_co_ctx(co, ctx);
    if (coctx == NULL) {
        coctx = ngx_stream_lua_create_co_ctx(r, ctx);
        if (coctx == NULL) {
            return luaL_error(L, "no memory");
        }

    } else {
        ngx_memzero(coctx, sizeof(ngx_stream_lua_co_ctx_t));
        coctx->co_ref = LUA_NOREF;
    }

    coctx->co = co;
    coctx->co_status = NGX_STREAM_LUA_CO_SUSPENDED;

#ifdef OPENRESTY_LUAJIT
    ngx_stream_lua_set_req(co, r);
    ngx_stream_lua_attach_co_ctx_to_L(co, coctx);
#else
    /* make new coroutine share globals of the parent coroutine.
     * NOTE: globals don't have to be separated! */
    ngx_stream_lua_get_globals_table(L);
    lua_xmove(L, co, 1);
    ngx_stream_lua_set_globals_table(co);
#endif

    lua_xmove(vm, L, 1);    /* move coroutine from main thread to L */

    lua_pushvalue(L, 1);    /* copy entry function to top of L*/
    lua_xmove(L, co, 1);    /* move entry function from L to co */

    if (pcoctx) {
        *pcoctx = coctx;
    }

#ifdef NGX_LUA_USE_ASSERT
    coctx->co_top = 1;
#endif

    return 1;    /* return new coroutine to Lua */
}


static int
ngx_stream_lua_coroutine_resume(lua_State *L)
{
    lua_State                           *co;
    ngx_stream_lua_request_t            *r;
    ngx_stream_lua_ctx_t                *ctx;
    ngx_stream_lua_co_ctx_t             *coctx;
    ngx_stream_lua_co_ctx_t             *p_coctx; /* parent co ctx */

    co = lua_tothread(L, 1);

    luaL_argcheck(L, co, 1, "coroutine expected");

    r = ngx_stream_lua_get_req(L);
    if (r == NULL) {
        return luaL_error(L, "no request found");
    }

    ctx = ngx_stream_lua_get_module_ctx(r, ngx_stream_lua_module);
    if (ctx == NULL) {
        return luaL_error(L, "no request ctx found");
    }

    ngx_stream_lua_check_context(L, ctx, NGX_STREAM_LUA_CONTEXT_CONTENT
                                 | NGX_STREAM_LUA_CONTEXT_TIMER
                                 | NGX_STREAM_LUA_CONTEXT_SSL_CLIENT_HELLO
                                 | NGX_STREAM_LUA_CONTEXT_SSL_CERT
                                 | NGX_STREAM_LUA_CONTEXT_PREREAD
                                 );

    p_coctx = ctx->cur_co_ctx;
    if (p_coctx == NULL) {
        return luaL_error(L, "no parent co ctx found");
    }

    coctx = ngx_stream_lua_get_co_ctx(co, ctx);
    if (coctx == NULL) {
        return luaL_error(L, "no co ctx found");
    }

    ngx_stream_lua_probe_user_coroutine_resume(r, L, co);

    if (coctx->co_status != NGX_STREAM_LUA_CO_SUSPENDED) {
        dd("coroutine resume: %d", coctx->co_status);

        lua_pushboolean(L, 0);
        lua_pushfstring(L, "cannot resume %s coroutine",
                        ngx_stream_lua_co_status_names[coctx->co_status].data);
        return 2;
    }

    p_coctx->co_status = NGX_STREAM_LUA_CO_NORMAL;

    coctx->parent_co_ctx = p_coctx;

    dd("set coroutine to running");
    coctx->co_status = NGX_STREAM_LUA_CO_RUNNING;

    ctx->co_op = NGX_STREAM_LUA_USER_CORO_RESUME;
    ctx->cur_co_ctx = coctx;

    /* yield and pass args to main thread, and resume target coroutine from
     * there */
    return lua_yield(L, lua_gettop(L) - 1);
}


static int
ngx_stream_lua_coroutine_yield(lua_State *L)
{
    ngx_stream_lua_request_t            *r;
    ngx_stream_lua_ctx_t                *ctx;
    ngx_stream_lua_co_ctx_t             *coctx;

    r = ngx_stream_lua_get_req(L);
    if (r == NULL) {
        return luaL_error(L, "no request found");
    }

    ctx = ngx_stream_lua_get_module_ctx(r, ngx_stream_lua_module);
    if (ctx == NULL) {
        return luaL_error(L, "no request ctx found");
    }

    ngx_stream_lua_check_context(L, ctx, NGX_STREAM_LUA_CONTEXT_CONTENT
                                 | NGX_STREAM_LUA_CONTEXT_TIMER
                                 | NGX_STREAM_LUA_CONTEXT_SSL_CLIENT_HELLO
                                 | NGX_STREAM_LUA_CONTEXT_SSL_CERT
                                 | NGX_STREAM_LUA_CONTEXT_PREREAD
                                 );

    coctx = ctx->cur_co_ctx;

    coctx->co_status = NGX_STREAM_LUA_CO_SUSPENDED;

    ctx->co_op = NGX_STREAM_LUA_USER_CORO_YIELD;

    if (!coctx->is_uthread && coctx->parent_co_ctx) {
        dd("set coroutine to running");
        coctx->parent_co_ctx->co_status = NGX_STREAM_LUA_CO_RUNNING;

        ngx_stream_lua_probe_user_coroutine_yield(r,
                                                  coctx->parent_co_ctx->co, L);

    } else {
        ngx_stream_lua_probe_user_coroutine_yield(r, NULL, L);
    }

    /* yield and pass retvals to main thread,
     * and resume parent coroutine there */
    return lua_yield(L, lua_gettop(L));
}


void
ngx_stream_lua_inject_coroutine_api(ngx_log_t *log, lua_State *L)
{
    int         rc;

    /* new coroutine table */
    lua_createtable(L, 0 /* narr */, 16 /* nrec */);

    /* get old coroutine table */
    lua_getglobal(L, "coroutine");

    /* set running to the old one */
    lua_getfield(L, -1, "running");
    lua_setfield(L, -3, "running");

    lua_getfield(L, -1, "create");
    lua_setfield(L, -3, "_create");

    lua_getfield(L, -1, "wrap");
    lua_setfield(L, -3, "_wrap");

    lua_getfield(L, -1, "resume");
    lua_setfield(L, -3, "_resume");

    lua_getfield(L, -1, "yield");
    lua_setfield(L, -3, "_yield");

    lua_getfield(L, -1, "status");
    lua_setfield(L, -3, "_status");

    /* pop the old coroutine */
    lua_pop(L, 1);

    lua_pushcfunction(L, ngx_stream_lua_coroutine_create);
    lua_setfield(L, -2, "__create");

    lua_pushcfunction(L, ngx_stream_lua_coroutine_wrap);
    lua_setfield(L, -2, "__wrap");

    lua_pushcfunction(L, ngx_stream_lua_coroutine_resume);
    lua_setfield(L, -2, "__resume");

    lua_pushcfunction(L, ngx_stream_lua_coroutine_yield);
    lua_setfield(L, -2, "__yield");

    lua_pushcfunction(L, ngx_stream_lua_coroutine_status);
    lua_setfield(L, -2, "__status");

    lua_setglobal(L, "coroutine");

    /* inject coroutine APIs */
    {
        const char buf[] =
            "local keys = {'create', 'yield', 'resume', 'status', 'wrap'}\n"
#ifdef OPENRESTY_LUAJIT
            "local get_req = require 'thread.exdata'\n"
#else
            "local getfenv = getfenv\n"
#endif
            "for _, key in ipairs(keys) do\n"
               "local std = coroutine['_' .. key]\n"
               "local ours = coroutine['__' .. key]\n"
               "local raw_ctx = ngx._phase_ctx\n"
               "coroutine[key] = function (...)\n"
#ifdef OPENRESTY_LUAJIT
                    "local r = get_req()\n"
#else
                    "local r = getfenv(0).__ngx_req\n"
#endif
                    "if r ~= nil then\n"
#ifdef OPENRESTY_LUAJIT
                        "local ctx = raw_ctx()\n"
#else
                        "local ctx = raw_ctx(r)\n"
#endif
                        "return ours(...)\n"
                    "end\n"
                    "return std(...)\n"
                "end\n"
            "end\n"
            "package.loaded.coroutine = coroutine"
#if 0
            "debug.sethook(function () collectgarbage() end, 'rl', 1)"
#endif
            ;

        rc = luaL_loadbuffer(L, buf, sizeof(buf) - 1, "=coroutine_api");
    }

    if (rc != 0) {
        ngx_log_error(NGX_LOG_ERR, log, 0,
                      "failed to load Lua code for coroutine_api: %i: %s",
                      rc, lua_tostring(L, -1));

        lua_pop(L, 1);
        return;
    }

    rc = lua_pcall(L, 0, 0, 0);
    if (rc != 0) {
        ngx_log_error(NGX_LOG_ERR, log, 0,
                      "failed to run the Lua code for coroutine_api: %i: %s",
                      rc, lua_tostring(L, -1));
        lua_pop(L, 1);
    }
}


static int
ngx_stream_lua_coroutine_status(lua_State *L)
{
    lua_State                     *co;  /* new coroutine to be created */
    ngx_stream_lua_request_t      *r;
    ngx_stream_lua_ctx_t          *ctx;
    ngx_stream_lua_co_ctx_t       *coctx; /* co ctx for the new coroutine */

    co = lua_tothread(L, 1);

    luaL_argcheck(L, co, 1, "coroutine expected");

    r = ngx_stream_lua_get_req(L);
    if (r == NULL) {
        return luaL_error(L, "no request found");
    }

    ctx = ngx_stream_lua_get_module_ctx(r, ngx_stream_lua_module);
    if (ctx == NULL) {
        return luaL_error(L, "no request ctx found");
    }

    ngx_stream_lua_check_context(L, ctx, NGX_STREAM_LUA_CONTEXT_CONTENT
                                 | NGX_STREAM_LUA_CONTEXT_TIMER
                                 | NGX_STREAM_LUA_CONTEXT_SSL_CLIENT_HELLO
                                 | NGX_STREAM_LUA_CONTEXT_SSL_CERT
                                 | NGX_STREAM_LUA_CONTEXT_PREREAD
                                 );

    coctx = ngx_stream_lua_get_co_ctx(co, ctx);
    if (coctx == NULL) {
        lua_pushlstring(L, (const char *)
                        ngx_stream_lua_co_status_names[NGX_STREAM_LUA_CO_DEAD]
                        .data,
                        ngx_stream_lua_co_status_names[NGX_STREAM_LUA_CO_DEAD]
                        .len);
        return 1;
    }

    dd("co status: %d", coctx->co_status);

    lua_pushlstring(L, (const char *)
                    ngx_stream_lua_co_status_names[coctx->co_status].data,
                    ngx_stream_lua_co_status_names[coctx->co_status].len);
    return 1;
}

/* vi:set ft=c ts=4 sw=4 et fdm=marker: */