简易协程库

协程

本项目Github地址

简介

使用C++和x86汇编设计并实现了一个简易的有栈协程库。通过保存和恢复Callee Save寄存器的方式实现协程的切换。通过admin协程进行协程调度，每次协程调度都会先切换回admin，由admin选择下一个上台运行的协程。使用显示运行时打桩机制Hook常用的库函数例如read和accept函数等，改为使用非阻塞的方式调用，从而避免线程被挂起。

从x86-64汇编开始

Callee Save和Caller Save

在x86-64体系结构中，寄存器可以分为两种：Callee Save和Caller Save。其中Callee Save的寄存器需要被调用者保存，也就是说如果调用者需要使用到这些寄存器，就需要先将这些寄存器的值保存（一般保存在栈中），函数返回之前再把这些寄存器的值还原。Caller Save的寄存器需要调用者主动保存，也就是说在调用子函数之前要先将这些寄存器的值保存，而被调用者可以随意覆盖这些寄存器，子函数返回之后，父函数将这些寄存器恢复。

图1

Callee Save的通用寄存器有%rbx，%rbp，%r12，%r13，%r14，%r15。

Caller Save的通用寄存器有%rax，%rcx，%rdx，%rsp，%rsi，%r8，%r9，%r10，%r11。

子函数调用与函数栈帧

看下面一个C程序：

//add.c
int add(int a, int b){
    int c;
    c = a + b;
    return c;
}

int main(){
    int a, b, c;
    a = 1;
    b = 1;
    c = add(a, b);
    return 0;
}

使用gcc编译成汇编语言的形式：

gcc -S add.c

	.file	"add.c"
	.text
	.globl	add
	.type	add, @function
add:
.LFB0:
	pushq	%rbp
	movq	%rsp, %rbp
	movl	%edi, -20(%rbp)
	movl	%esi, -24(%rbp)
	movl	-24(%rbp), %eax
	movl	-20(%rbp), %edx
	addl	%edx, %eax
	movl	%eax, -4(%rbp)
	movl	-4(%rbp), %eax
	popq	%rbp
	ret
.LFE0:
	.size	add, .-add
	.globl	main
	.type	main, @function
main:
.LFB1:
	pushq	%rbp
	movq	%rsp, %rbp
	subq	$16, %rsp
	movl	$1, -4(%rbp)
	movl	$1, -8(%rbp)
	movl	-8(%rbp), %edx
	movl	-4(%rbp), %eax
	movl	%edx, %esi
	movl	%eax, %edi
	call	add
	movl	%eax, -12(%rbp)
	movl	$0, %eax
	leave
	ret
.LFE1:
	.size	main, .-main
	.ident	"GCC: (GNU) 4.8.5 20150623 (Red Hat 4.8.5-36)"
	.section	.note.GNU-stack,"",@progbits

从上述汇编代码中可以看出，调用子函数使用的是call指令，子函数返回值保存在%rax寄存器中。所以在call指令后有一条指令：

movl	%eax, -12(%rbp)

将add函数的返回值（保存在%rax寄存器中）复制到-12(%rbp)中（这是局部变量c在函数栈帧中的地址）。

x86-64函数栈帧结构如下图所示：

图2

call指令调用子函数，等价于：

pushl %rip  //将返回地址（call指令的下一条地址）压栈
jmp Foo     //跳转指令

leave指令等价于：

movq %rbp, %rsp     //
popq %rbp           //恢复父函数栈帧

ret指令等价于：

popl %rip           //将%rip寄存器的值设为父函数返回地址

协程栈帧的保存与恢复

之前我们说到x86-64的通用寄存器分为Callee Save和Caller Save两种。假设我们需要调用一个函数来进行协程之间的切换，那么我们就需要保存当前协程的寄存器以及堆栈状态。因为寄存器中的%rbp和%rsp分别为栈帧的基地址和栈顶地址，所以实际上在保存寄存器状态的同时也保存了堆栈的状态。

我们需要保存的寄存器是Callee Save即被调用者需要保存的寄存器，因为Caller Save的寄存器编译器已经帮我们保存了。

协程栈帧保存函数save_context:

