實現一個簡單的C++協程庫

twoon發表於2014-03-22
C++

之前看協程相關的東西時，曾一念而過想著怎麼自己來實現一個給 C++ 用，但在儲存現場恢復現場之類的細節上被自己的想法嚇住，也沒有深入去研究，後面一丟開就忘了。近來微博上看人在討論怎麼實現一個 user space 上的執行緒庫，有人提到了 setcontext，swapcontext 之類的函式，說可以用來儲存和切換上下文，我忽然覺得這應該也能用來實現協程，回頭一搜，果然已經有人曾用這些函式做過相關的事情，略略看了幾個，覺得到底不大好用，還不如自己搞一個簡單點的。

說到 c++ 上的協程，boost 裡其實已經有相關的實現了，不過介面上看用起來有些麻煩，單純從語法上來說，我覺得 Lua 的協程最簡潔易用了，概念上也比較直接，為什麼不做一個類似的呢？所以我就打算照著 Lua 來山寨一個，只需要支援四個介面就夠了：

1）create coroutine。

2）run/resume coroutine。

3）Yield running corouinte。

4）IsCoroutineAlive。 

儲存與恢復上下文

實現協程/執行緒，最麻煩莫過於儲存和切換上下文了，好在 makecontext，swapcontext 這幾個函式相當好用，已經完全幫忙解決了這個難題：makecontext 可以幫我們建立起協程的上下文，swapcontext 則可以切換不同的上下文，從而實現那種把當前函式暫時停住，切換出去執行別的函式然後再切換回來繼續執行的效果：

#include <iostream>
#include <ucontext.h>
using namespace std;

static char g_stack[2048];
static ucontext_t ctx,ctx_main;

void func()
{
    // do something.
    cout << "enter func" << endl;

    swapcontext(&ctx, &ctx_main);

    cout << "func1 resume from yield" << endl;
    // continue to do something.
}

int main()
{
   getcontext(&ctx);
   ctx.uc_stack.ss_sp = g_stack;
   ctx.uc_stack.ss_size = sizeof g_stack;
   ctx.uc_link = &ctx_main;
    
   makecontext(&ctx, func, 0);

   cout << "in main, before coroutine starts" << endl;

   swapcontext(&ctx_main, &ctx);

   cout << "back to main" << endl;

   swapcontext(&ctx_main, &ctx);
   
   cout << "back to main again" << endl;
   return 0;
}

如上程式碼所示，顯然我們只要簡單包裝一下 swapcontext，很容易就可以實現 Yield 和 Resume，有了它們的幫助協程做起來就容易多了。

使用與實現

在使用 makecontext，swapcontext 的基礎上，我花了一個多小時簡單實現了一個協程庫，參看這裡，程式碼寫下來總共才200多行，出乎意料的簡單，用起來也很方便了：

#include "coroutine.h"

#include <iostream>

using namespace std;

CoroutineScheduler* sched = NULL;

void func1(void* arg)
{
    uintptr_t ret;
    cout << "function1 a now!,arg:" << arg << ", start to yield." << endl;
    ret = sched->Yield((uintptr_t)"func1 yield 1");
    cout << "1.fun1 return from yield:" << (const char*)ret << endl;
    ret = sched->Yield((uintptr_t)"func1 yield 2");
    cout << "2.fun1 return from yield:" << (const char*)ret << ", going to stop" << endl;

}

void func2(void* s)
{
    cout << "function2 a now!, arg:" << s << ", start to yield." << endl;
    const char* y = (const char*)sched->Yield((uintptr_t)"func2 yield 1");
    cout << "fun2 return from yield:" << y <<", going to stop" << endl;
}

int main()
{
    sched = new CoroutineScheduler();

    bool stop = false;
    int f1 = sched->CreateCoroutine(func1, (void*)111);
    int f2 = sched->CreateCoroutine(func2, (void*)222);

    while (!stop)
    {
        stop = true;
        if (sched->IsCoroutineAlive(f1))
        {
            stop = false;
            const char* y1 = (const char*)sched->ResumeCoroutine(f1, (uintptr_t)"resume func1");
            cout << "func1 yield:" << y1 << endl;
        }

        if (sched->IsCoroutineAlive(f2))
        {
            stop = false;
            const char* y2 = (const char*)sched->ResumeCoroutine(f2, (uintptr_t)"resume func2");
            cout << "func2 yield:" << y2 << endl;
        }
    }

    delete sched;
    return 0;
}

如上所示，Yield 裡傳的引數會在呼叫 Resume 時被返回，同理 Resume 裡的第二個引數，會在 Yield 裡被返回，這種機制也是模仿 Lua 來的，有些時候可以用來在協程間傳遞一些引數，很方便，看起來也挺酷的，但在實現上卻相當地簡潔，核心程式碼如下：

// static function
void CoroutineScheduler::SchedulerImpl::Schedule(void* arg)
{
    assert(arg);
    SchedulerImpl* sched = (SchedulerImpl*) arg;

    int running = sched->running_;

    coroutine* cor = sched->id2routine_[running];
    assert(cor);

    cor->func(cor->arg);

    sched->running_ = -1;
    cor->status = CO_FINISHED;
}

// resume coroutine.
uintptr_t CoroutineScheduler::SchedulerImpl::ResumeCoroutine(int id, uintptr_t y)
{
    coroutine* cor = id2routine_[id];
    if (cor == NULL || cor->status == CO_RUNNING) return 0;

    cor->yield = y;
    switch (cor->status)
    {
        case CO_READY:
            {
                getcontext(&cor->cxt);

                cor->status = CO_RUNNING;
                cor->cxt.uc_stack.ss_sp = cor->stack;
                cor->cxt.uc_stack.ss_size = stacksize_;
                // sucessor context.
                cor->cxt.uc_link = &mainContext_;

                running_ = id;
                makecontext(&cor->cxt, (void (*)())Schedule, 1, this);
                swapcontext(&mainContext_, &cor->cxt);
            }
            break;
        case CO_SUSPENDED:
            {
                running_ = id;
                cor->status = CO_RUNNING;
                swapcontext(&mainContext_, &cor->cxt);
            }
            break;
        default:
            assert(0);
    }

    uintptr_t ret = cor->yield;

    if (running_ == -1 && cor->status == CO_FINISHED) DestroyCoroutine(id);

    return ret;
}

uintptr_t CoroutineScheduler::SchedulerImpl::Yield(uintptr_t y)
{
    if (running_ < 0) return 0;

    int cur = running_;
    running_ = -1;

    coroutine* cor = id2routine_[cur];

    cor->yield = y;
    cor->status = CO_SUSPENDED;

    swapcontext(&cor->cxt, &mainContext_);
    return cor->yield;
}

單就程式碼量和程式結構而言，以上的實現很簡潔，但細節上看，每個協程都要分配一個一定大小的棧空間，空間效率上可能不大好，不夠輕量；執行效率上來說，swapcontext 的執行效率如何，現在也未知，只是出於學習的目的，就先這樣吧，可以再瞭解瞭解別人是怎麼做的。 

相關文章