背景
最近在研究协程相关的内容，用 Rust 实现了一个单线程有栈协程原型。(代码地址)
本文主要介绍有栈协程的基本原理和实现，希望可以帮到对有栈协程原理感兴趣的朋友。
后面内容基于 ARM64 汇编，所以需要一点 ARM 汇编知识
最终效果如下
let runtime = Runtime::new();
// 创建多个协程任务，每个任务会执行多次并主动让出控制权
for task_id in 1..=5 {
runtime.spawn({
let runtime = runtime.clone();
move || {
for step in 1..=3 {
println!("任务 {} 正在执行步骤 {}", task_id, step);
// 主动让出控制权，让其他协程有机会执行
runtime.schedule();
}
println!("任务 {} 完成", task_id);
}
});
}
// 开始运行所有协程...
runtime.wait();
// 所有协程执行完成!
api 只有 3 个
spawn 创建一个协程
schedule 退出当前的协程 ，交还执行权
wait 等待所有协程执行完毕，中间会不断的上下文切出和切回
输出如下，可以看到任务是交替执行的
有栈协程简介
Golang 是有栈协程的实现标杆, 以下面的程序为例
package main
import "runtime"
func init() {
runtime.GOMAXPROCS(1)
}
func task(taskId int) {
for range 3 {
println("Task", taskId, "is running")
runtime.Gosched() // 协程切出，执行权全交给 runtime，runtime 调度选择下一个协程运行
}
}
func main() {
for taskId := range 3 {
go task(taskId + 1)
}
}
上面的代码创建 3 个协程，每个协程打印 3 次“当前 taskId 正在运行”
输出如下
可以看到 task 是交替执行的。每个 task 都在打印一次后切出， 由 Golang 的运行时（Runtime）调度下一个协程继续执行。
因为我设置了runtime.GOMAXPROCS(1) ，Golang 只会在一个线程上调度 goroutine
这种单线程的并发效果是在用户态实现的调度，这种机制被称为协程（coroutine），更具体一点的说法是有栈式协程（stackful coroutine）。
下面我会讲一下有栈是什么意思, 为什么需要栈。
函数栈帧
针对下面的函数
fn foo() {
task1();
task2();
}
函数调用开始的时候，会在栈空间上初始化一块空间，具体表现是移动 sp 指针创建一块内存作为函数栈帧。
栈上的数据通过 sp 的偏移来确定，一些值被 push 到栈内存储。
函数栈帧的大小在编译期就确定了
上图是创建 foo 函数栈帧的过程
如果在 foo 里面继续调用 task1 函数，就会创建另外一个栈帧，继续移动当前的sp指针
当 task1 函数调用退出的时候，sp 指针会还原，foo 可以继续执行 task2
这是正常函数调用的时候会发生的事情，但是在协程并发的情况下就不行了。
比如下面 golang 的代码
func foo(){
go task1();
task2();
}
假设 task1 运行到一半中断了
sp 指针回到 foo, foo 继续执行 task2
因为此时 task1 栈空间还没有回收，而 task2 创建的函数栈帧正好在 task1 的区域
task1 和 task2 的栈帧就会重叠，对应栈上的数据就会错乱
有栈式协程的核心原理
在堆上分配栈空间
有栈协程的解决方案是给每一个协程在堆里面分配一块内存空间作为栈空间
这样每个协程的栈数据不会相互干扰
在这种情况下, 就可以交替运行两个协程，互相不会干扰数据
func foo() {
go task1()
go task2()
}
每次 task1 中断的时候还原 sp 指针, 由 runtime 决定运行哪一个协程，如下图所示
所以有栈的含义是协程需要独立的栈空间
上下文切换
除了独立的栈空间，协程切换的时候还需要保存和还原上下文信息
上下文信息包括下面的寄存器信息 (ARM64汇编)
callee-save的寄存器，包括从x19 - x28的寄存器fp,lr寄存器，分配对应x29,x30sp寄存器- 返回地址 ，可以简单理解为当前协程的
pc寄存器
上下文切换的实现
对应 ARM64 汇编的实现
上下文信息
#[repr(C, align(16))]
pub struct CoroutineContext {
x19: Register, //0
x20: Register, //8
x21: Register, //16
x22: Register, //24
x23: Register, //32
x24: Register, //40
x25: Register, //48
x26: Register, //56
x27: Register, //64
x28: Register, //72
fp: Register, // 80
lr: Register, // 88
sp: Register, // 96
pc: Register, // 104
task: usize,
runtime_ptr: usize,
}
保存上下文
#[unsafe(naked)]
pub(crate) unsafe extern "C" fn store_context(
context_ptr: *mut CoroutineContext,
context_lr: usize,
resume_address: usize,
) {
naked_asm!(
"mov x8, x0",
"stp x19, x20, [x8]",
"stp x21, x22, [x8, #16]",
"stp x23, x24, [x8, #32]",
"stp x25, x26, [x8, #48]",
"stp x27, x28, [x8, #64]",
"stp fp, x1, [x8, #80]", // x1 是 lr
"mov x10, sp",
"str x10, [x8, #96]",
"str x2, [x8, #104]", // 保存返回地址，x2 是 resume_address
"ret",
);
}
还原上下文
#[unsafe(naked)]
pub(crate) unsafe extern "C" fn restore_context(context_ptr: usize) {
naked_asm!(
"mov x8, x0", // x8 是 context 上下文地址
"ldp x19, x20, [x8]",
"ldp x21, x22, [x8, #16]",
"ldp x23, x24, [x8, #32]",
"ldp x25, x26, [x8, #48]",
"ldp x27, x28, [x8, #64]",
"ldp fp, lr, [x8, #80]",
"ldr x10, [x8, #96]", // sp 指针
"mov sp, x10",
"ldr x10, [x8, #104]", // pc
"mov x0, x8", // context 作为第一个参数
"br x10", // 没必要返回，直接跳转过去了
);
}
实现有栈协程一定要写汇编管理寄存器
为了排除干扰，用 naked_asm 防止 rust 生成函数序言和尾部
上面就是核心代码，还是挺简单的
下面是实现一个单线程有栈协程的完整流程图
补充部分
有栈协程的开销
- 初始化
初始化的时候需要堆分配一块内存空间，比如 4KB，在运行时进行栈增长
- 上下文切换
在上下文切换的时候需要保存和还原大量的寄存器数据
即便如此，有栈协程的上下文切换还是比线程的切换要小，数量级比例大概在 1:100 左右
对比无栈协程
优势
有栈协程没有传染性的问题
劣势
无栈协程的初始化开销和上下文切换的开销相比无栈协程较大
栈增长
因为 4kb 有栈溢出的风险，所以需要在运行时进行栈增长
检测是否需要栈增长有两种手段
- 使用页保护，在协程的栈空间上预留一部分内存，当命中到这一块内存的时候认为需要栈增长
- 在每个函数调用之前判断是否会发生栈溢出（因为在编译器可以知道每个函数的栈帧有多大）
栈增长也有两种方式
- 重新分配一块更大的内存，把当前栈上的数据复制过去
- 分配一块新的堆空间，和当前的内存空间拼接起来
如果使用方案 1，不能只复制栈上的数据，还需要更新栈上的引用关系（比较麻烦）
使用方案 2，有可能在栈断掉的地方发生cache miss
总结
以上就是全部内容了，本人能力有限，如果有不对的地方，欢迎指出。
希望读者都可以理解有栈协程的原理
感谢
