Skip to content

协程库实现

1. 背景与核心思想

1.1 协程的概念

  • 定义：协程（coroutine）是一种轻量级用户态协同程序（green threads/user-level threads），允许在一个线程内通过主动切换实现多个执行流。
  • 特点：
  • 不依赖操作系统线程，运行在用户态，资源占用少。
  • 协程通过显式调用 co_yield() 主动让出控制权，切换到其他协程。
  • 共享进程地址空间，支持共享内存多任务并发。
  • 与线程的对比：
  • 线程调度由操作系统和硬件控制，可能因中断或多核并行切换。
  • 协程调度由程序显式控制，切换点由 co_yield() 决定，适合状态机、actor模型、goroutine等场景。
  • 协程无需操作系统支持，资源占用更小，切换开销低。

1.2 状态机视角

  • 程序是状态机：每个协程是一个独立的状态机，状态由堆栈、寄存器和共享内存组成。
  • 协程切换：保存当前协程的寄存器和堆栈状态，恢复另一个协程的状态，实现状态机切换。
  • 操作系统模型：实验中的操作系统模型通过状态机管理程序（如 hello.py），协程库模拟类似的管理机制。

2. 协程库 API

2.1 API 定义

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

2.2 API 功能

  • co_start(name, func, arg)：
  • 创建新协程，返回 struct co 指针。
  • 新协程从 func 开始执行，传入参数 arg，但不立即运行，调用者继续执行。
  • 使用 malloc() 分配 struct co 内存，具体结构在 co.c 中定义。
  • co_yield()：
  • 当前协程暂停执行，保存状态。
  • 随机选择一个可运行协程（包括自身），恢复其状态并继续执行。
  • co_wait(co)：
  • 等待指定协程 co 执行完成。
  • 协程结束后，释放 co_start 分配的内存（使用 free()）。
  • 每个协程（除初始协程外）必须被 co_wait 恰好一次，否则可能导致内存泄漏。

2.3 使用示例

#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");
}
  • 输出：随机交替打印 a 和 b，计数从 1 递增到 10，如 b[1] a[2] b[3]...a[10] Done。
  • 特点：协程共享内存（如 count），通过 co_yield() 随机切换。

3. 实现原理

3.1 数据结构

协程状态

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

协程结构体

struct co {
    char *name;                     // 协程名称
    void (*func)(void *);           // 入口函数
    void *arg;                      // 函数参数
    enum co_status status;          // 协程状态
    struct co *waiter;              // 等待该协程的协程
    struct context context;         // 寄存器上下文
    uint8_t stack[STACK_SIZE];      // 独立堆栈（实验假设 ≤ 64 KiB）
};
  • context：保存寄存器状态（如 rax, rbx, rsp, rip 等）。
  • stack：每个协程有独立堆栈，栈顶由 context 中的 rsp 指向。

全局变量

struct co *current; // 当前运行的协程
  • 初始时为 main 协程，切换时更新为目标协程。

3.2 上下文切换

核心机制

  • 保存状态：将当前协程的寄存器（包括 rsp）保存到 current->context。
  • 恢复状态：选择目标协程，加载其 context 中的寄存器到 CPU。
  • 堆栈切换：更新 rsp 指向目标协程的独立堆栈。

使用 setjmp 和 longjmp

  • 实现 co_yield：
void co_yield() {
    int val = setjmp(current->context);
    if (val == 0) {
        // 保存当前状态，选择下一个协程
        choose_next_co();
        longjmp(next_co->context, 1);
    } else {
        // 从 longjmp 恢复，直接返回
        return;
    }
}
  • setjmp 保存当前寄存器状态，返回 0 时选择并切换到下一个协程。
  • longjmp 恢复目标协程的寄存器状态，setjmp 返回非 0 值，协程继续执行。
  • 优点：
  • 符合 x86-64 调用约定，自动处理部分寄存器（如 rdi, rsi 等 caller-saved 寄存器）。
  • 简化实现，无需保存浮点寄存器。

新协程启动

  • 新协程（CO_NEW）未保存上下文，需初始化堆栈并调用入口函数。
  • 使用 stack_switch_call 切换堆栈并执行函数：
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) : "memory"
#else
        "movl %0, %%esp; movl %2, 4(%0); jmp *%1"
        : : "b"((uintptr_t)sp - 8), "d"(entry), "a"(arg) : "memory"
#endif
    );
}
  • 设置 rsp 指向新协程的堆栈，rdi 传递参数，跳转到 entry 函数。

3.3 协程创建与初始化

  • co_start：
  • 使用 malloc 分配 struct co。
  • 初始化 name, func, arg, status（设为 CO_NEW）。
  • 分配独立堆栈，设置初始 rsp 和 rip（指向 func）。
  • 将协程加入可运行协程列表。
  • 初始化：
  • 可通过 __attribute__((constructor)) 定义初始化函数，在 main 前运行（如分配 main 协程）。
  • 确保 main 协程正确初始化，参与调度。

3.4 协程等待与资源释放

  • co_wait：
  • 检查目标协程状态：
    • 若已结束（CO_DEAD），释放资源（free(co)）并返回。
    • 若未结束，标记当前协程为 CO_WAITING，调用 co_yield 切换，直到目标协程结束。
  • 释放资源时，确保 co->stack 和 co 本身被正确回收。
  • 资源管理：
  • 每个协程（除 main）必须被 co_wait 一次，防止内存泄漏。
  • 使用 coroutine_wrapper 包装 func，在协程返回时执行清理：
void coroutine_wrapper(void (*func)(void *), void *arg) {
    func(arg);
    // 标记协程为 CO_DEAD，通知等待者，释放资源
}
  • 避免资源泄漏，确保长期运行程序的稳定性。

3.5 随机调度

  • 选择下一个协程：
  • 维护可运行协程列表（状态为 CO_NEW 或 CO_RUNNING）。
  • 使用随机算法（如 rand()）选择下一个协程。
  • 包括当前协程（可能切换回自身）。
  • 实现：
  • 在 co_yield 中调用 choose_next_co()，更新 current。
  • 新协程通过 stack_switch_call 启动，已运行协程通过 longjmp 恢复。