/*rdi保存传入的第一个参数：from线程上下文buffer基地址*/
save_context:
    pop %rsi                /*这里为什么要pop rsi？在x86-64的函数调用与栈帧原理中，
                            调用者call指令会将返回地址压栈，这里我们想要保存的是调用者原本的栈信息，
                            即esp应该是将返回地址压栈前的esp。这里先把返回地址pop出即可得到原本的esp*/
    xorl %eax, %eax         /*设置返回值为0*/
    movq %rbx, (%rdi)
    movq %rsp, 8(%rdi)
    push %rsi               /*上一条指令已经将调用者原esp保存，现将rsi入栈，使得函数能够正常返回*/
    movq %rbp, 16(%rdi)
    movq %r12, 24(%rdi)
    movq %r13, 32(%rdi)
    movq %r14, 40(%rdi)
    movq %r15, 48(%rdi)
    movq %rsi, 56(%rdi)     /*线程恢复执行的地址，即save_context调用的下一条地址*/
    ret                     /*ret指令会跳转到返回地址处，也是save_context调用的下一条地址*/

协程栈帧恢复函数restore_context:

restore_context:
    movl $1, %eax           /*设置返回值为1*/
    movq (%rdi), %rbx       /*恢复栈帧*/
    movq 8(%rdi), %rsp
    movq 16(%rdi), %rbp
    movq 24(%rdi), %r12
    movq 32(%rdi), %r13
    movq 40(%rdi), %r14
    movq 48(%rdi), %r15
    jmp *56(%rdi)           /*恢复执行*/

下面我们来看一下如何使用这两个汇编形式的函数来实现协程上下文的切换：

void do_switch(Context &from, Context &to){
    int ret = save_context(&form);
    if(ret == 0){
        restore_context(&to);
    }
}

第一次调用save_context时，当前协程的上下文被保存到from中，返回值为0。需要注意的一点是，此时保存在from中的%rip寄存器指向的是调用save_context的下一条指令的地址，在do_switch函数中对应的是给ret赋值返回值%rax那条指令，这是最关键的地方。之后调用restore_context，恢复保存在to中的寄存器上下文，当执行完最后一条jmp指令后，就跳转到新协程的代码上去执行了，实际上完成了协程的切换。当后续的协程调度其再次恢复保存在from中的上下文时，注意到的是restore_context设置的返回值为1，即%rax的值为1，当jmp指令跳转到from中的%rip指向的指令（之前的ret赋值指令）时，将返回值1赋值给ret，就又回到了do_switch函数中的if判断语句，条件不成立，函数返回。

这其中最关键的地方是巧妙地利用了save_context和restore_context的返回值来判断是协程切换前还是切换后。

协程对象的构建

enum register_tt{
    rbx,
    rsp,
    rbp,
    r12,
    r13,
    r14,
    r15,
    pc_addr,
};

typedef struct{
    long buffer[8];
}ctx_buf_t;

typedef void *(*task_handler_t)(int id);
typedef int coroutine_status;
const int INIT = 0;
const int RUNABLE = 1;
const int WAIT = 2;
const int DEAD = 3;

class Coroutine{
    int tid;                            //协程ID号
    int para1;                          //可带一个参数
    void *stack;                        //栈底指针
    void *stack_top;                    //栈顶指针
    task_handler_t handler;             //任务函数指针
public:
    ctx_buf_t ctx;                      //寄存器上下文
    coroutine_status status;
    Coroutine(){};
    Coroutine(int id, task_handler_t handler, int para);
    Coroutine(const Coroutine &t);            //拷贝构造函数
    Coroutine(Coroutine &&t) noexcept;      //移动构造函数
    Coroutine& operator=(Coroutine &&t) noexcept;
    Coroutine& operator=(const Coroutine &t);
    ~Coroutine();
    void start();
    int get_id();
};

Coroutine::Coroutine(int id, task_handler_t handler, int para){
    int stack_size = (1<<20);   //1MB
    stack_top = malloc(stack_size);       //注意，这里的stack_top是栈可用空间最顶部（栈由高到低生长）
    para1 = para;
    tid = id;
    stack = stack_top + stack_size;          //指向栈的底部，ebp
    this->handler = handler;
    memset(&ctx, 0, sizeof(ctx_buf_t));
    ctx.buffer[rsp] = (long)stack;           //初始时rsp和rbp都指向栈底
    ctx.buffer[rbp] = (long)stack;
    ctx.buffer[pc_addr] = (long)handler;     //rip指向函数地址

    status = INIT;
}

从上述代码可以看出，我们实现的是有栈协程，构造协程类时会malloc出一定大小的栈空间。

协程的调度

之前我们介绍了协程上下文的切换，现在来考虑一下协程的调度。一种简单的想法是设置一个admin协程，每次admin协程都从协程队列中选择出一个最适合上台运行的协程，切换到该协程上运行，每次协程需要挂起的时候，都切换回admin协程，再由admin选择下一个上台运行的协程。可以由admin保证协程调度的公平性。

Hook库函数

