100行代码实现一个RISC-V架构下的多线程管理框架

1. 摘要

本文将基于RISC-V架构和qemu仿真器实现一个简单的多线程调度和管理框架, 旨在通过简单的代码阐明如何实现线程的上下文保存和切换, 线程的调度并非本文的重点, 故线程调度模块只是简单地采用了轮询的方式.

2. 上下文是什么

对于一个运行中的程序, 我们可以把它看作一个状态机, CPU每次执行一条指令会导致这个状态机从一个状态转换到另一个状态, 而这里的状态就是CPU的寄存器值和当前时刻内存的状态, 即state = (Regs, Mem). 一个程序的执行过程就是状态机状态迁移的过程, 如果希望在一个线程执行过程到某条指令时切换到另一个线程, 我们必须额外实现一段代码去完成这种切换, 并把这段代码插入到希望的切换点的后面, 这段代码必须要满足两个条件: 1): 能够正确保存当前线程A的状态. 我们可以先假设完成这段切换的代码不存在, 线程A本该执行它的下一条指令, 而在执行下一条指令之前, 线程A有一个确定的状态, 即CPU的寄存器和内存中的值, 具体来说包括pc指针, 栈寄存器sp, 通用寄存器等以及线程使用到的内存的状态. 线程切换代码, 首先必须能够当前线程A的状态保存起来, 以便后续调度器再切换回当前的这个线程. 2): 能够正确恢复被选中的线程B的状态. 为了切换到另一个线程B去执行, 我们必须从保存线程B状态的内存地址处读取线程B的状态, 然后更新CPU寄存器以及内存的状态, 更新完成之后就可以继续线程B被中断的执行流了. 这里保存和恢复的线程状态就是线程的上下文, 在一个线程还没有运行过的情况下, 我们需要手工构造一个线程的上下文, 这个上下文的作用是使得线程一旦被调度, 就能够从创建线程时指定的入口地址处执行.

3. 保存上下文

以下都以RISC-V为基础进行讨论. 从线程A的角度看, 当它在自己的正常执行流中调用我们实现的上下文切换代码时, 只是相当于调用了一个普通函数, 因此编译器会根据RISC-V架构的ABI约定, 在调用这个函数之前去保存caller saved registers, 即ABI约定中应该由调用者去保存的寄存器, 并且在调用结束之后插入恢复这些寄存器的指令序列, 上下文切换代码段首先要做的就是保存ABI中约定由被调用者保存的寄存器, 我们可以在表示线程的结构体中开辟一块区域用来保存这些寄存器, 具体实现可以定义一个描述上下文的结构体, 并在线程结构体中嵌入一个上下文描述结构体.

4. 切换上下文

线程切换代码段的下一个任务是恢复线程B的状态到CPU和内存, 只要使用若干load指令去恢复线程B的寄存器状态即可, 最后使用调用返回指令ret就完成了线程的切换. ret指令会根据ra寄存器的值更新pc寄存器, 从而完成跳转, 而ra寄存器保存了线程B被中断的那条指令的地址, 需要注意的是, 在RISC-V架构下, 我们无法通过类似mv这样的指令直接修改pc寄存器的值.

5. 轮询调度

线程调度算法并非本文讨论重点, 在代码实现中只简单定义了一个线程数组, 调度器在调度时依次选择数组中处于READY状态的线程去执行, 当没有线程可调度时, 调度器会自动退出.

6. 代码实现

代码部分需要使用RISC-V工具链进行编译, 线程切换代码使用简单的汇编指令实现, 调试和运行需要qemu仿真工具, 软件工具的安装这里不进行讨论, 详情可参考相关网络资源.

6.1 整体流程

进入到main函数之后, 首先创将两个线程, 然后调用本文实现调度模块提供的接口对线程列表中的线程进行调度, 直到没有可供调度的线程.

int main(int argc, char **argv)
{
    qthread_create(thread1_start);
    qthread_create(thread2_start);
    printf("Start to run all threads...\n");
    qthread_run_all();
    printf("All threads are done!\n");

    return 0;
}

两个线程的入口函数分别为thread1_start和thread2_start, 线程完成的任务如下:

void *thread1_start(void *args)
{
    int i = 1;
    while (i < 1000)
    {
        if ((i % 100) == 0)
        {
            printf("Thread #1: running\n");
            qthread_yield();
        }
        i++;
    }
    qthread_exit();
    return 0;
}

