Context 上下文切换

最近将一个 PHP 项目整体从传统的 PHP-FPM 迁移到了 Swow,写篇文章记录一下协程实现相关内容作为备忘。相较于进程和线程来说,协程的底层实现是非常简单的,仅仅涉及到了用户态栈的切换。Swow 和 Swoole 一样,选择了性能优秀的 boost.context 作为底层实现。

boost.context 是基于汇编的,每个平台的实现都不相同,但总体思路是类似的,因此这里就只关注 Linux x86-64 架构下的 make_fcontextjump_fcontext 两个基础接口。

函数调用和栈帧

fcontext 的结构和函数调用的栈帧非常接近,所以不妨先熟悉一下普通函数调用时栈帧布局:

stack-frame

可以看到,被调函数栈帧的上方就是被调函数的参数和调用结束后的返回地址,被调函数栈帧开头先保存了调用函数的 %rbp 和一些由 Callee 保存的寄存器值。

fcontext 结构

fcontext 的结构可以直接参考 make_x86_64_sysv_elf_gas.S 中的示图:

fcontext

结构上和栈帧非常相似,翻译成 struct 大致相当于:

struct fcontext {
    int32_t fc_mxcsr;   //--------------
    int32_t fc_x87_cw;  //
    int64_t r12;        // context 负责保存
    int64_t r13;        // 的寄存器
    int64_t r14;        //
    int64_t r15;        //--------------
    int64_t rbx;        // context_func_t 函数指针
    int64_t rbp;        // context 返回地址
    int64_t rip;        // context 入口地址
}

当然,这么类比并不准确,还是直接看到代码更加方便。

make_fcontext

make_fcontext 函数的原型如下:

typedef void *context_t;
typedef struct
{
    context_t from_context;
    void *data;
} transfer_t;
typedef void (*context_function_t)(transfer_t transfer);
context_t make_fcontext(void *stack, long unsigned int stack_size, context_function_t func);

一共接受三个参数,stack 是为协程分配的栈底,stack_size 是栈的大小,func 是协程的函数指针。可以看到这里并没有 context_t 类型的参数,这是因为 context 结构是保存在栈的底部。

.file "make_x86_64_sysv_elf_gas.S"
.text
.globl make_fcontext
.type make_fcontext,@function
.align 16

/* make_fcontext
 * 
 * @param stack         stack 栈底(rdi)
 * @param stack_size    stack 大小(rsi)
 * @param func          入口函数指针(rdx)
 * 
 * @return              context_t(rax)
 */
make_fcontext:
    /* 将 stack 地址保存到 rax 寄存器 */
    movq  %rdi, %rax
    /* 对齐地址 */
    andq  $-16, %rax

    /* 将 rax 中的地址向低方向移动 64 个字节,相当于在栈底保留了 context 的位置
     *          -----------------------------------------------------------------
     * stack:   |                                  |         context         ||||
     *          -------------------------------------------------------------|---
     *              <-----低方向-----            %rax                    16-align
     */
    leaq  -0x40(%rax), %rax

    /* 将第三个参数保存到 0x28(%rax) 的位置
     *          -----------------------------------------------------------------
     * context: |   |   |       |       |       |       |  FUNC |       |       |
     *          -----------------------------------------------------------------
     *          0   4   8       16      24      32      40      48      56
     */
    movq  %rdx, 0x28(%rax)

    /*         
     *          -----------------------------------------------------------------
     * context: |MMX|x87|       |       |       |       |  FUNC |       |       |
     *          -----------------------------------------------------------------
     *          0   4   8       16      24      32      40      48      56
     */
    stmxcsr  (%rax)
    fnstcw   0x4(%rax)

    /* 保存 context 入口到 0x38(%rax)        
     *          -----------------------------------------------------------------
     * context: |MMX|x87|       |       |       |       |  FUNC |       | ENTRY |
     *          -----------------------------------------------------------------
     *          0   4   8       16      24      32      40      48      56
     */
    leaq  trampoline(%rip), %rcx
    movq  %rcx, 0x38(%rax)

    /* 保存返回地址到 0x30(%rax)        
     *          -----------------------------------------------------------------
     * context: |MMX|x87|       |       |       |       |  FUNC |  EXIT | ENTRY |
     *          -----------------------------------------------------------------
     *          0   4   8       16      24      32      40      48      56
     */
    leaq  finish(%rip), %rcx
    movq  %rcx, 0x30(%rax)

    /* 返回函数,返回值 %rax 中的地址即为 context 地址     
     *         %rax
     *          |
     *          -----------------------------------------------------------------
     * context: |MMX|x87|       |       |       |       |  FUNC |  EXIT | ENTRY |
     *          -----------------------------------------------------------------
     *          0   4   8       16      24      32      40      48      56
     */
    ret 

trampoline:
    /* 在以前的版本中并没有这部分代码,而是将 func 作为入口 */
    /* 提交记录中说是为了对齐 stack,修复某中特殊情况的 BUG,暂时不明 */
    push %rbp
    jmp *%rbx

finish:
    /* 退出 */
    xorq  %rdi, %rdi
    call  [email protected]
    hlt
.size make_fcontext,.-make_fcontext

/* Mark that we don't need executable stack. */
.section .note.GNU-stack,"",%progbits

jump_fcontext

typedef void *context_t;
typedef struct
{
    context_t from_context;
    void *data;
} transfer_t;
transfer_t jump_fcontext(context_t const to_context, void *data);

jump_fcontext 的原型如上,接受两个参数,目标 context 地址和参数指针 data,返回的 transfer_t 中保存了当前 context 的信息和数据指针 data。

