用户态调度要保存些什么

一、问题

我们希望可以在同一进程用户态下自由调度一些线程(threads of execution),这些是有各自独立ProgramCounter和栈空间的真正线程,而不是像protothreads那样的stackless伪线程.

举例来说:

   Thread_1                Thread_2

   func_A                  func_C
      |                      |
      |----func_B            |---func_D
             |                    |   
           P |-----switch()------>|
             | \    ......        |
             |  \<--switch()------|
             |                    |

如上图,Thread_1的函数调用链为?->A->B->?,在B中执行的某个时刻P,由于暂时缺乏某资源或数据,我们希望可以调到Thread_2之前执行的某个位置继续执行,当Thread_1条件成熟后再适时切回P处继续执行.

有什么用途?
比如Go语言中的goroutine,某些支持协程/微线程的网络框架,或者一些通用的用户线程库等等. 这里switch的目标不一定直接指向另一个线程,还可以经由某个中间统一调度模块做路由.

有什么好处?
相比OS的preemptive进线程调度,用户态调度的优势有:
1) 综合切换开销很低(同普通函数调用),完全不涉及OS上下文切换、刷TLB和cache污染等系统调度问题.
2) 相比preemptive调度,应用主动让出的时机可以自己控制,更有利于资源利用率(但无法强制保证公平).

也就是说:
a) 随着线程数量的增多,花在切换上的开销占比不会像OS线程那样增大;
b) 线程数量基本只受限于内存.

再加上我们偏爱线程模型的原始初衷:编程清晰简单,固然能移植到用户态也很好.

最后用户态程序有权限以及有必要用到的维护thread运行时关键资源是很小的一部分,因此可以有十分简单、高效的实现.

二、如何实现

总体思路比较简单,首先每个线程需要有单独的栈空间,例如可以从堆上为每个线程分配栈空间,然后从A switch 到B时先将A的执行相关上下文保存起来,再加载B之前保存的执行上下文信息,并替换当前的上下文,切换完成.

本文重点是说明该上下文里最少要保存哪些信息.

与内核调度不同,很多进程级的信息如pgt页表寄存器和内核栈帧等明显不需要关注,用户态所涉及的内容是其中一个较小的子集,用户态调度要简单得多.

在x86-64平台下,直接影响进程内用户态线程正确执行的内容有:
a) 指令地址PC
b) 栈指针rsp
c) 通用寄存器 rax、rbx、...、r15
d) 状态寄存器CC
e) sse/浮点数寄存器

其中a、b是明显需要的,它们是线程之所以能独立执行的根基.
c、d、e看上去好像也都是需要的,因为其他线程可能会做修改,但实际上还有进一步压缩的空间,下面依次来看看.

三、Caller-Saved Registers 和 Callee-Saved Registers.

上面说需要保存通用整数寄存器rax、rbx、...、r15的原因是,线程切走后可能会被别的线程更改,因而需要保存以便恢复. 但其实在x86-64下结合当前常见编译器的寄存器使用策略来看,并不需要全部保存.

先来看下在普通函数调用场景下对这个问题的处理. 比如函数P调用函数Q,也会有类似的问题:Q可能会更改掉P正在使用的一些寄存器的内容,从而导致Q返回后P出错. 解决的办法是编译器们约定一套统一的寄存器使用策略. 具体来说将这些整数寄存器划分为若干类:

返回值寄存器:rax
参数寄存器:rdi、rsi、rdx、rcx、r8、r9,分别对应函数的(从左往右)第1-6个参数(第7个以上参数通过压栈传递).
Callee-Saved Register: rbx、rbp、r12、r13、r14、r15
Caller-Saved Register: r10、r11

顾名思义返回值寄存器和参数寄存器一般专门用来存函数返回值和传递参数.
Callee-Saved Register: callee负责恢复原值的寄存器. 例如在Q中要修改rbx,修改前先push rbx保存,然后ret之前pop rbx恢复其值,从而在上层函数P看来rbx的内容是不变的.
Caller-Saved Register: callee可以随意更改,需要caller自己维护(保存到内存)的寄存器.

当返回值寄存器/参数寄存器有空余时,也可以当做Caller-Saved寄存器来做它用.

