协程上下文切换

liang2kl@outlook.com (Liang Yesheng) — Thu, 26 May 2022 13:02:01 +0800

最近在研究 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 的责任。