M2: 协程库 (libco)

def positive_integers():
    i = 0
    while i := i + 1:
        yield i # "output" i

def is_prime(i):
    return i >= 2 and True not in (i % j == 0 for j in range(2, i))

primes = (i for i in positive_integers() if is_prime(i)) # 所有素数的顺序集合
# primes: <generator object <genexpr> at 0x7f142c14c9e8>

struct co *co_start(const char *name, void (*func)(void *), void *arg);
void       co_yield();
void       co_wait(struct co *co);

#include <stdio.h>
#include "co.h"

void entry(void *arg) {
  while (1) {
    printf("%s", (const char *)arg);
    co_yield();
  }
}

int main() {
  struct co *co1 = co_start("co1", entry, "a");
  struct co *co2 = co_start("co2", entry, "b");
  co_wait(co1); // never returns
  co_wait(co2);
}

#include <stdio.h>
#include "co.h"

int count = 1; // 协程之间共享

void entry(void *arg) {
  for (int i = 0; i < 5; i++) {
    printf("%s[%d] ", (const char *)arg, count++);
    co_yield();
  }
}

int main() {
  struct co *co1 = co_start("co1", entry, "a");
  struct co *co2 = co_start("co2", entry, "b");
  co_wait(co1);
  co_wait(co2);
  printf("Done\n");
}

b[1] a[2] b[3] b[4] a[5] b[6] b[7] a[8] a[9] a[10] Done

void create(void (*func)(void *));
void join(void (*func)());

3. 正确性标准

请将你所有的代码实现在 co.c 中

不要修改 libco.h。Online Judge 在评测时仅拷贝你的 co.c 文件。

首先，我们预期你提交的代码能通过附带的测试用例，测试用例有两组：

1. (Easy) 创建两个协程，每个协程会循环 100 次，然后打印当前协程的名字和全局计数器 g_count 的数值，然后执行 g_count++。
2. (Hard) 创建两个生产者、两个消费者。每个生产者每次会向队列中插入一个数据，然后执行 co_yield() 让其他 (随机的) 协程执行；每个消费者会检查队列是否为空，如果非空会从队列中取出头部的元素。无论队列是否为空，之后都会调用 co_yield() 让其他 (随机的) 协程执行

执行 make test 会在 x86-64 和 x86-32 两个环境下运行你的代码——如果你看到第一个测试用例打印出数字 X/Y-0 到 X/Y-199、第二个测试用例打印出 libco-200 到 libco-399，说明你的实现基本正确；否则请调试你的代码。

Online Judge 上会运行类似的测试 (也会在 x86-64 和 x86-32 两个平台上运行)，但规模可能稍大一些。你可以假设：

1. 每个协程的堆栈使用不超过 64 KiB；
2. 任意时刻系统中的协程数量不会超过 128 个 (包括 main 对应的协程)。协程 wait 返回后协程的资源应当被回收——我们可能会创建大量的协程执行-等待-销毁、执行-等待-销毁。因此如果你的资源没有及时回收，可能会发生 Memory Limit Exceeded 问题。

还需要注意的是，提交的代码不要有任何多余的输出，否则将会被 Online Judge 判错。如果你希望在本地运行时保留调试信息并且不想在提交到 Online Judge 时费力地删除散落在程序中的调试信息，你可以尝试：

• 用环境变量控制调试开关，仅当 LIBCO_DEBUG 环境变量存在时才输出调试信息；
• 用编译选项，例如将 -DLIBCO_DEBUG 传递给 gcc，在程序中仅当 LIBCO_DEBUG 被定义时才输出调试信息。

在新手阶段，你很容易觉得做上面两件事会比较受挫 (又要去 STFW/RTFM 了)；但花点时间研究它们绝对是值得的——你已经知道，Online Judge 提交以后，返回的一般都是 Wrong Answer，来回注释调试信息的配置可不是个好主意，还有可能因为忘记注释掉调试信息出错。

4. 实验指南

不要慌。

这个实验有点难。以前的实验都是有明确目标的，比如 OJ 题给定输入和输出。但这次不一样，我们要 hack C 语言运行时的行为——写一个函数 “切换” 到另一个函数执行。听起来就无从下手。

不要慌，“计算机系统公理” 告诉我们，只要你想做的事情，就一定能做到。让我们一起分析这个问题。

4.1. 编译成动态链接库

和 pstree 把源文件编译成一个二进制文件不同，本实验把源文件编译成后缀名为 .so 的共享库 (shared object, 动态链接库)。使用 make 编译后，将会得到 libco-32.so 和 libco-64.so：

$ file libco-64.so 
libco-64.so: ELF 64-bit LSB shared object, x86-64, version 1 (SYSV), ...

