Python協程greenlet實現原理

greenlet是[stackless
Python](https://wiki.python.org/moin/StacklessPython)中剝離出來的一個專案，可以作為官方CPython的一個擴充套件來使用，從而支援Python協程。gevent正是基於greenlet實現。

協程實現原理

實現協程主要是在協程切換時，將協程當前的執行上下文儲存到協程關聯的context中。在c/c++這種native程式中實現協程，需要將棧內容和CPU各個暫存器的內容儲存起來。在Python這種VM中則有些不同。例如，在以下基於greenlet協程的python程式中：

def foo():
    bar()

def bar():
    a = 3 + 1
    gr2.switch()

def func():
    pass

gr1 = greenlet(foo)
gr2 = greenlet(func)
gr1.switch()

在bar中gr2.switch切換到gr2時，協程庫需要儲存gr1協程的執行上下文。這個上下文包括:

• Python VM的stack
• Python VM中解釋執行的上下文

理解以上兩點非常重要，至於為什麼呢？想象一下如何去實現一個Python
VM，去解釋執行一段Python程式碼。其實這在任何基於VM的語言中，原理都是一樣的（native程式可以把x86物理CPU也視作特殊的VM)。可以參考Python直譯器簡介-深入主迴圈。主要包含兩方面內容：

• VM在執行程式碼時，其自身呼叫棧通常都是遞迴的
• VM在執行程式碼時，通常會建立相應的資料結構來表示程式碼執行塊，例如通常會有個struct Frame來表示一個函式

在VM的實現中通常會有類似以下的程式碼:

struct Frame {
    unsigned char *codes; // 存放程式碼指令
    size_t pc; // 當前執行的指令位置
    int *stack; // stack-based的VM會有一個棧用於存放指令運算元
};

void op_call(frame) {
    switch (OP_CODE()) {
        case OP_CALL:
            child_frame = new_frame()
            op_call(child_frame)
                ...
        case OP_ADD:
            op_add(...)
    }
}

對應到前面的Python例子程式碼，在某一時刻VM的call stack可能是這樣的：

op_add
op_call
op_call

理解了以上內容後，就可以推測出greenlet本質上也是做了以上兩件事。

greenlet實現原理

greenlet庫中每一個協程稱為一個greenlet。greenlet都有一個棧空間，如下圖：

圖中未表達出來的，greenlet的棧空間地址可能是重疊的。對於活躍的（當前正在執行）的greenlet，其棧內容必然在c程式棧頂。而不活躍的被切走的greenlet，其棧內容會被copy到新分配的堆記憶體中。greenlet的棧空間是動態的，其起始地址是固定的，但棧頂地址不固定。以下程式碼展示一個greenlet的棧空間如何確定：

579         if (!PyGreenlet_STARTED(target)) { // greenlet未啟動，是一個需要新建立的greenlet
580             void* dummymarker; // 該區域性變數的地址成為新的greenlet的棧底
581             ts_target = target;
582             err = g_initialstub(&dummymarker); // 建立該greenlet並執行

以上greenlet->stack_stop確定了棧底，而棧頂則是動態的，在切換到其他greenlet前，對當前greenlet進行上下文的儲存時，獲取當前的RSP（程式實際執行的棧頂地址）：

410 static int GREENLET_NOINLINE(slp_save_state)(char* stackref)
411 {
412     /* must free all the C stack up to target_stop */
413     char* target_stop = ts_target->stack_stop;
414     PyGreenlet* owner = ts_current;
415     assert(owner->stack_saved == 0);
416     if (owner->stack_start == NULL)
417         owner = owner->stack_prev;  /* not saved if dying */
418     else
419         owner->stack_start = stackref; // stack_start指向棧頂

stackref是通過彙編獲取當前RSP暫存器的值：

     __asm__ ("movl %%esp, %0" : "=g" (stackref));

儲存棧內容到堆記憶體參看g_save的實現，沒什麼特別的。除了儲存棧內容外，如上一節講的，還需要儲存VM執行函式所對應的Frame物件，這個在g_switchstack中體現：

460         PyThreadState* tstate = PyThreadState_GET(); // 獲取當前執行緒的VM執行上下文
461         current->recursion_depth = tstate->recursion_depth;
462         current->top_frame = tstate->frame; // 儲存當前正在執行的frame到當前正在執行的greenlet
    ...
473         slp_switch(); // 做棧切換
    ...
487         PyThreadState* tstate = PyThreadState_GET();
488         tstate->recursion_depth = target->recursion_depth;
489         tstate->frame = target->top_frame; // 切換回來

上面的程式碼展示VM frame的切換。接下來看下最複雜的部分，當切換到目標greenlet時，如何恢復目標greenlet的執行上下文，這裡主要就是恢復目標greenlet的棧空間。假設有如下greenlet應用程式碼：

def test1():
    gr2.switch()

def test2():
    print(`test2`)

gr1 = greenlet(test1)
gr2 = greenlet(test2)
gr1.switch()

在gr1中切換到gr2時，也就是gr2.switch，會發生什麼事情。

// g_switch 實現
574         if (PyGreenlet_ACTIVE(target)) {
575             ts_target = target; // 找到目標greenlet，也就是gr2
576             err = g_switchstack(); // 開始切換

// g_switchstack 實現
462         current->top_frame = tstate->frame;
    ...
473     err = slp_switch();

// slp_switch 實現，根據不同平臺實現方式不同，原理相同
69         SLP_SAVE_STATE(stackref, stsizediff);
// 這個很重要，強行將當前的棧指標ESP/EBP (32位OS)通過加上一個與目標greenlet棧地址的偏移，而回到了
// 目標greenlet的棧空間。可以在下文看到stsizediff的獲取實現
70         __asm__ volatile (
71             "addl %0, %%esp
"
72             "addl %0, %%ebp
"
73             :
74             : "r" (stsizediff)
75             );
76         SLP_RESTORE_STATE();

// SLP_SAVE_STATE 實現
316 #define SLP_SAVE_STATE(stackref, stsizediff)            
317     stackref += STACK_MAGIC;                        
318     if (slp_save_state((char*)stackref)) return -1; 
319     if (!PyGreenlet_ACTIVE(ts_target)) return 1;    
// 獲取目標greenlet的棧空間與當前棧地址的偏移，用於稍後設定當前棧地址回目標greenlet的棧地址
320     stsizediff = ts_target->stack_start - (char*)stackref 

// slp_save_state 沒啥看的，前面也提過了，主要就是複製當前greenlet棧內容到堆記憶體

// SLP_RESTORE_STATE 也沒什麼看的，主要就是把greenlet堆記憶體複製回棧空間

以上，首先將ESP/EBP的值改回目標greenlet當初切換走時的ESP/EBP值，然後再把greenlet的棧空間記憶體（存放於堆記憶體中）全部複製回來，就實現了greenlet棧的回切。尤其注意的是，這個棧中是儲存了各種函式的return地址的，所以當slp_switch返回時，就完全恢復到了目標greenlet當初被切走時棧上的內容，包括各種函式呼叫棧。而當前greenlet的棧，則停留在了類似以下的函式呼叫棧：

g_switchstack
g_switch
...

參考

• Stackless Python 探祕
• python協程的實現（greenlet原始碼分析)
• 深度分析gevent執行流程