天地孤影任我行

python的Greenlet模块源码分析

Posted on 2015-01-26 14:58:26 c

python的Greenlet模块源码分析

1. greenlet协程简介

协程，即协作式程序，其思想是，一系列互相依赖的协程间依次使用CPU，每次只有一个协程工作，而其他协程处于休眠状态。协程可以在运行期间的某个点上暂停执行，并在恢复运行时从暂停的点上继续执行。协程已经被证明是一种非常有用的程序组件，不仅被python、lua、ruby等脚本语言广泛采用，而且被新一代面向多核的编程语言如golang rust-lang等采用作为并发的基本单位。协程可以被认为是一种用户空间线程，与传统的线程相比，有2个主要的优点：

与线程不同，协程是自己主动让出CPU，并交付他期望的下一个协程运行，而不是在任何时候都有可能被系统调度打断。因此协程的使用更加清晰易懂，并且多数情况下不需要锁机制。与线程相比，协程的切换由程序控制，发生在用户空间而非内核空间，因此切换的代价非常小。

协程可以认为是一种用户态的线程，与系统提供的线程不同点是，它需要主动让出CPU时间，而不是由系统进行调度，即控制权在程序员手上。既然看成是用户态线程，那必然要求程序员自己进行各个协程的调度，这样就必须提供一种机制供编写协程的人将当前协程挂起，即保存协程运行场景的一些数据，调度器在其他协程挂起时再将此协程运行场景的数据恢复，以便继续运行。这里我们将协程运行场景的数据称为上下文。

一个 greenlet 是一个很小的独立微线程。可以把它想像成一个堆栈帧，栈底是初始调用，而栈顶是当前greenlet的暂停位置。你使用greenlet创建一堆这样的堆栈，然后在他们之间跳转执行。跳转不是绝对的：一个greenlet必须选择跳转到选择好的另一个greenlet，这会让前一个挂起，而后一个恢复。两个greenlet之间的跳转称为切换(switch) 。

greenlet模块源码中的协程的堆栈:

当你创建一个greenlet，它得到一个初始化过的空堆栈；当你第一次切换到它，他会启动指定的函数，然后切换跳出greenlet。当最终栈底函数结束时，greenlet的堆栈又变成空的了，而greenlet也就死掉了。greenlet也会因为一个未捕捉的异常死掉。

greenlet又叫协程，是一个伪线程，实际上是串行执行的，其实greenlet不是一种真正的并发机制，而是在同一线程内，在不同函数的执行代码块之间切换，实施“你运行一会、我运行一会”，并且在进行切换时必须指定何时切换以及切换到哪。greenlet的接口是比较简单易用的，但是使用greenlet时的思考方式与其他并发方案存在一定区别。

greenlet本质是一种合理安排了的串行，greenlet能够得到比较好的性能表现，主要也是因为通过合理的代码执行流程切换，完全避免了死锁和阻塞等情况。因为greenlet本质是串行，因此在没有进行显式切换时，代码的其他部分是无法被执行到的，如果要避免代码长时间占用运算资源造成程序假死，那么还是要将greenlet与线程/进程机制结合使用（每个线程、进程下都可以建立多个greenlet，但是跨线程/进程时greenlet之间无法切换或通讯）。

2. C的栈帧

ia32程序用程序栈来支持过程调用。机器用栈来传递过程参数、存储返回信息、保存寄存器用于以后恢复，以及本地存储。为单个过程分配的那部分栈称为栈帧。下图描绘了linux下栈帧的通用结构。栈帧的最顶端以两个指针界定，寄存器%ebp为帧指针，而寄存器%esp为栈指针。当程序执行时，栈指针可以移动，因此大多数信息的访问都是相对于帧指针的。

栈帧结构

图：栈帧结构

这里我们可以看到，在调用一个函数前，都会先将各个参数、调用者在被调用函数返回时执行的下一条指令的地址–返回地址压栈，被调用函数在开始前会将%ebp的值保存，然后将当前%esp的值赋予%ebp。弄明白帧指针和栈指针的作用，以及返回地址等如何通过%ebp来获取的话，对下面将要介绍到的内容很有帮助。