回到我们设计协程的初衷，在高并发任务中需要开多个线程去执行任务，而线程的开销是比较大的（包括线程资源需要占的空间、线程调度需要系统调用陷入内核等等），这往往会成为系统的性能瓶颈。而协程调度则是完全在用户态完成，无需陷入内核态，开销是很小的，占据的内存空间也仅限于开辟的栈帧空间（可优化）。因此，单机可以开启的协程要比线程往往要多很多。

在线程中，遇到阻塞式的系统调用，例如Socket族的read函数，其在等待网络传输的数据达到时往往会导致线程的挂起。当然，在多线程的系统中，一个线程被网络阻塞而挂起，其它线程会被调度上台运行。但是在多协程的系统中，一个协程的阻塞会导致其宿主线程挂起，进而导致剩余的协程都无法被调度上台运行。那么该如何解决这个问题呢？

若第三方库无法正常使用协程，则对于整个程序协程便失去了意义。协程的出现就是为了减少IO操作阻塞而造成线程切换带来的开销，若例如mysql之类第三方库无法使用协程，则当调用其中的IO操作时仍然会造成线程切换。

一种简单的思路是将业务代码中的阻塞式的函数改为以非阻塞的方式运行。以read函数为例，在调用时底层尚未把数据准备好，正常情况下如果是以阻塞方式调用read函数，则该线程会被操作系统挂起。使用协程之后，我们需要将阻塞式的read改为非最的read，将对应的文件描述符fd交由epoll/poll去监听（意为此协程当前正在等待此fd的数据），然后将此协程挂起，控制权交换admin，admin选择下一个就绪的协程上台运行，当epoll/poll监听到fd数据就绪时，在稍后的调度中admin会选择对应的协程上台运行。

需要注意的是，若调用read时是以非阻塞的方式调用的，则应该遵循hook前后函数行为不变的原则，fd数据未就绪时直接返回即可。

HOOK是一个精细活，需要繁琐的边界条件测试，不但要保证返回值与原函数一致，相应的errno也要保持一致，做的与原函数越想，能够支持的第三方库就越多。

Hook read函数的示例代码：

int read(int fd, void *buf, size_t nbytes){
    int flags = fcntl(fd, F_GETFL, 0);

    size_t (*readcp)(int fd, void *buf, size_t nbytes);
    readcp = (size_t (*)(int fd, void *buf, size_t nbytes))dlsym(RTLD_NEXT, "read");
    
    if(flags&O_NONBLOCK){                   //若文件本身设置为非阻塞，则必须遵循hook前后函数行为不变的原则
        return readcp(fd, buf, nbytes);     //以非阻塞的方法调用read函数
        return 0;
    }

    std::cout<<"Task into Read WAIT status"<<endl;
    Scheduler::add_wait(fd, EPOLLIN);
    Scheduler::del_wait(fd);
    return readcp(fd, buf, nbytes);
}

将上述编译成为动态链接库myread.so

[root@xxxx] gcc -DRUNTIME -shared -fpic -o myread.so myread.cpp -ldl -std=c++11

运行主程序：

LD_PRELOAD='./myread.so' test

LD_PRELOAD的作用是指定一个共享库路径名的列表（以空格或者分号分隔），那么当你加载和执行一个程序，需要解析未定义的引用时，动态链接器会先去指定目录下搜索库，然后才搜索其它库。

关于dlsym函数的使用

通过dlsym()函数可以找到需要的符号。

void* dlsym(void *handle, char *symbol);

第一个参数是由dlopen()函数返回的动态库的句柄；第二个参数即所要查找的符号的名字。

当第一个参数为RTLD_NEXT时，其含义是：

If the first argument of dlsym' ordlvsym’ is set to RTLD_NEXT
the run-time address of the symbol called NAME in the next shared
object is returned. The “next” relation is defined by the order
the shared objects were loaded.

返回的是下一个被装载的对象中的符号地址。在库打桩机制中，一般是先生成包含dlsym的动态链接库，然后使用LD_PRELOAD指定该动态链接库的地址。这里面涉及到一个符号优先级的问题。

在本例中，运行主程序时，动态链接器首先去搜索LD_PRELOAD指定的动态链接库，也就是myread.so，然后再搜索其它的库（本例中是C语言标准库libc.so）。也就是说，myread.so中的read符号先被加载，而后libc中的read再被加载。这与我们在dlsym函数使用RTLD_NEXT的初衷一致，dlsym返回的就是之后加载的libc.so中的read符号的地址，这样就完成了一次hook。

而在test中调用的实际是myread.so中的read函数，为什么不是libc.so中的read呢？这是因为符号优先级。

符号优先级

当多个同名符号冲突时，先装入的符号优先，我们把这种优先级方式称为装载序列（Load Ordering）。