协程的简单实现

308 阅读2分钟

看极客时间专栏中的一个协程的简单的实现,将自己的理解记录下。原文如下 传送门

实现

#include <stdio.h>
#include <stdlib.h>

#define STACK_SIZE 1024

typedef void(*coro_start)();

class coroutine {
public:
    long* stack_pointer;
    char* stack;

    coroutine(coro_start entry) {
        if (entry == NULL) {
            stack = NULL;
            stack_pointer = NULL;
            return;
        }

        stack = (char*)malloc(STACK_SIZE);
        char* base = stack + STACK_SIZE;
        stack_pointer = (long*) base;
        stack_pointer -= 1;
        *stack_pointer = (long) entry;
        stack_pointer -= 1;
        *stack_pointer = (long) base;
    }

    ~coroutine() {
        if (!stack)
            return;
        free(stack);
        stack = NULL;
    }
};

coroutine* co_a, * co_b;

void yield_to(coroutine* old_co, coroutine* co) {
    __asm__ (
        "movq %%rsp, %0\n\t"
        "movq %%rax, %%rsp\n\t"
        :"=m"(old_co->stack_pointer):"a"(co->stack_pointer):);
}

void start_b() {
    printf("B");
    yield_to(co_b, co_a);
    printf("D");
    yield_to(co_b, co_a);
}

int main() {
    printf("A");
    co_b = new coroutine(start_b);
    co_a = new coroutine(NULL);
    yield_to(co_a, co_b);
    printf("C");
    yield_to(co_a, co_b);
    printf("E\n");
    delete co_a;
    delete co_b;
    return 0;
}

以上代码在编译的时候需要不使用编译时优化,否则将会内联yield_to方法,上面实现协程栈的切换其实巧妙的利用了在函数结束的治疗pop %rbp,retq,如果内联了就不生效了,因此需要使用如下命令进行编译:

g++ -g -o co -O0 coroutine.cpp

而执行程序可以看见打印顺序是A,B,C,D,E

理解

协程栈的建立

为了实现协程的切换,首先代码中定义了函数指针coro_start来代表切换的函数,而类coroutine在构造函数中封装了这一个指针,如果传递的不为null那么此时就会在堆中创建一个内容区域用于协程执行的栈,由于栈地址是从高到低的,因此stack_pointer成员变量需要经过简单的计算stack + STACK_SIZE来得到,创建好之后,这个栈很干净,之后进行了如下的操作。

  1. 将协程执行函数地址push到了栈中
  2. 将栈底指针值push到了栈中

以上两步是对协程栈做了初始化操作,具体如何利用需要查看yield_to方法。

协程栈的切换

查看main方法,协程的切换是通过yield_to方法来实现的,yield_to方法很简单,通过内联汇编的写法主要完成了以下操作

  1. 将老的(让出CPU)的协程的栈底指针保存到stack_pointer中。
  2. 将新的协程的栈指针保存到当前rsp寄存器中 这个时候函数执行结束,函数执行结束包含了两个退出栈的指令
  3. pop %rbp
  4. retq 查看当前内存布局

image.png

发现之前构造的栈刚好完了了rbp的赋值,retq也完成了新的rip指针的切换。

这样就完成了协程的切换。

再看协程的三个特点

  1. 切换和调度都发生在用户态。
  2. 占用的资源更少,开销更少
  3. 它的调度是协商式的,而不是抢占式的