C++20 Coroutine Task

Question

This project is the natural extension to my attempt to make a templatedgenerator coroutine. This time, I tried what I called a "task" (may be not the real term for this) : it generates a value or not, and is recursively awaitable. It can call other coroutines which each pause will also pause the main coroutine. It was really difficult to implement recursion, and I wonder if there is a better option, like one bundled with coroutine support directly. My first attempt was to make a variable innerCoroutine but I think it suffers from performances issues because when there is, like 100 coroutines on the stack, each call will recursively try to found the most inner coroutine, and with this method each resume() call is 0(1). I tested with unit cases, so there was many trials until all tests passes so it may be a bit complicated. I wonder if is it okay, especially from the point of view of undefined behavior and memory leaks.

When I started learning C++ coroutines, I found them overly unnecessary complicated, but I think I begin to figure out they are not that complicated, and highly powerful and customizable, at least I hope.

Godbolt link

#include <coroutine>
#include <optional>
#include <cassert>
#include <memory>
#include <vector>
#include <iostream>
struct promise_type_vars
{
 using stack = std::vector<std::coroutine_handle<>>;
 std::shared_ptr<stack> stack_;
};
template<typename self>
struct promise_type_base : promise_type_vars
{
 using promise_type = typename self::promise_type;
 std::suspend_always initial_suspend() { return {}; }
 std::suspend_always final_suspend() noexcept { return {}; }
 [[noreturn]] void unhandled_exception() { throw; }
 self get_return_object() { return {std::coroutine_handle<promise_type>::from_promise(static_cast<promise_type&>(*this))}; }
};
template<typename T>
struct add_reference
{
 using type = T&;
};
template<>
struct add_reference<void>
{
 using type = void;
};
template<typename T>
using add_reference_t = typename add_reference<T>::type;
template<typename T = void>
class task
{
 static constexpr bool is_void = std::is_same_v<T, void>;
public:
 struct promise_type_nonvoid : promise_type_base<task>
 {
 std::optional<T> t_;
 void return_value(T t)
 {
 t_ = std::move(t);
 this->stack_->pop_back();
 }
 };
 struct promise_type_void : promise_type_base<task>
 {
 void return_void()
 {
 this->stack_->pop_back();
 }
 };
 struct promise_type : std::conditional_t<is_void, promise_type_void, promise_type_nonvoid>
 {
 using stack = promise_type_vars::stack;
 };
private:
 std::coroutine_handle<promise_type> h_;
public:
 task(std::coroutine_handle<promise_type> h) : h_(h)
 {
 h_.promise().stack_ = std::make_shared<typename promise_type::stack>();
 h_.promise().stack_->push_back(h_);
 }
 task(const task&) = delete;
 task& operator=(const task&) = delete;
 task(task&& other)
 {
 swap(h_, other.h_);
 }
 task& operator=(task&& other)
 {
 swap(h_, other.h_);
 return *this;
 }
 ~task()
 {
 if(h_)
 {
 h_.destroy();
 h_ = {};
 }
 }
 bool is_resumable() const { return h_ && !h_.done(); }
 bool operator()() { return resume(); }
 auto& stack() const { return *h_.promise().stack_; }
 bool resume()
 {
 assert(is_resumable());
 auto& s = *h_.promise().stack_;
 auto back = s.back();
 back.resume(); //execute suspended point
 // at this point, back and s.back() could be different
 if(back.done())
 {
 if(s.size() > 0)
 {
 // Not root, execute co_await return expression
 resume();
 }
 else
 {
 if constexpr(!is_void)
 {
 assert(h_.promise().t_ && "Should return a value to caller");
 }
 }
 }
 return is_resumable();
 }
 // Make co_await'able (allow recursion and inner task as part of the parent task) ------------------------
 bool await_ready() const
 {
 return false;
 }
 template<typename X>
 bool await_suspend(std::coroutine_handle<X>& p)
 {
 // --- stack push
 h_.promise().stack_ = p.promise().stack_;
 p.promise().stack_->push_back(h_);
 // ---
 h_(); // never wait an inner function whether initially or finally (initially managed here)
 if(!h_.done())
 {
 // the inner coroutine has at least one suspend point
 return true;
 }
 else
 {
 return false; // don't wait if the coroutine is already ended (coroutine is NOOP)
 }
 }
 T await_resume()
 {
 if constexpr(!is_void)
 {
 assert(h_.promise().t_ && "Should return a value in a co_wait expression");
 return get();
 }
 }
 add_reference_t<T> get()
 {
 if constexpr(!is_void)
 {
 return h_.promise().t_.value();
 }
 }
};
task<double> one()
{
 std::cout << "one takes a coffee" << std::endl;
 co_await std::suspend_always{};
 co_return 1;
}
task<> noop()
{
 std::cout << "the noop does something" << std::endl;
 co_return;
}
task<short> two()
{
 co_return co_await one() + co_await one();
}
task<int> foo()
{
 int i = 0;
 i += co_await one();
 std::cout << "foo just run one()" << std::endl;
 i += co_await two();
 std::cout << "foo just run two()" << std::endl;
 std::cout << "foo makes a pause after his work" << std::endl;
 co_await std::suspend_always{};
 co_await noop();
 co_return i;
}
int main()
{
 auto handle = foo();
 
 for(int i = 1; handle(); ++i) {
 std::cout << i << ". hi from main" << std::endl;
 }
 std::cout << "Result: " << handle.get() << std::endl;
 return 0;
}