共享库可以有自己的代码、数据，且可以调用其他的库函数 (例如 libc)。共享库中全局的符号将能被加载它的应用程序调用。共享库中不需要入口 (main 函数)。我们的 Makefile 里已经写明了如何编译共享库：

$(NAME)-64.so: $(DEPS) # 64bit shared library
    gcc -fPIC -shared -m64 $(CFLAGS) $(SRCS) -o $@ $(LDFLAGS)

其中 -fPIC -fshared 就代表编译成位置无关代码的共享库。除此之外，共享库和普通的二进制文件没有特别的区别。虽然这个文件有 +x 属性并且可以执行，但会立即得到 Segmentation Fault (可以试着用 gdb 调试它)。当然，你的确可以让你的共享库真正可以执行 (并得到有意义的结果)，同时作为命令行工具和动态链接库，有兴趣的同学请 STFW。

4.2. 编写自己的程序链接共享库

为了降低大家 STFW/RTFM 的难度，我们提供了一组协程库的测试用例 (tests/ 目录下)，包含了编译和运行所需的脚本，其中编译的编译选项是

gcc -I.. -L.. -m64 main.c -o libco-test-64 -lco-64
gcc -I.. -L.. -m32 main.c -o libco-test-32 -lco-32

注意到 -I 和 -L 选项的使用：

• -I 选项代表 include path，使我们可以 #include <co.h>。在 threads.h 中我们已经介绍过这个编译选项；
• -L 选项代表增加 link search path。
• -l 选项代表链接某个库，链接时会自动加上 lib 的前缀，即 -lco-64 会依次在库函数的搜索路径中查找 libco-64.so 和 libco-64.a，直到找到为止。如果你将 libco-64.so 删除后用 strace 工具查看 gcc 运行时使用的系统调用，就能清晰地看到库函数解析的流程；

在运行时，使用 make test 直接运行，它执行的命令是：

LD_LIBRARY_PATH=.. ./libco-test-64

如果不设置 LD_LIBRARY_PATH 环境变量，你将会遇到 “error while loading shared libraries: libco-xx.so: cannot open shared object file: No such file or directory” 的错误。请 STFW 理解这个环境变量的含义。

至此，我们已经完成了共享库的编译，以及让一个 C 程序动态链接共享库执行。

小技巧：调试你的代码

直接运行 ./libco-test-64 会遇到 No such file or directory 的问题；当然，只需要用 make test 就可以解决这个问题了。但如果想要调试代码？gdb libco-test-64 同样也会遇到共享库查找失败的问题。大家可以在终端中使用 export 将当前 shell 进程的 LD_LIBRARY_PATH 设置好，这样就可以无障碍地运行 ./libco-test-64 了——这和大家 export AM_HOME=... 的原理是完全一致的。

4.3. 协程切换：分析

不妨让我们先写一小段协程的测试程序：

void foo() {
  int i;
  for (int i = 0; i < 1000; i++) {
    printf("%d\n", i);
    co_yield();
  }
}

这个程序在 co_yield() 之后可能会切换到其他协程执行，但最终它依然会完成 1, 2, 3, ... 1000 的打印。那么，co_yield 到底需要做什么呢？我们不妨先看一下这段代码编译后的汇编指令序列：

push   %rbp
push   %rbx
lea    <stdout>, %rbp       
xor    %ebx, %ebx
sub    $0x8, %rsp
mov    %ebx, %esi
mov    %rbp, %rdi
xor    %eax, %eax
callq  <printf@plt>
inc    %ebx
xor    %eax, %eax
callq  <co_yield>  // <- 切换到其他协程
cmp    $0x3e8, %ebx
jne    669 <foo+0xf>
pop    %rax
pop    %rbx
pop    %rbp
retq    

首先，co_yield 必须将当前的执行流 “封存” 下来，否则，如果你很着急地在 co_yield 中执行堆栈切换，例如

mov (next_rsp), $rsp

那么当前执行流的栈顶指针就永远地消失，co_yield 也就永远无法返回了。具体来说，我们必须保存的数据是：

• 所有被调用者保存的通用寄存器，因为后续指令执行可以任意修改 (例如循环变量保存在 %ebx 中，如果 %ebx 寄存器数据丢失，我们就丢失了 i 的数值)；但 %rdi 寄存器就不需要保存。
• 所有栈帧上的数据。当 co_yield() 返回后，上面的这段汇编代码会持续访问堆栈中的内容 (例如 pop 和 ret 指令)。

因此，为了实现 co_yield，我们需要做的事情其实是：

1. 为每一个协程分配独立的堆栈；堆栈顶的指针由 %rsp 寄存器确定；
2. 在 co_yield 发生时，将寄存器保存到属于该协程的 struct co 中 (包括 %rsp)；
3. 切换到另一个协程执行，找到系统中的另一个协程，然后恢复它 struct co 中的寄存器现场 (包括 %rsp)。