可以看出在此约定下,callee恢复所有其修改过的callee-saved-registers, caller自己维护caller-saved-registers,一个函数可以同时作为callee和caller,但寄存器值call/ret前后一致性都可以得到有效保障.

然而我们的线程调度是跨越栈的,不再是caller-callee关系,上述约定关系的某些部分会被破坏,具体来说:
1) 从线程A调度到线程B,B不会为A恢复callee-saved-register,因而这些寄存器需要保存到A的上下文中做恢复.
2) 从线程A调度到线程B,因为是通过siwtch()函数调用的形式,A仍然会自行(在栈上)保存caller-saved-register,不需要再次保存:

void thread_A_run()
{
   //调用switch_to()之前会保存caller-saved-register到当前栈上
   switch_to(thread_B)  //对编译器来说调度入口依然是函数调用的形式
   //后续再调度回这里,形式上等同于从上述函数返回,恢复caller-saved-register.
}

结论:整型通用寄存器中,只需要保存callee-saved-registers.

其他浮点数/sse向量寄存器等在某些规范(如LLVM)下也是callee-saved,原则上也需要保存,但可能实际使用的情况很少,因此可以提供单独的功能版本或不提供支持. 本文假设应用场景不涉及这些寄存器,不予以考虑.

四、Condition Code

再来看 d) 状态寄存器CC. 这个寄存器中保存一些条件分支指令执行所需要的一些Condition Code,例如:

CF: Carry Flag, 进位标志位.
ZF: Zero Flag, 最近一次运算结果是否为0.
SF: Sign Flag, 最近一次运算结果是否为负.
OF: Overflow Flag, 补码运算溢出位.

如何影响条件语句:

    if (a > b)
    {
       block_1;
    } 
    return;

    假设rdi保存a, rsi保存b, 上述语句一般编译为:

    cmpq %rsi, %rdi   //set CC according to a-b
    jle .Done         //if (a <= b) goto Done   ('jle' checks: (SF ^ OF) | ZF) 
    (block_1)         //do block_1
Done:
    ret

以上可以看出,if (expr) 这样的语句至少要分两条单独指令运行,先设置CC,再执行jxx跳转,而jxx需要访问CC的内容.

在普通OS环境下,这两条语句之间有可能发生软硬件中断,中断后当前cpu可能会调度其他进程执行:

   cmpq %rsi, %rdi  //set CC
   // ------ interruptted, other process may now change CC ------
   jle .Done        //jump according to CC

因此在OS调度时一定会保存/恢复CC的内容,是内核线程上下文的关键信息之一.

但是在用户态调度下,我们不用考虑CC,因为与内核抢占式调度不同,我们是主动让出cpu,没有(用户态)异步中断存在,只有我们调用switch()才会发生切换,因而正在执行if (a > b)这样的语句的时候是没有机会调用switch()的:

   if (a > b)   //no way to switch() here in the middle of 'if'
   {
      ...
      switch()  //can only switch() after 'cmp' and 'jxx' are done as whole.
   }
   ...

结论:在用户态下对CC的set_and_test操作一定是原子操作,因此不用保存CC

五、控制流

至此,在x86-64和通常的编译器规则下,可以得出用户态线程至少要保存以下内容作为调度上下文:

a) 指令地址PC
b) 栈指针rsp
c) 6个 callee-saved-registers: rbx、rbp、r12、r13、r14、r15

即至少要保存8个机器字长的内容.

再来看下调度控制的问题. 我们先假设switch实现方式如下:

   typedef long (*ctx_buf)[8];   //6个callee-saved寄存器,1个rsp,1个返回地址

   //从from(当前线程)切换到to线程.
   void switch(ctx_buf from, ctx_buf to)
   {
      save_context(from)     //保存from的上下文到某个位置, 待后续切回时做恢复
      restore_context(to)    //加载to的上下文并替换、jmp, 即恢复to线程的上下文
   }

这比较清晰,但有个问题,save_context(from)这里保存的是其下一条语句即restore_context函数调用的地址,假如后续从to再切回from时,又会跑到restore_context(to)执行,导致from永远无法再次执行,这显然是不正确的.