Output :

one takes a coffee
1. hi from main
foo just run one()
one takes a coffee
2. hi from main
one takes a coffee
3. hi from main
foo just run two()
foo makes a pause after his work
4. hi from main
the noop does something
Result: 3

• Michel: I am late to this. However if you are still using this approach, it seems you have over-complicated the recursive call. The C++ coroutine system already allocates one promise per task and you do not need to maintain a stack of handles in the promise. You only need the co_awaiter's handle that you return in the final_suspend.
• Louis Langholtz: @Michel no one has answered this review question yet but it sounds like you'd have some helpful insight to provide in an answer. I'd love to see a code example for instance of what you mean in terms of only needing the co_awaiter's handle.

Answer 1

@Louis-Langholtz This has been a while, but I certainly can do that.

After running a first attempt at fixing @rafoo's code, I've understood the real source of his code's problem. It does a co_await std::suspend_always{}; in two places. This is not the proper usage of std::suspend_always. The problem is that this call returns control to main, but main needs to restart the proper handle, which, without his stack of handles, is lost.

So instead, one needs to have an awaitable that schedules the handle of the awaiting coroutine for future resumption. This is done below by the object suspend{}. The variable execution_handle models a simple executor, main just resumes that handle until it is done. Note however, that the handle will be changed by any coroutine that co_await suspend{};. The rest of the flow is handled by the task continuation handle, allowing to return execution to the coroutine that awaits the task's result.

In a real code, suspend{} could pass the handle to a real multi-threaded executor. Or it could have an asynchronous I/O completion do the re-scheduling once the I/O has completed.

#include <iostream>
#include <coroutine>
#include <memory>
#include <optional>
namespace coroutines = std;
template<typename T = void>
struct [[nodiscard]] task {
 using value_type = T;
 struct task_promise;
 using promise_type = task_promise;
 explicit task(std::coroutine_handle<promise_type> handle) : my_handle(handle) {}
 task(task &&t) noexcept: my_handle(t.my_handle) { t.my_handle = nullptr; }
 /// Disable copy construction/assignment.
 task(const task &) = delete;
 task &operator=(const task &) = delete;
 ~task() {
 if (my_handle)
 my_handle.destroy();
 }
 task &operator=(task &&other) noexcept {
 if (std::addressof(other) != this) {
 if (my_handle) {
 my_handle.destroy();
 }
 my_handle = other.my_handle;
 other.my_handle = nullptr;
 }
 return *this;
 }
 /** \brief Provides for co_awaiting a task. */
 auto operator
 co_await() noexcept
 {
 struct awaiter {
 bool await_ready() const noexcept { return !coro_handle || coro_handle.done(); }
 std::coroutine_handle<>
 await_suspend(std::coroutine_handle<> awaiting_coroutine) noexcept {
 coro_handle.promise().set_continuation(awaiting_coroutine);
 return coro_handle;
 }
 T await_resume() noexcept {
 if constexpr (std::is_same_v<T, void>)
 return;
 else
 return std::move(coro_handle.promise().data.value());
 }
 std::coroutine_handle<promise_type> coro_handle;
 };
 return awaiter{my_handle};
 }
 auto get_handle() {
 return my_handle;
 }
private:
 std::coroutine_handle<promise_type> my_handle;
};
/** \brief Component of the promise handling the result.
 *
 * @tparam T Result type
 *
 * \details The C++20 coroutines' promises need to expose either a return_void or
 * return_value function, whether they return an actual value or a void.
 * They cannot expose both. Many implementations specialize the full promise type.
 * Here we only specialize the return value handling.
 */