例如，参考实现的 struct co 是这样定义的：

enum co_status {
  CO_NEW = 1, // 新创建，还未执行过
  CO_RUNNING, // 已经执行过
  CO_WAITING, // 在 co_wait 上等待
  CO_DEAD,    // 已经结束，但还未释放资源
};

struct co {
  char *name;
  void (*func)(void *); // co_start 指定的入口地址和参数
  void *arg;

  enum co_status status;  // 协程的状态
  struct co *    waiter;  // 是否有其他协程在等待当前协程
  jmp_buf        context; // 寄存器现场 (setjmp.h)
  uint8_t        stack[STACK_SIZE]; // 协程的堆栈
};

看起来就像是在《计算机系统基础》中实现的上下文切换！我们推荐大家使用 C 语言标准库中的 setjmp/longjmp 函数来实现寄存器现场的保存和恢复。在《计算机系统基础》实验中，我们已经用汇编代码实现了这两个函数。没有好好做实验的同学，要加油补上啦！

4.4. 实现寄存器现现场切换

分配堆栈是容易的：堆栈直接嵌入在 struct co 中即可，在 co_start 时初始化即可。但麻烦的是如何让 co_start 创建的协程，切换到指定的堆栈执行。AbstractMachine 的实现中有一个精巧的 stack_switch_call (x86.h)，可以用于切换堆栈后并执行函数调用，且能传递一个参数，请大家完成阅读理解 (对完成实验有巨大帮助)：

static inline void stack_switch_call(void *sp, void *entry, uintptr_t arg) {
  asm volatile (
#if __x86_64__
    "movq %0, %%rsp; movq %2, %%rdi; jmp *%1"
      : : "b"((uintptr_t)sp),     "d"(entry), "a"(arg)
#else
    "movl %0, %%esp; movl %2, 4(%0); jmp *%1"
      : : "b"((uintptr_t)sp - 8), "d"(entry), "a"(arg)
#endif
  );
}

理解上述函数你需要的文档：GCC-Inline-Assembly-HOWTO。当然，这个文档有些过时，如果还有不明白的地方，gcc 的官方手册是最佳的阅读材料。

警告：堆栈对齐

x86-64 要求堆栈按照 16 字节对齐 (x86-64 的堆栈以 8 字节为一个单元)，这是为了确保 SSE 指令集中 XMM 寄存器变量的对齐。如果你的程序遇到了神秘的 Segmentation Fault (可能在某个 libc 的函数中)，如果你用 gdb 确定到 Segmentation Fault 的位置，而它恰好是一条 SSE 指令，例如

movaps %xmm0,0x50(%rsp)
movaps %xmm1,0x60(%rsp)
...

那很可能就是你的堆栈没有正确对齐。我们故意没有说的是，System V ABI (x86-64) 对堆栈对齐的要求，是在 “何时” 做出的——在 call 指令之前按 16 字节对齐，在 call 指令之后就不对齐了。一方面你可以暴力地尝试一下；如果你想更深入地理解这个问题，就需要读懂 stack_switch_call，以及 STFW 关于 ABI 对对齐的要求，或是查看编译出的汇编代码。

每当 co_yield() 发生时，我们都会选择一个协程继续执行，此时必定为以下两种情况之一 (思考为什么)：

1. 选择的协程是新创建的，此时该协程还没有执行过任何代码，我们需要首先执行 stack_switch_call 切换堆栈，然后开始执行协程的代码；
2. 选择的协程是调用 yield() 切换出来的，此时该协程已经调用过 setjmp 保存寄存器现场，我们直接 longjmp 恢复寄存器现场即可。

当然，上述过程描述相当的抽象；你可能会花一点时间，若干次试错，才能实现第一次切换到另一个协程执行——当然，这会让你感到非常的兴奋。之后，你还会面对一些挑战，例如如何处理 co_wait，但把这些难关一一排除以后，你会发现你对计算机系统 (以及 “程序是个状态机”) 的理解更深刻了。

4.5. 实现协程

非常难理解？坚持住！

没错，的确很难理解。如果你没有完成《计算机系统基础》中的 setjmp/longjmp 实验，你需要多读一读 setjmp/longjmp 的文档和例子——这是很多高端面试职位的必备题目。如果你能解释得非常完美，就说明你对 C 语言有了脱胎换骨的理解。setjmp/longjmp 的 “寄存器快照” 机制还被用来做很多有趣的 hacking，例如实现事务内存、在并发 bug 发生以后的线程本地轻量级 recovery 等等。

