Lab2b 实验解析

tinyCoroLab2b 实验解析

⚠️tinyCoroLab 的实验强烈推荐实验者独自完成而非直接翻阅实验解析，否则这与读完题直接翻看参考答案无太大区别，实验解析仅供实验者参考。

本节将会以 tinyCoroLab 的官方实现tinyCoro为例，为大家分析并完成 lab2b，请实验者预先下载 tinyCoro 的代码到本地。

git clone https://github.com/sakurs2/tinyCoro

打开include/coro/context.hpp和src/context.cpp并大致浏览代码结构。

📖lab2b 任务参考实现

🧑‍💻Task #1 - 完善 context 初始化以及外部交互 API

context 的初始化与析构化比较简单：

auto context::init() noexcept -> void
{
    linfo.ctx = this;
    m_engine.init();
}
auto context::deinit() noexcept -> void
{
    linfo.ctx = nullptr;
    m_engine.deinit();
}

对于 context 的提交任务只需要调用 engine 的接口就可以了：

auto context::submit_task(std::coroutine_handle<> handle) noexcept -> void
{
    m_engine.submit_task(handle);
}

context 的引用计数可能涉及到多线程操作，因此使用原子变量存储计数，register_wait和unregister_wait方法只需要对计数增减即可。

atomic<size_t>      context::m_num_wait_task{0};
auto context::register_wait(int register_cnt = 1) noexcept -> void 
{
    m_num_wait_task.fetch_add(size_t(register_cnt), memory_order_acq_rel);
}
auto context::unregister_wait(int register_cnt = 1) noexcept -> void
{
    m_num_wait_task.fetch_sub(register_cnt, memory_order_acq_rel);
}

然后是通知工作线程停止的notify_stop，使用 jthread 提供的方法并注意调用wake_up，因为此时工作线程可能在阻塞态，代码如下：

auto context::notify_stop() noexcept -> void
{
    m_job->request_stop();
    m_engine.wake_up();
}

🧑‍💻Task #2 - 完善 context 核心任务循环

首先为 context 添加一些辅助函数比如执行任务、执行 IO 任务以及状态判断：

// 驱动 engine 从任务队列取出任务并执行
auto context::process_work() noexcept -> void
{
    auto num = m_engine.num_task_schedule();
    for (int i = 0; i < num; i++)
    {
        m_engine.exec_one_task();
    }
}
// 驱动 engine 执行 IO 任务
auto context::poll_work() noexcept -> void { m_engine.poll_submit(); }
// 判断是否没有 IO 任务以及引用计数是否为 0
auto context::empty_wait_task() noexcept -> bool
{
    return m_num_wait_task.load(memory_order_acquire) == 0 && m_engine.empty_io();
}

对于任务循环函数run，其核心目标便是在一遍遍的循环中驱动 engine 完成全部任务且优雅的退出，tinyCoro 给出的实现如下：

auto context::run(stop_token token) noexcept -> void
{
    while (true)
    {
        process_work();
        if (token.stop_requested() && empty_wait_task())
        {
            if (!m_engine.ready())
            {
                break;
            }
            else // 此时任务队列还有任务，优先执行任务队列里的任务
            {
                continue;
            }
        }
        poll_work();
        if (token.stop_requested() && empty_wait_task() && !m_engine.ready())
        {
            break;
        }
    }
}

tinyCoro 对run函数的实现并非标准实现，实验者只需要保证实现满足核心目标即可。

🧑‍💻Task #3 - 完善 scheduler 对 context 的管理能力

在正式讲解本节任务的具体实现前，我们可以明确一下核心问题，即scheduler 如何保证在所有 context 均完成任务后对外发送停止信号？ 这看起来需要依靠事件驱动来解决。

再具体化一下这个问题的难点，如何确保所有 context 均完成任务后向 scheduler 发送信号？

读者应该不难想到，能不能通过在 scheduler 内部维持一个引用计数，每当有一个 context 完成所有任务后便通知 scheduler 对该引用计数减 1，等到引用计数减至 0 时 scheduler 便通知所有 context 停止？这个方法核心思路可行，但每一个 context 在完成自身所有任务后均有可能收到 scheduler 派发的新任务， 所以单纯对引用计数减是不可行的，还要考虑引用计数的增加。

那么何时对引用计数增加呢？假设 scheduler 掌管多个 context 且除了 context A 之外其余全部完成任务，此时 scheduler 的引用计数为 1，但 context A 向 scheduler 派发了新任务，随后便请求降低引用计数，如果此时引用计数降至 0，那么 scheduler 会通知所有 context 停止，新任务便会被派发到已经停止了的 context，所以如何解决该问题？答案是在submit_to_scheduler这个用于派发新任务的函数中对 scheduler 引用计数加 1，这样 context A 派发完新任务并降低引用计数后 scheduler 的引用计数不为 0，就避免了将任务派发到已经停止的 context 中。

💡对程序的执行路径进行分析，在确定性的执行路径中对状态进行转移，这种思路在之后的 lab4 中也会有所体现

上述思路读者肯定会觉得存在诸多问题，因为这只是大致方案描述，具体落实需要补充很多细节，我们来看看 tinyCoro 是怎么实现的吧！

scheduler 内部新增了两个成员变量，代码如下：

