协程上下文切换

最近在研究 Arachne 做用户态线程调度的方法。

众所周知，一个线程的状态包括：

PC
一些寄存器（即“上下文”）

切换线程时，需要切换上下文；为了将来能够恢复上下文，在切换时我们需要将当前的上下文保存在某个地方。我们尝试用 C 写一个简单的上下文切换。

上下文切换⌗

进行上下文切换的汇编代码在上汇编课时就讲过：

swapcontext:
    # save current context
    pushq %r12
    pushq %r13
    pushq %r14
    pushq %r15
    pushq %rbx
    pushq %rbp
    movq %rsp, (%rsi)

    # restore previously saved context
    movq (%rdi), %rsp
    popq %rbp
    popq %rbx
    popq %r15
    popq %r14
    popq %r13
    popq %r12

上面的代码做了几件事情：

保存当前上下文
1. 将被调用者保存寄存器保存到当前的调用栈上
2. 将当前的栈指针 %rsp 保存到 %rsi 所指向的内存地址上
恢复之前保存的上下文
1. 将栈指针恢复为 %rdi 所保存的另一个协程的栈顶
2. 恢复另一个协程的栈上面保存的被调用者保存寄存器

这里需要理解一个问题：为什么需要保存被调用者保存寄存器？

“被调用者保存寄存器”指的是在调用子过程前后不发生改变的寄存器，即子过程需要在返回前将这些寄存器的值恢复。上面的 swapcontext 是一个过程调用，因此需要遵循这一惯例。然而，与一般的过程调用不同，我们在这里实际上是切换到了另一个上下文中，因此我们实际上是恢复了另一个协程的上下文。

这里可能会产生疑惑：调用者保存寄存器呢？实际上，调用 swapcontext 的过程在调用它之前就已经把调用者保存寄存器保存好了（如果有必要的话）。理解的关键是：swapcontext 是一个过程，而不是内联的汇编代码。

调用栈结构⌗

由于 swapcontext 中会将当前的上下文保存到某个内存地址，我们之后只需要将这个内存地址传入 swapcontext 就可以切换回来了。在这里，我们主要考虑在一开始如何切换到一个指定的函数上；我们需要初始化调用栈，使得 swapcontext 后能够跳转到指定的函数。

不难画出这样的一个调用栈：

      +-----------------------+
      |                       |
      +-----------------------+
      |     Return Address    |
      +-----------------------+
      |                       |
      |       Registers       |
sp->  |                       |
      +-----------------------+
      |                       |
      |                       |

我们将需要切换到的函数的地址保存在 Return Address 中。在使用 popq 指令恢复所有保存的寄存器后，rsp 指向此函数地址，swapcontext 返回（ret 指令）后便跳转到函数的入口地址。需要注意的是 Return Address 上方需要保留 8 字节的空间。

栈的初始化过程如下：

#define SPACE_FOR_SAVED_REGISTERS 48UL
#define STACK_SIZE 1024UL

void routine2() {
    // ...
}

void *sp2;
char stack[8 + SPACE_FOR_SAVED_REGISTERS + STACK_SIZE];

int main() {
    // Initialize stack for routine2
    sp2 = stack + SPACE_FOR_SAVED_REGISTERS + STACK_SIZE - 2 * sizeof(void *);

    // Set return address
    *(void **)sp2 = routine2;

    // Point to bottom of the register stack
    sp2 = (char *)sp2 - SPACE_FOR_SAVED_REGISTERS;

    //...
}

在这里，我们创建一个栈，并将 routine2 的地址写到相应位置；sp2 指向最后一个储存寄存器值的位置。

此时，stack、sp2 和调用栈的结构如下：

        +-----------------------+
        |                       |
        +-----------------------+
        | Address of `routine2` |
        +-----------------------+
        |                       |
        |       Registers       |
sp2->   |                       |
        +-----------------------+
        |                       |
        |                       |
        |      Stack Space      |
        |                       |
stack-> |                       |
        +-----------------------+