3. python的栈帧

由于Python本质上是解释性语言，虽然有字节码的编译过程，但最终解释的是字节码。Python的虚拟机依靠在内存中建立的数据结构，模拟了进程、线程、堆栈调用。这一套机制是Python的虚拟机实现的，不同于C的栈帧的实现。也就是说，C的栈帧间的切换，不等于Python栈帧的切换。但我们最终要切换的是Python代码的栈帧，但是我们在Python代码中没有办法可以切换栈帧，只有首先切换C的栈帧，然后再切换Python代码的EIP指针。Python中的EIP指针是Python对当前线程抽象的数据结构中的一个属性: top_frame，改变此属性，可以切换Python代码的执行流。但最终Python代码是依靠Python虚拟机运行的，而Python虚拟机是存在于C的栈帧上，所以还是需要依靠C的栈帧切换。

4. 新的greenlet启动流程

调用流程：g0->g1->g2，在g0中启动g1，然后在g1中启动g2，我们假设g0为main greenlet:

我们称需要切换到的greenlet为target greenlet, 而当前greenlet为current greenlet

流程如下：

4.1. 在g0中首先创建新的g1的greenlet对象，但只是在HEAP上创建这个对象，并设置了greenlet的run函数，还没有启动它。启动greenlet需要调用greenlet的switch方法。

4.2. g1的switch方法是C代码，所以调用g1的switch方法是在C的栈上建立switch过程的栈帧。switch方法方法中，会判断需要切换到的greenlet是否已经启动，如果没有启动，则调用g_initialstub方法启动目标greenlet(调用其run函数)。如果已经启动，则直接切换到目标greenlet:

while (target) {
    // 如果target已经启动，则调用target的g_switchstack()方法
    if (PyGreenlet_ACTIVE(target)) {
        ts_target = target;
        err = g_switchstack();
        break;
    }
    // 如果target没有启动，则调用target的g_initialstub()方法
    if (!PyGreenlet_STARTED(target)) {
        // dummymarker变量用户区分当前greenlet和目标greenlet的界限
        // 接下来会介绍这个变量
        void* dummymarker;
        ts_target = target;
        err = g_initialstub(&dummymarker);
        if (err == 1) {
            continue; /* retry the switch */
        }
        break;
    }
    target = target->parent;
}

4.3. 由于g1此时未启动，则我们会进入到g1的g_initialstub方法，此方法中会为尚未开始运行的g1，设置以下属性：

设置stack_stop为dummy_marker的地址。
设置g1的stack_prev为当前greenlet。

4.4. 接着执行EBP和ESP的切换，以及python的EIP的切换(top_frame)，将栈帧切换到taget greenlet运行:

/* perform the initial switch */
/* 执行栈帧的切换：保存当前python栈帧的信息，保存当前C栈帧的信息，并执行切换 */
/* 
    g_switchstack会调用slp_switch，
    对于从未运行过的greenlet，slp_switch中会直接返回1，不会执行C栈帧的切换，
    因为当g1的run函数执行完毕后，会切换回创建它的greenlet，也就是g0，
    所以没必要做C栈帧的切换，除非在g1的run函数中，用户代码执行了切换，
    否则会g1的run函数在执行完毕后，最终会切换回创建它的greenlet
    
    所以其实g_switchstack在判断slp_swtich返回1之后(slp_switch没有执行C栈帧切换)，
    会执行Python执行栈帧的切换
*/
err = g_switchstack();

/* 如果返回1，说明目标greenlet未执行，则执行g1的run函数 */
if (err == 1) {
    result = PyEval_CallObjectWithKeywords(run, args, kwargs);
    
    /* 如果执行完，切换回parent */
    /* jump back to parent */
    self->stack_start = NULL;  /* dead */
    for (parent = self->parent; parent != NULL; parent = parent->parent) {
        result = g_switch(parent, result, NULL);
    }
}

4.5. 在g_switchstack中，会直接调用嵌入GCC汇编的slp_switch函数。而在当目标greenlet未执行时，slp_switch会直接返回1，而不会执行C栈帧的切换：