// atomic_ref_wrapper 具体定义在 include/detail/atomic_helper.hpp 中
template<typename T>
struct alignas(config::kCacheLineSize) atomic_ref_wrapper
{
    // alignas 是 std::atomic_ref 要求的地址对齐
    alignas(std::atomic_ref<T>::required_alignment) T val;
    // 这里为何不直接用 std::atomic<T>呢？
    // 因为我们要在容器里存储多个该类型，由于 std::atomic<T>禁用了
    // 拷贝构造和拷贝赋值，所以直接用 std::atomic<T>会导致容器无法初始化
};

// 上文提到的引用计数
using stop_token_type = std::atomic<int>;

// 每个 context 对应的状态，只有 0 和 1 两个值，0 表示 context 已完成所有任务，1 表示 context 还在执行任务中
using stop_flag_type  = std::vector<detail::atomic_ref_wrapper<int>>;

stop_flag_type m_ctx_stop_flag; // 存储各个 context 的执行状态
stop_token_type m_stop_token; // 引用计数成员变量

这里为何要使用两个成员变量呢？稍后读者便能知晓。

对于 context，我们为其新增了一个成员变量以及方法：

using stop_cb = std::function<void()>;
stop_cb m_stop_cb;

auto set_stop_cb(stop_cb cb) noexcept -> void{
    m_stop_cb = cb;
}

m_stop_cb定义了 context 完成所有任务后应该执行的停止逻辑，set_stop_cb使得外部可以将停止逻辑注入到 context 中。此时我们需要将 context 的start函数修改为如下：

auto context::start() noexcept -> void
{
    m_job = make_unique<jthread>(
        [this](stop_token token)
        {
            this->init();
            // 如果外部没有注入 stop_cb，那么自行为其添加逻辑
            if (!(this->m_stop_cb))
            {
                m_stop_cb = [&]() { m_job->request_stop(); };
            }
            this->run(token);
            this->deinit();
        });
}

上述代码保证了 context 在独立运行的情况下也能正确执行停止逻辑，然后修改 context 的run函数：

auto context::run(stop_token token) noexcept -> void
{
    while (!token.stop_requested())
    {
        process_work();
        if (empty_wait_task())
        {
            if (!m_engine.ready())
            {
                // 此处表明 contetx 已执行完所有任务，那么调用停止逻辑
                m_stop_cb();
            }
            else
            {
                continue;
            }
        }

        poll_work();
    }
}

此时 context 侧的修改已结束了，我们回到 scheduler，在 init_impl 中为新增的两个成员变量进行初始化：

auto scheduler::init_impl(size_t ctx_cnt) noexcept -> void
{
    // other codes...
    m_ctx_stop_flag = stop_flag_type(m_ctx_cnt, detail::atomic_ref_wrapper<int>{.val = 1});
    m_stop_token    = m_ctx_cnt;
}

然后是核心的loop函数：

inline static auto loop() noexcept -> void { get_instance()->loop_impl(); }
auto scheduler::start_impl() noexcept -> void
{
    for (int i = 0; i < m_ctx_cnt; i++)
    {
        m_ctxs[i]->set_stop_cb(
            [&, i]()
            {
                // context 将其关联的状态设置为 0 即已执行完所有任务，cnt 总是为 1
                auto cnt = std::atomic_ref(this->m_ctx_stop_flag[i].val).fetch_and(0, memory_order_acq_rel);
                // 将 scheduler 的引用计数减 1，如果引用计数降至 0，那么触发 scheduler 发送停止信号
                if (this->m_stop_token.fetch_sub(cnt) == cnt)
                {
                    this->stop_impl();
                }
            });
        m_ctxs[i]->start();
    }
}
auto scheduler::stop_impl() noexcept -> void
{
    for (int i = 0; i < m_ctx_cnt; i++)
    {
        m_ctxs[i]->notify_stop();
    }
}

上述代码主要是负责降低引用计数，下面我们看怎么在submit_to_scheduler中增加引用计数：

// submit_to_scheduler 本质是调用该方法
auto scheduler::submit_task_impl(std::coroutine_handle<> handle) noexcept -> void
{
    // 不要在 scheduler::loop 结束后再添加新任务
    assert(this->m_stop_token.load(std::memory_order_acquire) != 0 && "error! submit task after scheduler loop finish");
    size_t ctx_id = m_dispatcher.dispatch();
    // 直接增加引用计数是不合理的，
    // 根据 context 的运行状态来增加引用计数，避免冗余增加
    m_stop_token.fetch_add(
        1 - std::atomic_ref(m_ctx_stop_flag[ctx_id].val).fetch_or(1, memory_order_acq_rel), memory_order_acq_rel);
    m_ctxs[ctx_id]->submit_task(handle);
}

读者不难发现，scheduler 发送停止信号的条件是所有 context 的运行状态均为 0，添加m_stop_token可以简化这个判断过程。

综上，通过注入回调函数并搭配原子变量，成功使得scheduler::loop达到预期的行为。

💡上述代码对引用计数的增加和降低涉及到多组原子操作，读者可以自行分析是否存在执行顺序会导致异常的行为

实验总结

• 通过完善 context 驱动 engine 执行任务使得 tinyCoro 正式拥有完整的对外提供服务的能力

• 巧妙利用原子和回调函数来增强 scheduler 对 context 的管理能力