两个线程的入口函数功能类似, 都会在使用占用一段时间的CPU之后主动放弃使用权, 任务完成之后主动退出, 通过qemu仿真器可以得到运行结果:

junan@u1:~/Documents/random_coding/01_context_swtch$ qemu-riscv64 -L /usr/riscv64-linux-gnu/ ./main
Start to run all threads...
Thread #1: running
Thread #2: running
Thread #1: running
Thread #2: running
Thread #1: running
Thread #2: running
Thread #1: running
Thread #2: running
Thread #1: running
Thread #2: running
Thread #1: running
Thread #2: running
Thread #1: running
Thread #2: running
Thread #1: running
Thread #2: running
Thread #1: running
Thread #2: running
All threads are done!

6.2 线程调度模块接口

对于本文实现的线程调度模块qthread, 其对外的接口包含:

int qthread_create(qthread_entry_t entry);
void qthread_yield();
void qthread_exit();
void qthread_run_all();

分别完成线程的创建, 线程释放CPU, 线程退出以及线程调度.

6.3 线程描述结构体

为了管理创建出来的线程, 需要一个描述线程的结构体, 记录线程的id, 调度状态, 线程上下文等信息. 需要注意的是, 每个线程的运行都需要栈的支持, 这里在线程结构体中定义了一个1024字节大小的数组, 用于线程函数运行使用的栈, 在创建线程时, 会手动设置线程上下文中的sp指针的值.

typedef struct
{
    uint64 tid;
    enum thread_state state;
    qthread_entry_t qthread_entry;
    char stack[1024];
    struct context context;
} qthread_t;

创建线程的接口实现如下:

int qthread_create(qthread_entry_t entry)
{
    int index_found = -1;
    for (int i = 0; i < MAX_THREADS_NR; ++i)
    {
        if (qthread_table[i].state == UNUSED)
        {
            index_found = i;
            break;
        }
    }

    if (index_found < 0 || index_found >= MAX_THREADS_NR)
    {
        return -1;
    }

    qthread_table[index_found].tid = get_next_tid();
    qthread_table[index_found].state = READY;
    qthread_table[index_found].qthread_entry = entry;

    qthread_table[index_found].context.ra = (uint64)entry;
    qthread_table[index_found].context.sp = (uint64)qthread_table[index_found].stack + 1024;

    return 0;
}

因为栈的增长方向是由高地址向低地址, 所以线程上下文中的sp指针赋值时需要加上栈的大小, 否则会出现栈溢出的问题.

6.4 上下文描述结构体

对于RISC-V架构, 在进行上下文切换时需要保存ABI约定中callee维护的寄存器外加ra寄存器, 因为ra寄存器记录了调用完线程切换代码段对应的函数之后的返回地址, 即线程A被打断的指令地址. 上下文描述结构体如下:

struct context
{
    uint64 ra;
    uint64 sp;
    uint64 s0;
    uint64 s1;
    uint64 s2;
    uint64 s3;
    uint64 s4;
    uint64 s5;
    uint64 s6;
    uint64 s7;
    uint64 s8;
    uint64 s9;
    uint64 s10;
    uint64 s11;
};

6.5 线程切换代码段

如上文所述, 线程切换代码段需要保存一个上下文, 并恢复另一个上下文, 这部分代码使用汇编指令去实现:

    .section .text
    .global swtch
swtch:
    sd ra, 0(a0)
    sd sp, 8(a0)
    sd s0, 16(a0)
    sd s1, 24(a0)
    sd s2, 32(a0)
    sd s3, 40(a0)
    sd s4, 48(a0)
    sd s5, 56(a0)
    sd s6, 64(a0)
    sd s7, 72(a0)
    sd s8, 80(a0)
    sd s9, 88(a0)
    sd s10, 96(a0)
    sd s11, 104(a0)

    ld ra, 0(a1)
    ld sp, 8(a1)
    ld s0, 16(a1)
    ld s1, 24(a1)
    ld s2, 32(a1)
    ld s3, 40(a1)
    ld s4, 48(a1)
    ld s5, 56(a1)
    ld s6, 64(a1)
    ld s7, 72(a1)
    ld s8, 80(a1)
    ld s9, 88(a1)
    ld s10, 96(a1)
    ld s11, 104(a1)

    ret