int err;
{   /* save state */
    PyGreenlet* current = ts_current;
    /* 保存当前Python栈帧的数据 */
    PyThreadState* tstate = PyThreadState_GET();
    current->recursion_depth = tstate->recursion_depth;
    current->top_frame = tstate->frame; /* 这里保存了Python的EIP */
    current->exc_type = tstate->exc_type;
    current->exc_value = tstate->exc_value;
    current->exc_traceback = tstate->exc_traceback;
}

/*
    切换到目标greenlet
    其实当前通过过程的调用，控制流其实已经在target greenlet中
    只是帧指针和栈指针还需要切换
*/
err = slp_switch();

/* 如果返回值小于0，说明出错 */
if (err < 0) {   /* error */
    PyGreenlet* current = ts_current;
    current->top_frame = NULL;
    current->exc_type = NULL;
    current->exc_value = NULL;
    current->exc_traceback = NULL;

    assert(ts_origin == NULL);
    ts_target = NULL;
}

/* 
    等于0，或者大于0的情况，
    等于0时，说明目标greenlet已经在运行中，且已经执行了C栈帧的切换到目标greenlet
    等于1时，说明目标greenlet从未运行，slp_switch会直接返回1，不会执行C栈帧切换
    
    这两种情况都会执行Python栈帧的切换，这是必须的
*/
else {
    PyGreenlet* target = ts_target;
    PyGreenlet* origin = ts_current;
    PyThreadState* tstate = PyThreadState_GET();
    tstate->recursion_depth = target->recursion_depth;
    /* Python栈帧的切换 */
    tstate->frame = target->top_frame;
    target->top_frame = NULL;
    tstate->exc_type = target->exc_type;
    target->exc_type = NULL;
    tstate->exc_value = target->exc_value;
    target->exc_value = NULL;
    tstate->exc_traceback = target->exc_traceback;
    target->exc_traceback = NULL;

    assert(ts_origin == NULL);
    Py_INCREF(target);
    ts_current = target;
    ts_origin = origin;
    ts_target = NULL;
}
return err;

4.6. slp_switch函数是含有GCC嵌入汇编的C代码，它执行真正的C栈帧切换：

static int
slp_switch(void)
{
    void* ebp;
    void* ebx;
    unsigned int csr;
    unsigned short cw;
    register int *stackref, stsizediff;
    
    /* 保存寄存器到本地变量ebp,esp */
    __asm__ volatile ("" : : : REGS_TO_SAVE);
    __asm__ volatile ("fstcw %0" : "=m" (cw));
    __asm__ volatile ("stmxcsr %0" : "=m" (csr));
    __asm__ volatile ("movl %%ebp, %0" : "=m" (ebp));
    __asm__ volatile ("movl %%ebx, %0" : "=m" (ebx));
    
    /* 将当前函数的esp保存到stackref */
    __asm__ ("movl %%esp, %0" : "=g" (stackref));
    
    {
        /*
        切换目标栈帧之前从当前g开始，按照stack_pref的关系
        保存每个g的数据
        */
        SLP_SAVE_STATE(stackref, stsizediff);
        
        
        /*
         SLP_SAVE_STATE展开后变成下面的代码
        */
        
        stackref += STACK_MAGIC;
        // 保存当前greenlet和之前的greenlet，出错返回-1
        if (slp_save_state((char*)stackref)) return -1;
        
        // 如果目标greenlet没有启动，返回1到g_initialstub中继续执行，不会执行C栈帧切换
        // 对于从未执行的greenlet，到此处会函数会直接返回1，不会执行下面的C栈帧切换
        // 所以我们先不解析下面代码
        if (!PyGreenlet_ACTIVE(ts_target)) return 1;
        
        // 如果保存greenlet的栈没有出错，并且目标greenlet已经启动
        // 则stsizediff
        stsizediff = ts_target->stack_start - (char*)stackref
        
        
        /* 将ebp和esp指向目标栈，从这里开始从目标栈开始运行 */
        /* 将ebp和esp都设置为stsizediff */
        __asm__ volatile (
            "addl %0, %%esp\n"
            "addl %0, %%ebp\n"
            :
            : "r" (stsizediff)
            );

        /* 切换到新的greenlet之后，需要恢复之前保存的栈上的数据 */
        SLP_RESTORE_STATE();
    }

    /* 切换到新的greenlet之后，需要恢复之前保存的寄存器 */
    __asm__ volatile ("movl %0, %%ebx" : : "m" (ebx));
    __asm__ volatile ("movl %0, %%ebp" : : "m" (ebp));
    __asm__ volatile ("ldmxcsr %0" : : "m" (csr));
    __asm__ volatile ("fldcw %0" : : "m" (cw));
    __asm__ volatile ("" : : : REGS_TO_SAVE);
    return 0;
}