.file "jump_x86_64_sysv_elf_gas.S"
.text
.globl jump_fcontext
.type jump_fcontext,@function
.align 16

/* jump_fcontext
 *
 * @param context_t         目标 context(rdi)
 * @param data              数据参数指针(rsi) 
 *
 * @return transfer_t       from_context(rax),data(rdx)
 */
jump_fcontext:
    /* 将栈顶向下移动 0x38 个字节,开辟保存当前 context 数据的空间
     * 由于这里是函数调用,call 会自动将返回地址入栈,因此只需要 0x38 个字节即可
     *
     *                                   |<----context---->|
     *                                   |<---0x38--->|    |
     *          -------------------------|------------|----|---------------------
     * stack:   |                        |            | RA |     ......      ||||
     *          -------------------------|------------|----------------------|---
     *              <-----低方向-----   %rsp       old_rsp               16-align
     */
    leaq  -0x38(%rsp), %rsp /* prepare stack */

    /* 依次将需要保存的寄存器值保存到 context 对应的位置中     
     *         %rsp
     *          |
     *          -----------------------------------------------------------------
     * context: |MMX|x87|  R12  |  R13  |  R14  |  R15  |  RBX  |  RBP  |  RA   |
     *          -----------------------------------------------------------------
     *          0   4   8       16      24      32      40      48      56
     */
#if !defined(BOOST_USE_TSX)
    stmxcsr  (%rsp)     
    fnstcw   0x4(%rsp)  
#endif
    movq  %r12, 0x8(%rsp) 
    movq  %r13, 0x10(%rsp)
    movq  %r14, 0x18(%rsp)
    movq  %r15, 0x20(%rsp)
    movq  %rbx, 0x28(%rsp)
    movq  %rbp, 0x30(%rsp)

    /* 保存当前 %rsp 到 %rax 寄存器 */
    movq  %rsp, %rax

    /* 将 %rsp 切换到 to_context 的地址     
     *         %rsp
     *          |
     *          -----------------------------------------------------------------
     * context: |MMX|x87|  R12  |  R13  |  R14  |  R15  |  RBX  |  RBP  |  RIP  |
     *          -----------------------------------------------------------------
     *          0   4   8       16      24      32      40      48      56
     */
    movq  %rdi, %rsp

    /* 将返回地址(RIP)保存到 %r8 寄存器 */ 
    movq  0x38(%rsp), %r8

    /* 依次将 to_context 中的数据恢复到寄存器中 */
#if !defined(BOOST_USE_TSX)
    ldmxcsr  (%rsp) 
    fldcw    0x4(%rsp)
#endif
    movq  0x8(%rsp), %r12 
    movq  0x10(%rsp), %r13
    movq  0x18(%rsp), %r14
    movq  0x20(%rsp), %r15
    movq  0x28(%rsp), %rbx
    movq  0x30(%rsp), %rbp

    /* 将栈顶向上移动 0x40 个字节,回收 context 占用的空间
     *
     *                                                                     %rsp
     *                                                                       |
     *          -------------------------------------------------------------|---
     * stack:   |                                          |     context     ||||
     *          -------------------------------------------------------------|---
     *                       <-----低方向-----          old_rsp          16-align
     */
    leaq  0x40(%rsp), %rsp

    /* 构造返回数据 transfer,此时 %rax 为 from_context 地址 */
#if !defined(_ILP32)
    /* RAX == fctx, RDX == data */
    movq  %rsi, %rdx
#else
    /* RAX == data:fctx */
    salq  $32, %rsi
    orq   %rsi, %rax
#endif

    /* 将 transfer 作为参数传递给 to_context 入口地址 */
#if !defined(_ILP32)
    /* RDI == fctx, RSI == data */
#else
    /* RDI == data:fctx */
#endif
    movq  %rax, %rdi

    /* 跳转到返回地址开始执行 */
    jmp  *%r8
.size jump_fcontext,.-jump_fcontext

/* Mark that we don't need executable stack.  */
.section .note.GNU-stack,"",%progbits

这里需要注意一下在 jump_fcontext 中保存当前 context 时只在栈上留了 0x38 个字节,并没有预留返回地址的空间。这是因为在调用 jump_fcontextcall 指令会自动把返回地址入栈,可以参考前面的栈帧分布图。

测试

写一个简单的程序测序一下上下文切换:

// test.c
#include <stdio.h>
#include <stdlib.h>

#define STACK_SIZE (4 * 1024)

typedef void *context_t;
typedef struct
{
    context_t from_context;
    void *data;
} transfer_t;
typedef void (*context_function_t)(transfer_t transfer);

context_t make_fcontext(void *stack, size_t stack_size, context_function_t function);
transfer_t jump_fcontext(context_t const to_context, void *data);

void func(transfer_t transfer) {
    printf("I am in func.\n");
    transfer = jump_fcontext(transfer.from_context, NULL);
    printf("I am in func again!\n");
    jump_fcontext(transfer.from_context, NULL);
}

int main() {
    char stack[STACK_SIZE];
    context_t to_context = make_fcontext(stack + STACK_SIZE, STACK_SIZE, func);
    transfer_t transfer;
    printf("I am in main.\n");
    transfer = jump_fcontext(to_context, NULL);
    printf("I am in main again!\n");
    jump_fcontext(transfer.from_context, NULL);
    printf("End of main\n");
    return 0;
}

编译并运行:

$ gcc -g -c make_x86_64_sysv_elf_gas.S -o make.o
$ gcc -g -c jump_x86_64_sysv_elf_gas.S -o jump.o
$ gcc -g malloc.c make.o jump.o -o test
$ ./test
I am in main.
I am in func.
I am in main again!
I am in func again!
End of main