协程学习：协程的实现（二）

本文以云风的协程框架来讲述协程的实现，学习协程的实现有利于了解协程的工作机制。
云风协程的实现框架： https://github.com/cloudwu/coroutine/

准备知识： ucontext机制 [1]
ucontext机制是GNU C库提供的一组用于创建、保存、切换用户态“上下文”(context)的API。
首先要了解的是结构体ucontext_t，这个结构体的作用是用来保存上下文的：

typedef struct ucontext {
  struct ucontext *uc_link;
    sigset_t      uc_sigmask;
    stack_t        uc_stack;
    mcontext_t uc_mcontext;
    ...
} ucontext_t;

还有4个相应的操作函数：

* int makecontext(ucontext_t *ucp, void (*func)(), int argc, ...) 该函数用以初始化一个ucontext_t类型的结构，也就是我们所说的用户执行上下文。函数指针func指明了该context的入口函数，argc指明入口参数个数，该值是可变的，但每个参数类型都是int型，这些参数紧随argc传入。 另外，在调用makecontext之前，一般还需要显式的指明其初始栈信息（栈指针SP及栈大小）和运行时的信号屏蔽掩码（signal mask）。 同时也可以指定uc_link字段，这样在func函数返回后，就会切换到uc_link指向的context继续执行。可以通过getcontext函数来获取初始化的ucontext_t。

* int setcontext(const ucontext_t *ucp) 该函数用来将当前程序执行线索切换到参数ucp所指向的上下文状态，在执行正确的情况下，该函数直接切入到新的执行状态，不再会返回。比如我们用上面介绍的makecontext初始化了一个新的上下文，并将入口指向某函数entry()，那么setcontext成功后就会马上运行entry()函数。

* int getcontext(ucontext_t *ucp) 该函数用来将当前执行状态上下文保存到一个ucontext_t结构中，若后续调用setcontext或swapcontext恢复该状态，则程序会沿着getcontext调用点之后继续执行，看起来好像刚从getcontext函数返回一样。 这个操作的功能和setjmp所起的作用类似，都是保存执行状态以便后续恢复执行，但需要重点指出的是：getcontext函数的返回值仅能表示本次操作是否执行正确，而不能用来区分是直接从getcontext操作返回，还是由于setcontext/swapcontex恢复状态导致的返回，这点与setjmp是不一样的。

* int swapcontext(ucontext_t *oucp, ucontext_t *ucp) 理论上，有了上面的3个函数，就可以满足需要了，但由于getcontext不能区分返回状态，因此编写上下文切换的代码时就需要保存额外的信息来进行判断，显得比较麻烦。 为了简化切换操作的实现，ucontext 机制里提供了swapcontext这个函数，用来“原子”地完成旧状态的保存和切换到新状态的工作

#include <stdio.h>
#include <ucontext.h>
#include <unistd.h>

int main(int argc, char *argv[]) {
  ucontext_t context;

  getcontext(&context);
  puts("Hello world");
  sleep(1);
  setcontext(&context);
  return 0;
}

上面程序的运行结果是大概每隔一秒中就会输出一次Hello World。程序通过getcontext获取当前的上下文，然后在通过setcontext切换到之前的上下文（类似于goto的作用）。

云风协程框架的实现分析[2]

整体思路：通过ucontext机制来进行上下文的切换，框架定义一个协程调度器schedule，用于管理协程。在schedule中定义了一个char stack[STACK_SIZE]用于保存运行时的协程的栈信息，在协程调度的过程中，每个协程都是从main函数切换过去的（系统是先切换到main函数，然后再从main函数切换到协程，协程返回的时候也切换到main函数，不能直接切换到其他协程），并且在协程切换的过程中，需要进行协程栈信息的拷贝，resume协程，需要把协程的栈信息拷贝到调度器的stack数组中，yield的时候需要把栈信息从schedule的stack拷贝到协程的stack。协程的stack是一个指针，用于指向协程的栈信息的地址，这个空间是动态分配的，不同的协程或者是同一个协程在不同的运行时间大小有可能不一样。

在继续阅读下面的内容之前，强烈建议读者先认真阅读下云风的协程框架的实现源码： https://github.com/cloudwu/coroutine/

数据结构：

struct schedule {               // 协程调度器，用于管理协程
    char stack[STACK_SIZE];     // 用于保存当前的栈信息
    ucontext_t main;            // 用于保存主函数的上下文
    int nco;                    // 当前协程的数量
    int cap;                    // 当前schedule容器协程的容量（最大的协程数量，系统会动态扩展，每次*2）
    int running;                // 当前正在运行的协程的id
    struct coroutine **co;      // 二位数组，保存当前系统的协程对应的结构体
};

struct coroutine {              // 协程结构体
    coroutine_func func;        // 协程运行的函数，相当于多线程中指定的工作函数
    void *ud;                   // 工作函数的参数
    ucontext_t ctx;             // 协程的上下文
    struct schedule * sch;      // 对应的schedule
    ptrdiff_t cap;              // 当前栈的容量
    ptrdiff_t size;             // 栈的使用空间大小
    int status;                 // 协程的状态，宏定义COROUTINE_DEAD，COROUTINE_READY，COROUTINE_RUNNING，COROUTINE_SUSPEND
    char *stack;                // 栈指针，栈是通过malloc动态分配的，用于在协程休眠的时候保存协程的栈信息，就是说需要自己去管理每个协程的上下文
};

创建新的协程 coroutine_new