4.7. slp_switch函数中执行slp_save_state，slp_save_state会调用g_save函数。会将存在与C栈上的，并且位于当前greenlet之前的所有greenlet的C栈帧保存，用于当在目标greenlet执行中或者执行完毕后，C栈帧切换回来：

static int GREENLET_NOINLINE(slp_save_state)(char* stackref)
{
    /* stackref是调用者(slp_switch函数)的栈顶(ESP)的地址 */
    /* must free all the C stack up to target_stop */
    char* target_stop = ts_target->stack_stop;  /* 获取目标greenlet的stack_stop */
    
    PyGreenlet* owner = ts_current;
    
    assert(owner->stack_saved == 0);
    
    /* 如果owner是main greenlet，这个不会是NULL，在创建时已经设置为1*/
    if (owner->stack_start == NULL)
        owner = owner->stack_prev;  /* not saved if dying */
    else
        /* 设置当前greenlet的stack_start为slp_switch的栈顶 */
        owner->stack_start = stackref;
    
        
    /* 循环，在当前greenlet的stack_stop小于目标greenlet的stack_stop的条件下 */
    while (owner->stack_stop < target_stop) {
        /* ts_current is entierely within the area to free */
        /*
        owner是一个greenlet
        将target_stop之前的每个greenlet的stack_stop到stack_stop部分保存到HEAP
        */
        if (g_save(owner, owner->stack_stop))   
            /* 如果g_save返回-1，则当前函数也返回-1，用于错误判断 */
            return -1;  /* XXX */
            
        /* 将owner指向当前greenlet的前一个 */
        owner = owner->stack_prev;
    }
    if (owner != ts_target) {
        if (g_save(owner, target_stop))
            return -1;  /* XXX */
    }
    return 0;
}

4.8. g_save函数中，根据当前greenlet之前设置的stack_stop和当前的ESP，计算出需要保存到HEAP上的大小，并将这些数据保存到HEAP上：