template<typename T>
struct promise_setter {
 template<typename U>
 void return_value(U &&u) {
 data.template emplace(std::forward<U>(u));
 }
 std::optional<T> data;
};
/** \brief Specialization of promise_setter for void returning tasks. */
template<>
struct promise_setter<void> {
 void return_void() {}
};
/** \brief Promise type associated with a task
 *
 * @tparam T the return type of the task.
 */
template<typename T>
struct task<T>::task_promise : public promise_setter<T> {
 /// \brief Invoked when we first enter a coroutine, build the task object with the coroutine
 /// handle.
 task get_return_object() noexcept {
 return task{std::coroutine_handle<task_promise>::from_promise(*this)};
 };
 auto initial_suspend() const noexcept {
 return std::suspend_always{};
 }
 /// \brief Get an awaiter with the continuation if any.
 /// \details If T is void, co_await final_suspend(); will destroy the coroutine frame.
 auto final_suspend() const noexcept {
 struct awaiter {
 // Return false here to return control to the thread's event loop. Remember that we're
 // running on some async thread at this point.
 /// \brief The current coroutine is done, so it is not ready to restart.
 /// \details await_suspend will be called next allowing to return the handle of the
 /// continuation.
 bool await_ready() const noexcept { return false; }
 void await_resume() const noexcept {}
 /** \brief Suspension method for the ending co-routine.
 * @param h Handle of the current and ending co-routine.
 * @return The handle of the coroutine that continues the work if any.
 */
 std::coroutine_handle<>
 await_suspend(std::coroutine_handle<task::task_promise> h) noexcept {
 return h.promise().continuation;
 }
 };
 return awaiter{};
 }
 void unhandled_exception() noexcept { std::terminate(); }
 void set_continuation(std::coroutine_handle<> handle) { continuation = handle; }
private:
 std::coroutine_handle<> continuation = std::noop_coroutine();
};
/* Simple mock-up of an executor. main() will continuously resume execution_handle.
 A coroutine can be re-scheduled for execution by writing its handle to this variable. */
std::coroutine_handle<> execution_handle;
/// An awaitable for suspending a coroutine and, in this demo, immediately
/// scheduling it to be resumed next.
struct suspend {
 /// \brief The current coroutine wants to pause and is not ready to restart.
 /// \details await_suspend will be called next allowing to schedule for future execution.
 bool await_ready() const noexcept { return false; }
 void await_resume() const noexcept {}
 /** \brief Suspension method for the ending co-routine.
 * @param h Handle of the current and ending co-routine.
 * @return The handle of the coroutine that continues the work if any.
 */
 auto
 await_suspend(std::coroutine_handle<> h) noexcept {
 // Schedule for resumption
 execution_handle = h;
 // Pause the execution
 return std::noop_coroutine();
 }
};
task<double> one()
{
 std::cout << "one takes a coffee" << std::endl;
 // Suspend (and reschedule)
 co_await suspend{};
 co_return 1;
}
task<> noop()
{
 std::cout << "the noop does something" << std::endl;
 co_return;
}
task<short> two()
{
 co_return co_await one() + co_await one();
}
task<int> foo()
{
 int i = 0;
 i += co_await one();
 std::cout << "foo just run one() " << i << std::endl;
 i += co_await two();
 std::cout << "foo just run two() " << i << std::endl;
 std::cout << "foo makes a pause after his work" << std::endl;
 // Suspend (and reschedule)
 co_await suspend{};
 co_await noop();
 co_return i;
}
int main()
{
 // Build the task foo() but does not start any work
 auto tsk = foo();
 execution_handle = tsk.get_handle();
 for(int i = 0; !execution_handle.done(); ++i) {
 std::cout << "Resume # " << i << std::endl;
 execution_handle.resume();
 }
 std::cout << "Result: " << tsk.get_handle().promise().data.value() << std::endl;
 return 0;
}
```