Lab4c 实验解析

tinyCoroLab4c 实验解析

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

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

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

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

📖lab4c 任务参考实现

🧑‍💻Task #1 - 实现 wait_group

wait_group 相比 latch 仅多了一个可以添加计数的功能，但却不能使用 event 来简化实现了，因为 event 一旦被 set 后续的协程调用wait均不会陷入 suspend 状态，而 wait_group 在引用计数降至 0 后可以重新增加计数，相当于将 event 重置为 unset 状态，而 tinyCoro 的 event 不具备这个功能，因此需要自行实现逻辑，但实验者也可以将 event 实现为可以被 unset 这样 wait_group 又可以被 event 用来简化实现了。

首先看 wait_group 的声明：

class wait_group
{
public:
    struct awaiter
    {
        awaiter(context& ctx, wait_group& wg) noexcept : m_ctx(ctx), m_wg(wg) {}

        constexpr auto await_ready() noexcept -> bool { return false; }

        auto await_suspend(std::coroutine_handle<> handle) noexcept -> bool;

        auto await_resume() noexcept -> void;

        auto resume() noexcept -> void;

        context&                m_ctx; // 绑定的 context
        wait_group&             m_wg; // 绑定的 wait_group
        awaiter*                m_next{nullptr}; // suspend awaiter 链表的 next 指针
        std::coroutine_handle<> m_await_coro{nullptr}; // 待 resume 的协程句柄
    };

    explicit wait_group(int count = 0) noexcept : m_count(count) {}

    auto add(int count) noexcept -> void;

    auto done() noexcept -> void;

    auto wait() noexcept -> awaiter;

private:
    friend awaiter;
    std::atomic<int32_t>     m_count; // 保存计数
    std::atomic<awaiter_ptr> m_state; // 与 event 的 m_state 功能一致
};

下面分析 awaiter 的实现，对于await_ready其恒定返回 false，则await_suspend一定会被调用，当然实验者自己对await_ready的实现也可以不返回恒定值。

对于await_suspend其实现如下：

auto wait_group::awaiter::await_suspend(std::coroutine_handle<> handle) noexcept -> bool
{
    m_await_coro = handle; // 保存协程句柄
    m_ctx.register_wait(); // 增加引用计数
    while (true)
    {
        // 如果引用计数为 0 那么协程不用陷入 suspend 状态
        if (m_wg.m_count.load(std::memory_order_acquire) == 0)
        {
            return false;
        }

        // 下面的操作主要是当前 awaiter 将自身置为 wait_group 的链表头，
        // 然后将 next 指向原先的链表头
        auto head = m_wg.m_state.load(std::memory_order_acquire);
        m_next    = static_cast<awaiter*>(head);
        if (m_wg.m_state.compare_exchange_weak(
                head, static_cast<awaiter_ptr>(this), std::memory_order_acq_rel, std::memory_order_relaxed))
        {
            // 引用计数不为 0 且成功挂载至 wait_group 的 awaiter 链表
            return true;
        }
    }
}

上述代码需要注意的是 awaiter 对于 wait_group 的 m_state 即链表头的操作，由于整个系列操作是非原子的，因此用原子变量的 cas 来判断过程中是否有其他线程改变了状态，因为是循环操作，所以使用了compare_exchange_weak而不是compare_exchange_strong，这部分代码与 lab4a 中 event 的实现逻辑是相同的。

await_resume主要实现降低引用计数的逻辑：

auto wait_group::awaiter::await_resume() noexcept -> void
{
    m_ctx.unregister_wait();
}

另外 awaiter 还提供了resume函数用来封装恢复协程执行的逻辑：

auto wait_group::awaiter::resume() noexcept -> void
{
    m_ctx.submit_task(m_await_coro); // 将协程句柄提交到其绑定的 context 的任务队列
}

对于 wait_group 面向用户的接口，其实现如下：

auto wait_group::add(int count) noexcept -> void
{
    m_count.fetch_add(count, std::memory_order_acq_rel); // 增加计数
}

auto wait_group::done() noexcept -> void
{
    if (m_count.fetch_sub(1, std::memory_order_acq_rel) == 1)
    {
        // 如果计数降至 0，那么恢复所有的 suspend awaiter
        auto head = static_cast<awaiter*>(m_state.exchange(nullptr, std::memory_order_acq_rel));
        while (head != nullptr)
        {
            head->resume();
            head = head->m_next;
        }
    }
}

auto wait_group::wait() noexcept -> awaiter
{
    return awaiter{local_context(), *this};
}

上述的代码实验者在完成 lab4a 和 lab4b 应当相当熟悉，因为逻辑是一样的。

实验总结

• 通过实现 wait_group 完成了拓展版的 latch，不管功能差异多大，其本质实现是一样的