代码的前半部分用于保存上下文, 后半部分用于恢复上下文, 按照ABI约定, 寄存器a0和a1分别用来传递函数调用的第一个和第二个参数, 这里在调用时传递的就是两段保存和恢复各自上下文的内存地址. 最后的ret指令, 会用寄存器ra的值更新寄存器pc, 因为ra已经被恢复成了线程B被中断的指令地址, 所以ret指令执行过后CPU就会执行线程B对应的代码. 在恢复上下文时也更新了栈指针寄存器sp, 所以线程B在执行时使用的时自己的栈内存. 线程主动释放CPU资源时会调用qthread_yield(), 在这个函数中会调用线程切换函数swtch:

void qthread_yield()
{
    current->state = READY;
    swtch(&current->context, &qthread_context);
}

对于调用qthread_yield()的线程A来说, 当其内部调用swtch函数时, 实际上CPU已经切换到了另一个线程B去执行, 但线程A是不能感知到的, 当这个swtch函数返回之后, 对线程A来说就像调用了一个普通的函数一样, 线程A会继续自己的执行流.

6.6 轮询调度实现

调度采用轮询的方式, 每次选取一个可用线程去切换.

void qthread_run_all()
{
    int count = 0;
    while (1)
    {
        for (int i = 0; i < MAX_THREADS_NR; ++i)
        {
            if (qthread_table[i].state == READY)
            {
                current = &qthread_table[i];
                current->state = RUNNING;
                count++;
                swtch(&qthread_context, &current->context);
            }
        }

        if (count == 0)
            break;
        count = 0;
    }
}

7. gdb调试上下文切换

• 首先启动qemu-riscv64, 指定gdb调试端口:

junan@u1:~/Documents/random_coding/01_context_swtch$ qemu-riscv64 -L /usr/riscv64-linux-gnu/ -g 1234 ./main

• 使用gdb-multiarch连接调试端口, 下发调试命令:

junan@u1:~/Documents/random_coding/01_context_swtch$ gdb-multiarch main
GNU gdb (Ubuntu 12.0.90-0ubuntu1) 12.0.90
Copyright (C) 2022 Free Software Foundation, Inc.
License GPLv3+: GNU GPL version 3 or later <http://gnu.org/licenses/gpl.html>
This is free software: you are free to change and redistribute it.
There is NO WARRANTY, to the extent permitted by law.
Type "show copying" and "show warranty" for details.
This GDB was configured as "x86_64-linux-gnu".
Type "show configuration" for configuration details.
For bug reporting instructions, please see:
<https://www.gnu.org/software/gdb/bugs/>.
Find the GDB manual and other documentation resources online at:
    <http://www.gnu.org/software/gdb/documentation/>.

For help, type "help".
Type "apropos word" to search for commands related to "word"...
Reading symbols from main...
(gdb) target remote :1234

调用swtch之前各个寄存器的状态如下:



此时正在运行的是线程1, 即将进入swtch函数, 在swtch函数执行过程中寄存器的状态如下:



可以看到在swtch函数中, ret指令执行之前, ra寄存器的值保存的是调度函数qthread_run_all()中某条指令的地址, 因此ret执行完之后CPU首先跳转到调度函数, 调度函数接着被中断的循环继续运行, 并从线程列表中选取下一个READY状态的线程, 然后再使用swtch函数切换到被选中的线程, 同时把自己的上下文保存到一个全局的上下文处. 即用户线程的切换需要依赖调度函数qthread_run_all, 线程释放CPU时会先切换到qthread_run_all的上下文, 然后由这个调度器去选取下一个应该运行的线程并进行切换.

8. 总结

本文使用尽可能简单的代码阐述了上下文切换的原理. 个人认为简单的情况更有利于理解一项技术背后的原理, 当然, 文中讲述的上下文切换不仅能够用来实现一个demo. 设想以下, 当一个用户进程通过一个系统调用进入到内核态时, 假如这个syscall想访问某个此时不可用硬件资源, 那么内核通常情况下需要将当前的这个进程休眠, 即需要调度器选择另一个可用进程占据当前的这个CPU, 这里同样需要进行上下文切换. 内核需要保存即将休眠的进程在内核态中的上下文, 然后切换到调度器, 再由调度器去完成进程调度并切换到被选中的进程. 实际上这个过程涉及到的技术原理和本文讨论的内容时一致的. 对于用户态切换到内核态的上下文保存后续文章会展开讨论.

本文的完整代码链接如下: https://github.com/kfggww/random_coding/tree/main/01_context_swtch, 欢迎"一键三连" , 代码实现或者文中表述若有不妥之处还清批评指正.

9. 参考资源

[1] https://riscv.org/technical/specifications/

[2] https://pdos.csail.mit.edu/6.S081/2021/