static int g_save(PyGreenlet* g, char* stop)
{
    /* g是一个greenlet，stop是它的stack_stop */
    
    /* Save more of g's stack into the heap -- at least up to 'stop'

    // 获取当前greenlet已经在堆上保存的数据大小
    intptr_t sz1 = g->stack_saved;

    // 用当前greenlet的stack_stop - ESP(start_start)
    // 得出需要将当前greenlet在栈上的, 并且需要保存到HEAP上的数据的大小
    // stop和g->stack_start都是指针，指针相减得出中间指定类型的元素个数
    // sz2指示当前greenlet实际的栈大小
    intptr_t sz2 = stop - g->stack_start;
    
    // 断言当前greenlet的stack_start(ESP)不为NULL
    assert(g->stack_start != NULL);

    // 如果当前greenlet需要保存的数据大小，大于已经保存到HEAP上的数据大小
    // 说明当前greenlet自从上次运行到本次switch到别的greenlet之前
    // 栈又增长了，说白了就是当前greenlet的栈又有新的数据
    // 需要在HEAP上扩大空间，用来保存这些栈上新增的数据
    if (sz2 > sz1) {
        // 扩大当前greenlet在HEAP上的空间
        // 将堆上的内存空间扩大至sz2的大小
        char* c = (char*)PyMem_Realloc(g->stack_copy, sz2);
        if (!c) {
            PyErr_NoMemory();
            return -1;
        }

        // sz1是当前HEAP上已经保存的数据大小
        // c是realloc返回的指针，它等于stack_copy
        // 所以c+sz1就等于stack_saved
        // 但stack_copy已经为容纳新的数据做好准备
        //
        // g->stack_start是greenlet当前的ESP
        //
        // 复制的长度是greenlet栈的总长度 - 已经保存(saved)在HEAP上的长度
        // 等于当前greenlet的栈相对于上一次切换，增长的长度
        memcpy(c+sz1, g->stack_start+sz1, sz2-sz1);
        g->stack_copy = c;
        g->stack_saved = sz2;
    }
    return 0;
}

如下图：

这样当前greenlet和其之前的所有greenlet的C栈上的数据，都会保存到HEAP上。当然由于g1从未运行过，所以也就是说，当从g0切换到g1时，只需要保存g0本身的数据到HEAP即可。

4.9. 目前C的栈上保存着g0,g1,g2,g3的栈帧，如下图：

g0,g1,g2,g3的栈帧

图：g0,g1,g2,g3的栈帧

注意：每一个greenlet和前一个greenlet之间，在界限上有重叠的地方，从dummy_marker到slp_switch。所以在保存greenlet的栈帧时，会和相邻的greenlet保存重复的栈数据，而正是这样才实现了所谓的’fork堆栈’。作用是，当从g2切换回g0时，g0可以沿着slp_switch函数的执行流继续执行。这和g_initialstub中的注释不谋而合：正是在g_initialstub函数中，target greenlet和currnet greenlet分道扬镳，当从target greenlet切换回currnet greenlet时，currnet greenlet才能沿着之前的执行流继续执行。这个方法也同样适用于target greenlet：当在target greenlet创建新的greenlet，而新的greenlet需要切换回target greenlet时。

注意以下几点：

任何时候只有一个greenlet在运行
在运行的greenlet一定是在栈顶

5. 切换到被挂起的greenlet

调用流程：g0->g1->g2->g0，在g0中启动g1，然后在g1中启动g2，在g2中切换到g0，我们假设g0为main greenlet

5.1. 由于从g0->g1->g2，这部分流程已经在上面介绍过，下面介绍的是从g2切换到g0的流程

5.2. 在g_switch函数中，如果判断目标greenlet之前运行过，但目前是被挂起的状态，则会直接执行g_switchstack函数:

if (PyGreenlet_ACTIVE(target)) {
    ts_target = target;
    err = g_switchstack();
    break;
}

5.3. g_switchstack函数中，会调用slp_switch函数，它执行真正的C栈帧切换：

static int
slp_switch(void)
{
    void* ebp;
    void* ebx;
    unsigned int csr;
    unsigned short cw;
    register int *stackref, stsizediff;
    /* 保存寄存器到本地变量ebp,esp */
    __asm__ volatile ("" : : : REGS_TO_SAVE);
    __asm__ volatile ("fstcw %0" : "=m" (cw));
    __asm__ volatile ("stmxcsr %0" : "=m" (csr));
    __asm__ volatile ("movl %%ebp, %0" : "=m" (ebp));
    __asm__ volatile ("movl %%ebx, %0" : "=m" (ebx));
    
    /* 将当前函数的esp保存到stackref */
    __asm__ ("movl %%esp, %0" : "=g" (stackref));
    {
        /*
        切换目标栈帧之前从当前g开始，按照stack_pref的关系
        保存每个g的数据
        */
        SLP_SAVE_STATE(stackref, stsizediff);
        
        
        /*
         SLP_SAVE_STATE展开后变成下面的代码
        */
        
        stackref += STACK_MAGIC;
        // 保存当前greenlet和之前的greenlet，出错返回-1
        if (slp_save_state((char*)stackref)) return -1;
        
        // 如果目标greenlet没有启动，返回1到g_initialstub中继续执行，不会执行栈帧切换
        if (!PyGreenlet_ACTIVE(ts_target)) return 1;
        
        // 如果保存greenlet的栈没有出错，并且目标greenlet已经启动
        // 则stsizediff
        stsizediff = ts_target->stack_start - (char*)stackref
        
        
        /* 将ebp和esp指向目标栈，从这里开始从目标栈开始运行 */
        /* 将ebp和esp都设置为stsizediff */
        __asm__ volatile (
            "addl %0, %%esp\n"
            "addl %0, %%ebp\n"
            :
            : "r" (stsizediff)
            );

        /* 切换到新的greenlet之后，需要恢复之前保存的栈上的数据 */
        SLP_RESTORE_STATE();
    }

    /* 切换到新的greenlet之后，需要恢复之前保存的寄存器 */
    __asm__ volatile ("movl %0, %%ebx" : : "m" (ebx));
    __asm__ volatile ("movl %0, %%ebp" : : "m" (ebp));
    __asm__ volatile ("ldmxcsr %0" : : "m" (csr));
    __asm__ volatile ("fldcw %0" : : "m" (cw));
    __asm__ volatile ("" : : : REGS_TO_SAVE);
    return 0;
}

