如何将线程池引用传入C++20协程co_await的await_suspend函数

要在await_suspend中使用线程池，核心是让awaiter结构体能够访问到线程池的引用。这里提供两种直接可行的实现方式：

方法一：直接传递线程池引用给Awaiter #

这种方法最直观，通过修改operator co_await的参数，让它同时接收std::future<int>和线程池的引用，再让awaiter持有这个引用。

步骤1：修改Awaiter结构体 #

让awaiter保存线程池的引用，并在await_suspend中用它提交任务：

// 假设你的线程池类名为ThreadPool，需提前定义
template<typename ThreadPool>
struct PoolAwaiter : std::future<int> {
    ThreadPool& pool;

    // 构造函数传入future和线程池引用
    PoolAwaiter(std::future<int> fut, ThreadPool& p)
        : std::future<int>(std::move(fut)), pool(p) {}

    bool await_ready() { return false; } // 始终挂起协程

    void await_suspend(std::coroutine_handle<> handle) {
        // 将等待future+恢复协程的任务提交到线程池
        pool.enqueue([this, handle]() mutable {
            this->wait(); // 在线程池线程中等待future完成
            handle.resume(); // 完成后恢复协程
        });
    }

    int await_resume() { return this->get(); } // 返回future的结果
};

步骤2：重载operator co_await #

添加一个接受std::future<int>和线程池引用的重载版本：

template<typename ThreadPool>
auto operator co_await(std::pair<std::future<int>, ThreadPool&> task_with_pool) {
    return PoolAwaiter<ThreadPool>(
        std::move(task_with_pool.first),
        task_with_pool.second
    );
}

// 可选：封装成更易读的辅助函数
template<typename ThreadPool>
auto with_pool(std::future<int> fut, ThreadPool& pool) {
    return std::make_pair(std::move(fut), pool);
}

步骤3：在协程中使用 #

调用co_await时传入future和线程池：

// 假设已经初始化了ThreadPool pool;
int c = co_await with_pool(async_add(a, b), pool);

方法二：通过Promise类型共享线程池上下文 #

如果希望线程池成为协程的全局上下文（比如所有协程共享同一个线程池），可以把线程池的引用存在promise_type中，再通过coroutine_handle获取。

步骤1：修改Promise类型 #

在coroutine_traits的promise_type中添加线程池的静态引用（用thread_local避免线程间干扰）：

// 提前声明ThreadPool类
class ThreadPool;

template <>
struct std::coroutine_traits<std::future<int>> {
    struct promise_type : std::promise<int> {
        std::future<int> get_return_object() { return this->get_future(); }
        std::suspend_never initial_suspend() noexcept { return {}; }
        std::suspend_never final_suspend() noexcept { return {}; }
        void return_value(int value) { this->set_value(value); }
        void unhandled_exception() {
            this->set_exception(std::current_exception());
        }

        // 用thread_local存储当前线程绑定的线程池
        static inline thread_local ThreadPool* current_pool = nullptr;
    };
};

// 同理修改带int参数的coroutine_traits版本
template <>
struct std::coroutine_traits<std::future<int>, int> {
    struct promise_type : std::promise<int> {
        // 重复上述相同的promise_type成员...
        static inline thread_local ThreadPool* current_pool = nullptr;
    };
};

步骤2：在启动协程前设置线程池 #

创建协程前，把线程池的地址赋值给current_pool：

int runfib(int arg) {
    ThreadPool pool; // 假设这里初始化你的线程池
    // 设置当前线程的协程上下文线程池
    std::coroutine_traits<std::future<int>>::promise_type::current_pool = &pool;
    
    auto fut = test_async_fib();
    fut.wait();
    return 0;
}

步骤3：修改Awaiter的await_suspend #

在await_suspend中通过promise_type获取线程池：

auto operator co_await(std::future<int> future) {
    struct awaiter : std::future<int> {
        bool await_ready() { return false; }

        void await_suspend(std::coroutine_handle<> handle) {
            // 获取当前协程的promise_type中的线程池
            auto* pool = std::coroutine_traits<std::future<int>>::promise_type::current_pool;
            
            if (pool) {
                // 提交任务到线程池
                pool->enqueue([this, handle]() mutable {
                    this->wait();
                    handle.resume();
                });
            } else {
                //  fallback：没有线程池时直接在当前线程等待
                this->wait();
                handle.resume();
            }
        }

        int await_resume() { return this->get(); }
    };
    return awaiter{std::move(future)};
}

注意事项 #

1. 确保线程池的enqueue方法能接收无参数的可调用对象（如果你的线程池需要参数，调整lambda即可）。
2. awaiter的生命周期由协程帧管理，只要协程未销毁，awaiter就会存在，因此lambda中捕获this是安全的。
3. 方法一的灵活性更高，支持不同协程使用不同的线程池；方法二更适合全局共享线程池的场景。

内容的提问来源于stack exchange，提问作者Samuel Squire