解决办法:
(1) save_context不保存其下一条语句即restore_context的地址,而是保存上层switch返回后的下一条语句地址.
(2) save_context后面不直接跟restore,而是插入条件判断返回值(rax),返回值根据控制流是从本线程save_context(from)返回还是通过其他线程调用restore_context(from)返回而设置不同的值,从而可以避免再次调用restore_context(to).

方法(1) 需要跨两层栈帧传递返回地址,不太方便,方法(2)足够简单,因而这里采用方法(2)后重写switch如下:

   void switch(ctx_buf from, ctx_buf to)
   {
       int ret = save_context(from);   //保存上下文, 首次返回会设置eax = 0
       if (ret == 0) 
       {
           restore_context(to);         //加载to的上下文, jmp之前会设置eax = 1
       }
       else
       {
           //restored from other threads, just return and continue
       }
   }

这样从其他线程恢复from的执行时,控制流会从if语句继续运行,而此时ret值为1(被其他线程通过调用restore_context设置),从而落入else分支从switch返回继续from线程的执行.

了解setjmp/longjmp的同学可能一眼就认出来了吧, 确实是差不多的.

最后补上save_context和restore_context,基本就是在save/load线程的context buffer


/*
 * ctx.S: x86-64 only
 *
 * ctx_buf[0]: rbx
 * ctx_buf[1]: rsp
 * ctx_buf[2]: rbp
 * ctx_buf[3]: r12
 * ctx_buf[4]: r13
 * ctx_buf[5]: r14
 * ctx_buf[6]: r15
 * ctx_buf[7]: Program Counter
*/

//rdi保存参数: from线程上下文buffer
save_context:
	pop  %rsi	    //将rsi设置为返回地址,即save_context调用的下一条语句; 注意这里修改了rsp	
	xorl %eax,%eax	    //设置返回值 eax = 0
	movq %rbx,(%rdi)    //from[0] = rbx
	movq %rsp,8(%rdi)   //from[1] = rsp, 在下一条语句之前保存rsp,此时rsp正好为调用save_context之前的值
	push %rsi	    //前面把返回地址pop掉了,这里再push回去,让函数能正常返回; 注意这里修改了rsp	
	movq %rbp,16(%rdi)  //from[2] = rbp
	movq %r12,24(%rdi)  //from[3] = r12
	movq %r13,32(%rdi)  //from[4] = r13
	movq %r14,40(%rdi)  //from[5] = r14
	movq %r15,48(%rdi)  //from[6] = r15
	movq %rsi,56(%rdi)  //from[7] = rsi, 即线程恢复执行的地址	
	ret

//rdi保存参数: to线程上下文buffer
restore_context:
	movl $1,%eax	    //设置返回值 eax = 1		
	movq (%rdi),%rbx    //rbx = to[0]
	movq 8(%rdi),%rsp   //rsp = to[1], 恢复rsp
	movq 16(%rdi),%rbp  //rbp = to[2]
	movq 24(%rdi),%r12  //r12 = to[3]
	movq 32(%rdi),%r13  //r13 = to[4]
	movq 40(%rdi),%r14  //r14 = to[5]
	movq 48(%rdi),%r15  //r15 = to[6]
	jmp *56(%rdi)       //jmp 到 to[7]返回地址 恢复执行
        

六、Put It All Together

下面给一个可以运行的简单demo,两个user-level threads相互调度:

#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include <sys/mman.h>

enum register_t
{
    rbx,
    rsp,
    rbp,
    r12,
    r13,
    r14,
    r15,
    pc_addr,
};

typedef struct 
{
    /*
     * buffer[0]: rbx
     * buffer[1]: rsp
     * buffer[2]: rbp
     * buffer[3]: r12
     * buffer[4]: r13
     * buffer[5]: r14
     * buffer[6]: r15
     * buffer[7]: Program Counter
     */
    long buffer[8];
} ctx_buf_t;

typedef void *(*thread_handler_t)(void);
typedef struct 
{
  int id;
  void *stack;
  thread_handler_t handler;
  ctx_buf_t ctx; 
} thread_t;

//symbol from ctx.S 
extern int save_context(ctx_buf_t *from);
extern int restore_context(ctx_buf_t *to);