int coroutine_new(struct schedule *S, coroutine_func func, void *ud) {
    struct coroutine *co = _co_new(S, func , ud);
    if (S->nco >= S->cap) {  // 当协程调度器的空间不够的时候进行扩展
        int id = S->cap;
        S->co = realloc(S->co, S->cap * 2 * sizeof(struct coroutine *));  // 扩展为原来的2倍
        memset(S->co + S->cap , 0 , sizeof(struct coroutine *) * S->cap);
        S->co[S->cap] = co;
        S->cap *= 2;     // 扩展为原来的两倍
        ++S->nco;
        return id;
    } else {
        int i;
        for (i=0;i<S->cap;i++) {   // 找到当前可用的id，这里通过for循环一个个查找，当协程数量比较多的时候效率相对比较低
            int id = (i+S->nco) % S->cap;
            if (S->co[id] == NULL) {
                S->co[id] = co;
                ++S->nco;
                return id;
            }
        }
    }

    assert(0);
    return -1;
}

恢复协程的运行状态： coroutine_resume

void coroutine_resume(struct schedule * S, int id) {
    assert(S->running == -1);
    assert(id >=0 && id < S->cap);
    struct coroutine *C = S->co[id];
    if (C == NULL)
        return;

    int status = C->status;
    switch(status) {
    case COROUTINE_READY:        // 新建的协程处于READY状态，并没有开始运行
        getcontext(&C->ctx);     // 获取上下文，这里主要的作用是通过getcontext获得一个初始化话的ucontext_t结构体，下面会继续对其进行初始化
        C->ctx.uc_stack.ss_sp = S->stack;      // 指定协程运行时的栈（所有的栈信息都会保存在这里）
        C->ctx.uc_stack.ss_size = STACK_SIZE;  // 栈的大小
        C->ctx.uc_link = &S->main;             // 协程退出时返回的上下文，一般协程退出时需要返回到main函数上次执行的地方
        S->running = id;                       // 记录当前运行的协程id
        C->status = COROUTINE_RUNNING;         // 标识为运行状态
        uintptr_t ptr = (uintptr_t)S;          
        makecontext(&C->ctx, (void (*)(void)) mainfunc, 2, (uint32_t)ptr, (uint32_t)(ptr>>32));  // 初始化ucontext结构，并且指定对应的运行函数mainfunc，mainfunc会调用协程对应的工作函数
        swapcontext(&S->main, &C->ctx);        // 切换到协程C的上下文，并且把旧的状态的上下文保存到S->main中（用于返回main函数）
        break;
    case COROUTINE_SUSPEND:
        memcpy(S->stack + STACK_SIZE - C->size, C->stack, C->size); // 将协程C的上下文（栈信息）拷贝到S->stack，因为C->ctx.uc_stack.ss_sp = S->stack，所以要进行恢复
        S->running = id;
        C->status = COROUTINE_RUNNING;
        swapcontext(&S->main, &C->ctx);        // 切换到协程C的上下文，并且把旧的状态的上下文保存到S->main中（用于返回main函数）
        break;
    default:
        assert(0);
    }
}

挂起协程： coroutine_yield

void coroutine_yield(struct schedule * S) {
    int id = S->running;
    assert(id >= 0);
    struct coroutine * C = S->co[id];
    assert((char *)&C > S->stack);
    _save_stack(C,S->stack + STACK_SIZE);   // 保存协程C的上下文 
    C->status = COROUTINE_SUSPEND;          // 把协程的状态标志为挂起
    S->running = -1;
    swapcontext(&C->ctx , &S->main);        // 切换会main函数
}

保存协程栈的函数： _save_stack

static void _save_stack(struct coroutine *C, char *top) {
    char dummy = 0;                         // 定义一个临时变量，用于获取当前栈的顶端的地址，要知道的是栈是往下增长的
    assert(top - &dummy <= STACK_SIZE);     // 获取当前栈所用的空间大小
    if (C->cap < top - &dummy) {
        free(C->stack);
        C->cap = top-&dummy;
        C->stack = malloc(C->cap);          // 申请空间用于保存协程C自己的栈信息，因为要切换到其他协程，所以每个协程在挂起的时候都要各自保存自己的栈信息
    }

    C->size = top - &dummy;
    memcpy(C->stack, &dummy, C->size);      // 保存协程C的栈信息，栈信息原本保存在S->statc中，并且从&dummy地址到S->stack + STACK_SIZE地址是当前栈的情况，从这里足以看出云风大哥深厚的编程功力
}

_save_stack感觉是整个框架中最有意思的一个函数了，通过程序运行栈从上往下增长的原理，从而获取到当前协程的栈信息。不过这个实现方式有个缺点就是，在切换协程的时候需要进行栈信息的拷贝，在yield的时候需要把栈信息从S->stack拷贝到协程C->stack中，在resume的时候也需要把协程的栈信息C->stack拷贝到S->statck中，这样会降低切换的效率，不过这个基本可以忽略。要解决这个问题可以在coroutine中也定义一个char stack[STACK_SIZE]，然后把协程的运行栈指向这个地址，这样resume和yield的时候就不用进行栈信息拷贝了，但这个办法同样会引入另外一个问题，就是如果协程很多，那么会占用大量的内存，而且每个协程需要的栈大小可能不一样，如果STACK_SIZE定义过大就会造成空间的浪费。所以两种方案都是有优缺点的，要根据实际情况进行取舍了。

参考资料：
[1] ucontext-人人都可以实现的简单协程库
[2] C的coroutine库

协程学习：协程的实现（二）

猜你喜欢