C/C++实现协程API包括上下文切换

协程可以很方便在一个线程中实现多个函数间的切换执行，如果某个函数需要等待，则可以切换到其他函数，这可以很大程度提高一个线程的利用率，因此现在很流行；c++很多版本原生并没有提供协程，但是由于c++支持内联汇编，所以实现一套协程api并非难事，只要实现函数间的上下文切换，其他的就容易多了，上下文切换通常就是保存当前函数的寄存器，由于是主动切换所以需要保存的寄存器也就很少了，下面给出一种协程的实现方式：

coroutine.h

#ifndef M_COROUTINE_H
#define M_COROUTINE_H
//coroutine info
typedef struct
{
    unsigned int id;
    int flag;
    int sleep_end_time;

    void* stack_start;
    void* stack_end;
    void* stack_top_p;

    void* parameters;
    void (*task)(void*);
}MCoroutine;

MCoroutine* alloc_co();
void append_co(MCoroutine* co);
MCoroutine* pop_co();
bool CreateCoroutine(void (*task)(void*), void* args);
void Schedule();
void RunTask();

#define co_yield Schedule();
#define print(str,...) printf(str##"\n",__VA_ARGS__)

#endif

coroutine.cpp

#include <iostream>
#include <windows.h>
#include <list>
#include "coroutine.h"
//Define the size of the coroutine stack
#define COROUTINE_STACK_SIZE 1024*1024*1 // 1M stack
#define MAX_COROUTINE 1000 // max coroutine num

//Coroutine status flag
enum FLAGS {
    COROUTINE_CREATE = 0x1,
    COROUTINE_READY = 0x2,
    COROUTINE_EXIT = 0x3,
    COROUTINE_MAIN = 0x9
};

std::list<MCoroutine *> GlobalCoroutineList; // Global Coroutine List
static int CoroutineRunningCount = 0; // the Running count of Coroutine
static MCoroutine cor_main = {0, COROUTINE_MAIN}; //main thread coroutine
static MCoroutine *CurrentCoroutine = &cor_main; //the current running coroutine

//Assign a coroutine
MCoroutine *alloc_co() {
    if (GlobalCoroutineList.size() > MAX_COROUTINE) {
        return nullptr;
    }
    auto *co = new MCoroutine();
    GlobalCoroutineList.push_back(co);
    return co;
}

// append a coroutine to Global Coroutine List
void append_co(MCoroutine *co) {
    GlobalCoroutineList.push_back(co);
}

// pop a coroutine from Global Coroutine List
MCoroutine *pop_co() {
    if (GlobalCoroutineList.empty()) {
        return nullptr;
    }
    MCoroutine *co;
    co = GlobalCoroutineList.front();
    GlobalCoroutineList.pop_front();
    return co;
}

//This function will be run for the first time, and the task will be called by this function
static void coroutine_startup(MCoroutine *coroutine) {
    CoroutineRunningCount++;
    coroutine->task(coroutine->parameters);
    coroutine->flag = COROUTINE_EXIT;
    CoroutineRunningCount--;
    Schedule();
}

static void push_stack(unsigned int **stack_top_p_p, unsigned int val) {
    *stack_top_p_p -= 1;
    **stack_top_p_p = val;
}

static bool init_stack(MCoroutine *coroutine) {
    unsigned char *stack_pages;
    unsigned int *stack_top_p;
    stack_pages = (unsigned char *) VirtualAlloc(nullptr, COROUTINE_STACK_SIZE, MEM_COMMIT, PAGE_READWRITE);
    if (stack_pages == nullptr) return false;
    coroutine->stack_start = stack_pages + COROUTINE_STACK_SIZE;
    coroutine->stack_end = stack_pages;
    stack_top_p = (unsigned int *) coroutine->stack_start;
    push_stack(&stack_top_p, (unsigned int) coroutine);//
    push_stack(&stack_top_p, 0);//padding
    push_stack(&stack_top_p, (unsigned int) coroutine_startup);//
    push_stack(&stack_top_p, 1);//ebp
    push_stack(&stack_top_p, 2);
    push_stack(&stack_top_p, 3);
    push_stack(&stack_top_p, 4);
    push_stack(&stack_top_p, 5);
    push_stack(&stack_top_p, 6);
    push_stack(&stack_top_p, 8);//eax
    coroutine->stack_top_p = stack_top_p;
    coroutine->flag = COROUTINE_READY;
    return true;
}

