【动手写协程库】系列笔记是学习sylar的协程库时的记录,参考了从零开始重写sylar C++高性能分布式服务器框架和代码随想录中的文档。文章并不是对所有代码的详细解释,而是为了自己理解一些片段所做的笔记。

Coroutine类中其他函数的定义可以在这里查看:Github: src/coroutine.cpp

对于什么是协程,为什么要使用协程,可以看看之前的笔记:【协程】C++20协程初体验

对于我们自己来实现协程,其实在之前Xv6的Lab中就有做过:【MIT6.S081】Lab6 multithreading,当初做这个lab的时候没有意识到这就是协程。协程的切换最重要的就是要保存和恢复上下文,在这个lab中,我们通过保存每个协程在切换之前的寄存器的值,以此可用来恢复原来的执行流。

在sylar的协程库实现中,使用的是Linux原生提供的ucontext来保存协程的上下文和切换。对于协程切换,最重要的两个API就是yieldresume,分别对应协程让出执行权和恢复协程。利用 ucontext 提供的四个函数getcontext()setcontext()makecontext()swapcontext()可以在一个进程中实现协程切换。(这里就不介绍这几个函数的用法来,详细文档可使用man命令查看)

Coroutine的相关API如下:

#ifndef COROUTINE_H_
#define COROUTINE_H_

#include <sys/ucontext.h>
#include <ucontext.h>
#include <cstdint>
#include <functional>
#include <memory>

class Coroutine : public std::enable_shared_from_this<Coroutine> {
public:
    using Ptr = std::shared_ptr<Coroutine>;

    enum State {
        READY,
        RUNNING,
        FINISH
    };

    Coroutine(std::function<void()> callback, size_t stack_size = 0, bool run_in_scheduler = true);
    ~Coroutine();

    void Yield();

    void Resume();

    void Reset(std::function<void()> callback);

    uint64_t GetId() const {
        return id_;
    }

    State GetState() const {
        return state_;
    }

public:
    static void SetNowCoroutine(Coroutine* co);

    static Coroutine::Ptr GetNowCoroutine();

    static uint64_t TotalCoNums();

    static void Task();

    static uint64_t GetCurrentId();

private:
    Coroutine();

private:
    uint64_t id_ = 0;
    uint32_t stack_size_ = 0;
    State state_ = READY;
    bool is_run_in_sched_;

    ucontext_t ctx_;         // 协程上下文
    void* pstack_ = nullptr; // 协程的栈地址

    std::function<void()> callback_;
};

#endif /* COROUTINE_H_ */

一个线程可以运行多个协程,但是在某一个时刻只能运行一个协程。我们需要为每个线程设置当前运行的协程的指针和表示线程的主协程的指针。这里用到了C++中的thread_local关键字:

static thread_local Coroutine* cur_coroutine = nullptr;
static thread_local Coroutine::Ptr main_coroutine = nullptr;

每当我们需要创建一个协程来执行任务时,我们必须要传入的参数为要执行的函数,而协程的栈大小有默认值,默认使用调度器进行调度:

Coroutine::Coroutine(std::function<void()> callback, size_t stack_size, bool run_in_scheduler)
    : callback_(callback)
    , is_run_in_sched_(run_in_scheduler) {
    co_count++;
    stack_size_ = stack_size > 0 ? stack_size : CO_STACK_SIZE;
    pstack_ = malloc(stack_size_);

    // getcontext()用于保存当前上下文,以便将来可以从这个点恢复执行
    if (getcontext(&ctx_) != 0) {
        std::cout << "err: Coroutine::getcontext\n";
        exit(1);
    }
    // 初始化协程上下文
    ctx_.uc_link = nullptr;
    ctx_.uc_stack.ss_sp = pstack_;
    ctx_.uc_stack.ss_size = stack_size_;
    // 为已经初始化的上下文设置一个将在该上下文被激活时执行的函数,并为该函数传递参数。
    // 为什么这里的第二个参数不直接设置成callback呢?是因为我们自己写协程的话不仅仅只要将任务函数执行完成就行了,执行完成后还要设置协程的状态
    makecontext(&ctx_, &Coroutine::Task, 0);
}

// 每个协程会运行它所绑定的callback,并且在执行完成后将重置该协程的状态,并让出执行权
void Coroutine::Task() {
    auto cur = GetNowCoroutine();
    assert(cur);

    cur->callback_();
    cur->callback_ = nullptr;
    cur->state_ = FINISH;

    auto raw_ptr = cur.get();
    cur.reset();
    raw_ptr->Yield();
}

YieldResume函数利用swapcontext来进行协程的切换。由于我们在之后会需要将协程添加到调度器中而不是手动调度,所以要注意协程有使用调度器标志时要与调度器协程进行切换:

// 让出该协程的执行权,转交到主协程
void Coroutine::Yield() {
    assert(state_ == FINISH || state_ == RUNNING);
    SetNowCoroutine(main_coroutine.get());
    if (state_ != FINISH) {
        state_ = READY;
    }
    if (is_run_in_sched_) {
        if (swapcontext(&ctx_, &(Scheduler::GetSchedCoroutine()->ctx_)) != 0) {
            std::cout << "err: Yield::swapcontext\n";
            assert(false);
        }
    } else {
        if (swapcontext(&ctx_, &(main_coroutine->ctx_)) < 0) {
            std::cout << "err: Yield::swapcontext\n";
            exit(1);
        }
    }
}

// 从当前运行的协程恢复到该协程
void Coroutine::Resume() {
    assert(state_ != FINISH && state_ != RUNNING);
    // 每次恢复时需要将当前运行的协程设置为自身
    SetNowCoroutine(this);
    state_ = RUNNING;
    if (is_run_in_sched_) {
        if (swapcontext(&(Scheduler::GetSchedCoroutine()->ctx_), &ctx_) != 0) {
            std::cout << "err: Resume::swapcontext\n";
            assert(false);
        }
    } else {
        if (swapcontext(&(main_coroutine->ctx_), &ctx_) != 0) {
            std::cout << "err: Resume::swapcontext\n";
            exit(1);
        }
    }
}