自己动手实现一个简单的协程模型

协程又叫用户级轻量线程，它不需要像线程那样占用大量系统资源，但却能像线程那样并发地运行多个函数，它是怎样实现的呢？让我们先搞清楚它的实现细节，然后再动手自己做一个。

在CPU中有个IP寄存器，它的值决定了下一条将要执行的指令地址，出于安全起见，通常我们不能手动设置它的值，但是有一些指令，如CALL和RET能够间接的改变它。比如CALL指令将函数地址赋值给IP寄存器，使CPU接下来跳到函数中执行，在函数执行完成之后，RET指令从堆栈中弹出CALL之后的那条指令，并且赋值给IP寄存器，那么函数返回后就能按原来的流程继续执行。

还有两个比较重要的寄存器：BP和SP，他们控制着当前函数的局部变量，以及所有的子函数调用。BP指向当前函数的栈顶，在函数运行期间是固定的，SP指向函数的栈底，它们之间的内存区域标识当前函数的执行状态。如下图所示：

image

如果我们能想到办法，通过某种机制保存函数这时候的状态，是不是可以在多个函数之间切换呢？在c语言库函数中，真有这样的函数：setjmp()和longjmp()，setjmp()负责保存当前函数的执行状态，longjmp()可以在其它函数中随时返回之前的现场，是不是很神奇。看下面的例子：

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

jmp_buf env;

void fun()
{
    printf("fun() enter \n");
    longjmp(env, 1);
    printf("fun() exit \n");
}

int main()
{
    int ret = setjmp(env);

    if (ret == 0)
    {
        printf("from setjmp \n");
        fun();
    }
    else
    {
        printf("from longjmp \n");
    }

    return 0;
}

编译：

g++ -o setjmp setjmp.c

执行结果如下：

[kdjie@localhost test.d]$ ./setjmp
from setjmp
fun() enter
from longjmp

上面是一个函数的例子，如果有很多函数，我们为每个函数都调用setjmp()保存一个现场，是不是可以随意在多个函数之间切换呢？答案是否定的，因为使用setjmp()和longjmp()有个约束，只能跳一次，而且只能从栈的深处向浅处跳！为什么呢？因为正常的进程栈栈，无论有多少个函数，全部共用一块连续的内存区域，一旦发生跳转，跳转之间的区域全部丢失，这将导致堆栈被破坏，再也跳不回来了！

有朋友会提出问题：可以为每个函数独立设置堆栈吗？答案是可以！但是c函数库的setjmp()和longjmp()不行，我们必须自己动手写。

首先自己实现这两个函数，并保存为setjmp.s：

.section .text

.globl setjmp, longjmp

setjmp:
    movq %rbp, 0(%rdi)
    movq %rsp, 8(%rdi)
    movq 0(%rsp), %rdx
    movq %rdx, 16(%rdi)
    movq $0, %rax
    retq

longjmp:
    movq 0(%rdi), %rbp
    movq 8(%rdi), %rsp
    movq 16(%rdi), %rdx
    movq %rdx, 0(%rsp)
    movq $1, %rax
    retq

在setjmp中，只保存了BP和SP这两个寄存器，以及从堆栈中取出函数的返回地址（即原本应该执行的下一条语句！），并且设置返回值为0，加以区别。

在longjmp中，将入参的BP、SP、函数返回地址加以恢复，实际上是通过破坏堆栈的方式欺骗了RET指令！让RET误以为函数正常返回，实际上却切换到了别的函数中。

然后再实现一个简单协程的例子，保存为jmpcall.c：

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

// 定义上下文内存区域，用于保存BP、SP、返回地址
// 在64位环境，寄存器名称为rbp、rsp、rip
struct context {
    unsigned long rbp;
    unsigned long rsp;
    unsigned long rip;
};

// 定义上下文数据，用于模拟多个协程
struct context ctEntry = {0};
struct context ctArray[2] = { {0}, {0} };
int ct_count = 2;
nt ct_idx = -1;

// 保存当前函数现场，并切换到其它函数
void yield()
{
    if (setjmp( &ctArray[ct_idx] ) == 0)
    {
        longjmp(&ctEntry);
    }
}

// 协程函数1
void fun1()
{
    int number = 0;
    
    for (;;)
    {
        printf("fun1() %d \n", number++);
        sleep(1);

        yield();
    } 
}

// 协程函数2
void fun2()
{
    int number = 0;
    for (;;)
    {
        printf("fun2() %d \n", number++);
        sleep(1);

        // 模拟协程退出，不再参与调度
        if (number == 3)
        {
            ct_count = 1;
        }

        yield();
    }
}

int main()
{
    // 创建协程1
    // 注意：这里创建了函数自己的堆栈区域
    ctArray[0].rbp = (unsigned long)((char*)malloc(4096) + 4000);
    ctArray[0].rsp = ctArray[0].rbp;
    ctArray[0].rip = (unsigned long)&fun1;
    *( (unsigned long*)(ctArray[0].rbp) ) = 0LL;
    
    // 创建协程2
    // 注意：这里创建了函数自己的堆栈区域
    ctArray[1].rbp = (unsigned long)((char*)malloc(4096) + 4000);
    ctArray[1].rsp = ctArray[1].rbp;
    ctArray[1].rip = (unsigned long)&fun2;
    *( (unsigned long*)(ctArray[1].rbp) ) = 0LL;
    
    // 保存调度现场
    int ret = setjmp(&ctEntry);
    if (ret == 1)
    {
        // 调度协程
        ct_idx = (ct_idx + 1) % ct_count;
        longjmp( &ctArray[ct_idx] );
    }

    // 启动调度
    longjmp(&ctEntry);
    return 0;
}

协程模型流程图示：

image

编译：

as --64 -gstabs -o setjmp.o setjmp.s
gcc -ggdb -c -o jmpcall.o jmpcall.c
gcc -ggdb -o jmpcall setjmp.o jmpcall.o

执行结果如下：

[kdjie@localhost test.d]$ ./jmpcall
fun1() 0
fun2() 0
fun1() 1
fun2() 1
fun1() 2
fun2() 2
fun1() 3
fun1() 4
fun1() 5
fun1() 6
……

OK，实现成功！