//user thread switch
void do_switch(thread_t *from, thread_t *to)
{
    int ret = save_context(&from->ctx); //保存上下文, 首次返回会设置eax = 0
    if (ret == 0) 
    {
        restore_context(&to->ctx);     //加载to的上下文, jmp之前会设置eax = 1
    }
    else
    {
        //restored from other threads, just return and continue
    }
}

//create a thread
int thread_create(thread_t *t, int id, thread_handler_t handler)
{
    int stack_size = (1 << 20);
    void *stack_top = malloc(stack_size);  //or mmap with stack type sepcified

    t->id = id;
    t->stack = stack_top + stack_size; //栈是从上往下增长的,因此栈基址要指向最大处
    t->handler = handler;
    memset(&t->ctx, 0, sizeof(ctx_buf_t));
    t->ctx.buffer[rsp] = (long)t->stack;  //initialize stack pointer
    t->ctx.buffer[pc_addr] = (long)t->handler;  //initialize program counter
    return id;
}

void thread_destory(thread_t *t)
{
    free(t->stack);
}

//start a thread from main
int thread_start(thread_t *t)
{
    restore_context(&t->ctx);

    //never returns
    return -1;
}


thread_t g_thread_A;
thread_t g_thread_B;

void *func_A()
{
    printf("Thread [A] Start running in func_A...\n");

    //infinite loop for demo
    while (1)
    {
        printf("Thread [A] Switching to another thread...\n");

        do_switch(&g_thread_A, &g_thread_B);  //switch to thread B

        printf("Thread [A] returned to func_A, Stack Base[%p]\n", g_thread_A.stack);
        printf("Thread [A] doing stuff...\n");
        sleep(1);  //pretends to be busy in this thread. But this actually blocks the whole process.
    }
    return NULL;
}

void *func_B()
{
    printf("Thread [B] Start running in func_B...\n");

    while (1)
    {
        printf("Thread [B] Switching to another thread...\n");

        do_switch(&g_thread_B, &g_thread_A);  //switch to thread A

        printf("Thread [B] returned to func_B, Stack Base[%p]\n", g_thread_B.stack);
        printf("Thread [B] doing stuff...\n");
        sleep(1);  //pretends to be busy in this thread. But this actually blocks the whole process.
    }
    return NULL;
}


int main()
{
    thread_create(&g_thread_A, 1, func_A);     //A is ready
    thread_create(&g_thread_B, 2, func_B);     //B is ready

    thread_start(&g_thread_A);                 //A is started

    //never reaches here in this demo
    thread_destory(&g_thread_A);
    thread_destory(&g_thread_B);
    return 0;    
}

/*

gcc ctx.S demo.c -o demo

./demo output:

Thread [A] Start running in func_A...
Thread [A] Switching to another thread...
Thread [B] Start running in func_B...
Thread [B] Switching to another thread...
Thread [A] returned to func_A, Stack Base[0x7fd9d9bcc010]
Thread [A] doing stuff...
Thread [A] Switching to another thread...
Thread [B] returned to func_B, Stack Base[0x7fd9d95fa010]
Thread [B] doing stuff...
Thread [B] Switching to another thread...
Thread [A] returned to func_A, Stack Base[0x7fd9d9bcc010]
Thread [A] doing stuff...
...

*/

(https://github.com/woofyzhao/tiny_sched)

看上去很像是两个function相互调用并没有做什么有趣的事情,但它们确实是处于两个完整的拥有独立栈空间的执行线程内,因而调度的位置可以是任意深的,不会影响每个栈上的函数调用链的逻辑. 即我们实现了用户态调度的基本切换机制.

这个例子中原始main线程实际上被永久挂起了,因为我们没有保存其状态所以理论上也回不来了;在实际场景中我们可以把main线程当做一个单独的统一管理调度策略的模块,比如所有user thread统一switch到main,main再负责按某种策略选取一个thread执行,如果大家都不能执行则在main中挂起进程等待相应的事件发生即可.

工业级实现参考:微信的libco

七、参考资料
[1] CSAPP 3rd, chapter 3
[2] Glibc setjmp/longjmp
[3] SPP

发表评论

电子邮件地址不会被公开。