Coroutine Nonsense

coexist.lua

local co_pool = {};

local function co_main(fn, ...)
    -- Main loop for a thread pool coroutine. Each resumption of the coroutine
    -- as acquired from the pool should pass in a function and arguments to be
    -- invoked.
    --
    -- The function is expected to return either true or false as its first
    -- return value; if true then the thread will stay alive, otherwise it
    -- will die. A dead thread should not be released into the pool.

    local function return_or_yield(ok, ...)
        assert(type(ok) == "boolean", "task function must return a boolean");

        if ok then
            return co_main(coroutine.yield(...));
        else
            return ...;
        end
    end

    return return_or_yield(fn(...));
end

local function co_acquire()
    local thread = table.remove(co_pool);

    if not thread then
        thread = coroutine.create(co_main);
    end

    return thread;
end

local function co_release(thread)
    assert(coroutine.status(thread) ~= "dead", "attempted to release a dead coroutine");

    if #co_pool < 10 then
        table.insert(co_pool, thread);
    end
end

local function co_launch(fn, ...)
    -- The supplied function will be run on a thread acquired from the thread
    -- pool. If the function returns without yielding the coroutine, the
    -- results are passed back to the caller and the thread released back
    -- into the pool.
    --
    -- If the given function yields during execution, we'll set a flag to
    -- indicate that this has occurred and when it finally finishes executing,
    -- we won't release it into the pool and will also tell the "co_main"
    -- function to wrap up and kill the thread.
    --
    -- This allows a function to essentially fork a thread from the pool
    -- and take ownership of it by yielding without running any risks of it
    -- being re-entered into the pool or resumed unexpectedly.

    local thread = co_acquire();
    local done = false;
    local suspended = false;

    local function finish_and_return(...)
        done = true;
        return not suspended, ...;
    end

    local function invoke_and_finish(...)
        return finish_and_return(fn(...));
    end

    local function process_result(ok, ...)
        if not ok then
            return false, (...);
        elseif done then
            -- finish_and_return was invoked prior to unwinding
            if not suspended then
                co_release(thread);
            end

            return "done", ...;
        else
            suspended = true;
            return "suspended", thread, ...;
        end
    end

    return process_result(coroutine.resume(thread, invoke_and_finish, ...));
end

local function co_check(errfn, ok, ...)
    if not ok then
        return false, errfn((...));
    end

    return ok, ...;
end

local function co_raise(err)
    error(err, 4);
end

function coroutine.exec(fn, ...)
    return co_check(co_raise, co_launch(fn, ...));
end

function coroutine.pcall(fn, ...)
    return co_launch(fn, ...);
end

function coroutine.xpcall(fn, errfn, ...)
    return co_check(errfn, co_launch(fn, ...));
end

---

coroutine.exec(function() print(1); end);
coroutine.exec(function(...) print(...); end, 1, 2, 3);
print(coroutine.exec(function(a, b) return a + b, a - b; end, 10, 5));
print(pcall(coroutine.exec, function() error("exec pcall"); end));
print(coroutine.pcall(error, "pcall"));
coroutine.xpcall(error, print, "xpcall");
print(coroutine.xpcall(error, function(a) return "changed", a; end, "xpcall ret"));

do
    local status, inner, message = coroutine.exec(function()
        print("Starting");
        print("Finishing", coroutine.yield("inner yield"));
    end);

    assert(status == "suspended", type(inner) == "thread");

    coroutine.exec(function() print("Second thread"); end);

    print(coroutine.resume(inner, message .. " with more text"));
    assert(coroutine.status(inner) == "dead");
end