程序⌗

最后，我们写几个简单的上下文切换的 demo。

例子 #1⌗

在这里，我们进行两次上下文切换，从 routine 切换到 routine2，再切换回 routine。

其中，swapcontext 需要加上 noinline 的 attribute，这是为了保证 swapcontext 作为过程被调用。

另外，编译时需要加上 -fomit-frame-pointer，否则编译器在编译时会在过程前后分别加上 pushq %rbp 和 popq %rbp，导致栈结构遭到破坏。

#include <stdio.h>

#define SPACE_FOR_SAVED_REGISTERS 48UL
#define STACK_SIZE 1024UL

void __attribute__((noinline)) swapcontext(void** saved, void** target) {
    asm("pushq %r12\n\t"
        "pushq %r13\n\t"
        "pushq %r14\n\t"
        "pushq %r15\n\t"
        "pushq %rbx\n\t"
        "pushq %rbp\n\t"
        "movq %rsp, (%rsi)");

    asm("movq (%rdi), %rsp\n\t"
        "popq %rbp\n\t"
        "popq %rbx\n\t"
        "popq %r15\n\t"
        "popq %r14\n\t"
        "popq %r13\n\t"
        "popq %r12");
}

void *sp2;
void *sp;
char stack[8 + SPACE_FOR_SAVED_REGISTERS + STACK_SIZE];

void routine() {
    printf("routine()\n");
    // Switch to routine2
    swapcontext(&sp2, &sp);
    printf("routine() end\n");
}

void routine2() {
    printf("routine2()\n");
    // Switch back
    swapcontext(&sp, &sp2);
    // Unreachable
    printf("routine2() end\n");
}

int main() {
    // Initialize stack for routine2
    sp2 = stack + SPACE_FOR_SAVED_REGISTERS + STACK_SIZE - 2 * sizeof(void *);

    // Set return address
    *(void **)sp2 = routine2;

    // Point to bottom of the register stack
    sp2 = (char *)sp2 - SPACE_FOR_SAVED_REGISTERS;

    routine();

    printf("main routine\n");
    return 0;
}

输出如下：

routine()
routine2()
routine() end
main routine

例子 #2⌗

继续尝试一些其他的用法。在这里，我们计算 0 累加至 9 的结果，并在每次循环时都切换一次上下文，然后立刻切换回来。

void routine() {
    printf("routine()\n");
    int sum = 0;

    for (int i = 0; i < 10; i++) {
        sum += i;
        swapcontext(&sp2, &sp);
    }

    printf("sum: %d\n", sum);
}

void routine2() {
    for (;;) {
        printf("routine2()\n");
        swapcontext(&sp, &sp2);
    }
}

输出为：

routine()
routine2()
routine2()
routine2()
routine2()
routine2()
routine2()
routine2()
routine2()
routine2()
routine2()
sum: 45

汇编结果⌗

使用 -O1 编译上面第二个例子（去掉所有 printf），routine 的汇编代码如下：

0000000000000639 <routine>:
 639:	55                   	push   %rbp
 63a:	53                   	push   %rbx
 63b:	bb 0a 00 00 00       	mov    $0xa,%ebx
 640:	48 8d 2d f9 09 20 00 	lea    0x2009f9(%rip),%rbp        # 201040 <sp>
 647:	48 89 ee             	mov    %rbp,%rsi
 64a:	48 8d 3d 47 0e 20 00 	lea    0x200e47(%rip),%rdi        # 201498 <sp2>
 651:	e8 af ff ff ff       	callq  605 <swapcontext>
 656:	83 eb 01             	sub    $0x1,%ebx
 659:	75 ec                	jne    647 <routine+0xe>
 65b:	5b                   	pop    %rbx
 65c:	5d                   	pop    %rbp
 65d:	c3                   	retq

可以看到，for 循环使用 %ebx（被调用者保存）进行计数，在调用 swapcontext 前后并没有保存其值。保存这个寄存器是 swapcontext 的责任。