static unsigned int ID_NUM = 1;

static unsigned int get_id() {
    return ID_NUM++;
}

bool CreateCoroutine(void (*task)(void *), void *args) {
    MCoroutine *coroutine = alloc_co();
    if (coroutine == nullptr) return false;
    coroutine->flag = COROUTINE_CREATE;
    coroutine->id = get_id();
    coroutine->task = task;
    coroutine->parameters = args;
    return init_stack(coroutine);
}

void __fastcall release_coroutine(MCoroutine *coroutine){
    if(coroutine->flag!=COROUTINE_EXIT) return;
    VirtualFree(coroutine->stack_end,COROUTINE_STACK_SIZE,MEM_DECOMMIT);
    delete coroutine;
}

__declspec(naked) void switch_context(MCoroutine *cur_coroutine, MCoroutine *dst_coroutine) {
    __asm {
    push ebp
    mov ebp, esp
    push edi
    push esi
    push ebx
    push ecx
    push edx
    push eax

    mov esi, cur_coroutine
    mov edi, dst_coroutine
    mov[esi + MCoroutine.stack_top_p], esp
    /// Classic thread switch, another coroutine resurrection
    mov esp,[edi + MCoroutine.stack_top_p]
    /// release coroutine memory
    mov ecx,esi
    call release_coroutine
    /// resume dst_coroutine register
    pop eax
    pop edx
    pop ecx
    pop ebx
    pop esi
    pop edi
    pop ebp
    ret
    }
}

void Schedule() {
    MCoroutine *src_coroutine;
    MCoroutine *dst_coroutine;
    dst_coroutine = pop_co();
    if (dst_coroutine == nullptr) {
        dst_coroutine = &cor_main;
    }
    switch(CurrentCoroutine->flag){
        case COROUTINE_CREATE:
        case COROUTINE_READY:
            append_co(CurrentCoroutine);
            break;
        case COROUTINE_EXIT:
        case COROUTINE_MAIN:
            break;
    }
    src_coroutine = CurrentCoroutine;
    CurrentCoroutine = dst_coroutine;
    switch_context(src_coroutine, dst_coroutine);
    //print("switch_context over: src_coroutine:%d,dst_coroutine:%d", src_coroutine->id, dst_coroutine->id);
}

void RunTask() {
    do {
        Schedule();
        //print("current coroutine num: %d", CoroutineRunningCount);
    } while (CoroutineRunningCount);
}

main.cpp

#include "coroutine.h"
#include <iostream>

void task(void* num) {
    print("Coroutine%d start...", (int)num);
    for (int i = 0; i < 1000; ++i) {
        print("Coroutine%d_%d", (int)num, i);
        co_yield
    }
}

int main() {
    print("start...");
    for (int i = 0; i < 10; i++)
    {
        CreateCoroutine(task,(void*)i);
    }
    RunTask();
    return 0;
}

CMakeLists.txt

cmake_minimum_required(VERSION 3.16)
project(MyThread)

set(CMAKE_CXX_STANDARD 14)

add_executable(test00 main.cpp coroutine.cpp coroutine.h)

运行结果：

start...
Coroutine0 start...
Coroutine0_0
Coroutine1 start...
Coroutine1_0
Coroutine2 start...
Coroutine2_0
Coroutine3 start...
Coroutine3_0
Coroutine4 start...
Coroutine4_0
Coroutine5 start...
Coroutine5_0
Coroutine6 start...
Coroutine6_0
Coroutine7 start...
Coroutine7_0
Coroutine8 start...
Coroutine8_0
Coroutine9 start...
Coroutine9_0
Coroutine0_1
Coroutine1_1
Coroutine2_1