setjmp/longjmp 类似于保存寄存器现场/恢复寄存器现场的行为，其实模拟了操作系统中的上下文切换。因此如果你彻底理解了这个例子，你们一定会觉得操作系统也不过如此——我们在操作系统的进程之上又实现了一个迷你的 “操作系统”。类似的实现还有 AbstractMachine 的native，它是通过 ucontext.h 实现的，有兴趣的同学也可以尝试阅读 AbstractMachine 的代码。

在参考实现中，我们维护了 “当前运行的协程” 的指针 (这段代码非常类似于操作系统中，为每一个 CPU 维护一个 “当前运行的进程”)：

struct co *current;

这样，在 co_yield 时，我们就知道要将寄存器现场保存到哪里。我们使用的代码是

void co_yield() {
  int val = setjmp(current->context);
  if (val == 0) {
    // ?
  } else {
    // ?
  }
}

在上面的代码中，setjmp 会返回两次：

• 在 co_yield() 被调用时，setjmp 保存寄存器现场后会立即返回 0，此时我们需要选择下一个待运行的协程 (相当于修改 current)，并切换到这个协程运行。
• setjmp 是由另一个 longjmp 返回的，此时一定是因为某个协程调用 co_yield()，此时代表了寄存器现场的恢复，因此不必做任何操作，直接返回即可。

最后，框架代码里有一行奇怪的 CFLAGS += -U_FORTIFY_SOURCE，用来防止 __longjmp_chk 代码检查到堆栈切换以后报错 (当成是 stack smashing)。Google 的 sanitizer 也遇到了相同的问题。

4.6. 资源初始化、管理和释放

需要初始化？

如果你希望在程序运行前完成一系列的初始化工作 (例如分配一些内存)，可以定义 __attribute__((constructor)) 属性的函数，它们会在 main 执行前被运行。我们在课堂上已经讲解过。

这个实验最后的麻烦是管理 co_start 时分配的 struct co 结构体资源。很多时候，我们的库函数都涉及到资源的管理，在面向 OJ 编程时，大家养成了很糟糕的习惯：只管申请、不管释放，依赖操作系统在进程结束后自动释放资源。但如果是长期运行的程序，这些没有释放但又不会被使用的泄露资源就成了很大但问题，例如在 Windows XP 时代，桌面 Windows 是没有办法做到开机一星期的，一周之后机器就一定会变得巨卡无比。

管理内存说起来轻巧——一次分配对应一次回收即可，但协程库中的资源管理有些微妙 (但并不复杂)，因为 co_wait 执行的时候，有两种不同的可能性：

1. 此时协程已经结束 (func 返回)，这是完全可能的。此时，co_wait 应该直接回收资源。
2. 此时协程尚未结束，因此 co_wait 不能继续执行，必须调用 co_yield 切换到其他协程执行，直到协程结束后唤醒。

希望大家仔细考虑好每一种可能的情况，保证你的程序不会在任何一种情况下 crash 或造成资源泄漏。然后你会发现，假设每个协程都会被 co_wait 一次，且在 co_wait 返回时释放内存是一个几乎不可避免的设计：如果允许在任意时刻、任意多次等待任意协程，那么协程创建时分配的资源就无法做到自动回收了——即便一个协程结束，我们也无法预知未来是否还会执行对它的 co_wait，而对已经回收的 (非法) 指针的 co_wait 将导致 undefined behavior。C 语言中另一种常见 style 是让用户管理资源的分配和释放，显式地提供 co_free 函数，在用户确认今后不会使用时释放资源。

资源管理一直是计算机系统世界的难题，至今很多系统还受到资源泄漏、use-after-free 的困扰。例如，顺着刚才资源释放的例子，你可能会感觉 pthread 线程库似乎有点麻烦：pthread_create() 会修改一个 pthread_t 的值，线程返回以后资源似乎应该会被释放。那么：

• 如果 pthread_join 发生在结束后不久，资源还未被回收，函数会立即返回。
• 如果 pthread_join 发生在结束以后一段时间，可能会得到 ESRCH (no such thread) 错误。
• 如果 pthread_join 发生在之后很久很久很久很久，资源被释放又被再次复用 (pthread_t 是一个的确可能被复用的整数)，我不就 join 了另一个线程了吗？这恐怕要出大问题。

实际上，pthread 线程默认是 “joinable” 的。joinable 的线程只要没有 join 过，资源就一直不会释放。特别地，文档里写明： Failure to join with a thread that is joinable (i.e., one that is not detached), produces a "zombie thread". Avoid doing this, since each zombie thread consumes some system resources, and when enough zombie threads have accumulated, it will no longer be possible to create new threads (or processes).

这就是实际系统中各种各样的 “坑”。在《操作系统》这门课程中，我们尽量不涉及这些复杂的行为，而是力图用最少的代码把必要的原理解释清楚。当大家对基本原理有深入的理解后，随着经验的增长，就会慢慢考虑到更周全的系统设计。