在开始执行C栈帧切换时，保存了EBP,EBX,ESP和其他寄存器的值到本地变量，用于将栈帧从target greenlet切换回来之后，恢复相关寄存器。在执行下面的语句后，EBP和ESP切换到target greenlet执行，目标greenlet的执行流还是从这里继续执行，但是当前的greenlet的执行流停止在这里：

/* 将ebp和esp指向目标栈，从这里开始从目标栈开始运行，而当前greenlet的执行流在这里停止 */
/* 将ebp和esp都设置为stsizediff */
__asm__ volatile (
    "addl %0, %%esp\n"
    "addl %0, %%ebp\n"
    :
    : "r" (stsizediff)
);

所以当从目标greenlet切换回来之后，之前被挂起的执行流，在这里得以继续执行:

/* 切换到目标greenlet之后，需要恢复目标greenlet之前保存在HEAP上的数据 */
SLP_RESTORE_STATE();

/* 切换到新的greenlet之后，需要恢复之前保存的各个寄存器的值 */
__asm__ volatile ("movl %0, %%ebx" : : "m" (ebx));
__asm__ volatile ("movl %0, %%ebp" : : "m" (ebp));
__asm__ volatile ("ldmxcsr %0" : : "m" (csr));
__asm__ volatile ("fldcw %0" : : "m" (cw));
__asm__ volatile ("" : : : REGS_TO_SAVE);
return 0;

最后的返回值为0，执行流从target greenlet又切换回current greenlet，这也是一次完整的greenlet之间切换的过程。

在从target切换回current时，由于target greenlet的栈帧是位于current greenlet之上的。在执行current.switch()时，由于按照C的过程调用规则，就是说直接调用current.switch，是在当前greenlet(target greenlet)的栈帧上建立current.switch的调用，而这个current.switch的栈帧，不是之前执行流被挂起的那个current.switch。在这个新的current.switch栈帧上，才会执行栈帧的切换到旧的，且执行流被停止的那个current.switch上。而当前的target greenlet，上由于调用了current.switch，在其栈帧上存在一个currenet.switch的调用轨迹。这是无法避免的，是C的过程调用规则。

5.4. 在调用slp_switch的g_switchstack函数中，恢复目标greenlet的Python栈帧：

PyGreenlet* target = ts_target;
PyGreenlet* origin = ts_current;
PyThreadState* tstate = PyThreadState_GET();
tstate->recursion_depth = target->recursion_depth;
tstate->frame = target->top_frame;  /* 设置Python的栈顶 */
target->top_frame = NULL;
tstate->exc_type = target->exc_type;
target->exc_type = NULL;
tstate->exc_value = target->exc_value;
target->exc_value = NULL;
tstate->exc_traceback = target->exc_traceback;
target->exc_traceback = NULL;

assert(ts_origin == NULL);
Py_INCREF(target);
ts_current = target;
ts_origin = origin;
ts_target = NULL;

6. 总结

greenlet的本质是C栈上的一些过程调用，当切换到其他greenlet时，保存一部分或者全部数据到HEAP上。当切换回来时，再从HEAP上复制数据到栈上，释放HEAP上申请的内存。 greenlet之间的切换主要依靠EBP、ESP寄存器和Python自身堆栈的EIP指针(top_frame)来实现python代码执行流的转移。