stdexec-server / src / task.hpp

#pragma once

#include <cassert>
#include <coroutine>
#include <exception>
#include <optional>
#include <stdexec/execution.hpp>
#include <type_traits>
#include <utility>

namespace kev
{

template <typename T> struct task;

/**
 * @brief Common data for all promise types.
 *
 */
struct promise_data
{
    std::exception_ptr m_exception{};
};

/**
 * @brief Promise base for tasks with return type T.
 *
 * @tparam T The return type of the task.
 */
template <typename T> struct promise_base : promise_data
{
    /**
     * @brief Store the returned value in the promise.
     *
     * @param value
     */
    void return_value(T value)
    {
        this->m_value = std::move(value);
    }
    std::optional<T> m_value{std::nullopt};
};

/**
 * @brief Specialization of promise base for void return type.
 *
 */
template <> struct promise_base<void> : promise_data
{
    /**
     * @brief Handle the void return type.
     *
     */
    static constexpr void return_void()
    {
    }
};

/**
 * @brief Promise type for tasks.
 *
 * @tparam T The return type of the task.
 */
template <typename T> struct task_promise : stdexec::with_awaitable_senders<task_promise<T>>, promise_base<T>
{
    using promise_type = task_promise<T>;
    using coroutine_handle = std::coroutine_handle<promise_type>;

    /**
     * @brief Get the return object object
     *
     * @return task<T>
     */
    auto get_return_object() -> task<T>
    {
        return task<T>{coroutine_handle::from_promise(*this)};
    }
    /**
     * @brief The task type always suspends at the beginning.
     *
     * @return auto
     */
    auto initial_suspend() -> std::suspend_always
    {
        return std::suspend_always{};
    }

    /**
     * @brief Final awaiter to resume the continuation.
     *
     */
    struct final_awaiter
    {
        /**
         * @brief Always suspend to resume the continuation.
         *
         */
        static constexpr auto await_ready() noexcept -> bool
        {
            // Always suspend to allow resuming the continuation.
            return false;
        }

        std::coroutine_handle<> await_suspend(std::coroutine_handle<promise_type> h) noexcept
        {
            // cppreference: "if await_suspend returns a coroutine handle for some other coroutine, that handle is
            // resumed (by a call to handle.resume())"
            return h.promise().m_continuation;
        }

        void await_resume() noexcept
        {
            // Nothing to do here.
        }
    };

    /**
     * @brief Final suspend point to resume the continuation.
     *
     * @return final_awaiter
     */
    auto final_suspend() noexcept -> final_awaiter
    {
        return {};
    }

    /**
     * @brief Store the current exception in the promise.
     *
     */
    auto unhandled_exception() -> void
    {
        // cppreference:
        // If the coroutine ends with an uncaught exception, it performs the following:
        // catches the exception and calls promise.unhandled_exception() from within the catch-block
        this->m_exception = std::current_exception();
    }

    std::coroutine_handle<> m_continuation;
};

/**
 * @brief Awaiter for tasks. Allows awaiting the task and retrieving its result.
 *
 * @tparam T The return type of the task.
 */
template <typename T> struct task_awaiter
{
    using coroutine_handle = std::coroutine_handle<task_promise<T>>;

    explicit task_awaiter(coroutine_handle handle) noexcept : m_coroutine_handle(handle)
    {
        // Take ownership of the coroutine handle.
    }

    task_awaiter(task_awaiter &&other) noexcept : m_coroutine_handle(std::exchange(other.m_coroutine_handle, {}))
    {
        // Awaiters are move only.
    }

    task_awaiter(const task_awaiter &) = delete;
    task_awaiter &operator=(const task_awaiter &) = delete;

    task_awaiter &operator=(task_awaiter &&other) noexcept
    {
        if (this != &other)
        {
            if (m_coroutine_handle)
            {
                m_coroutine_handle.destroy();
            }
            m_coroutine_handle = std::exchange(other.m_coroutine_handle, {});
        }
        return *this;
    }

    ~task_awaiter()
    {
        // Destroy the coroutine if we still own it. There's a possibility of transferring ownership via move
        // construction/assignment.
        if (m_coroutine_handle)
        {
            m_coroutine_handle.destroy();
        }
    }

    auto await_ready() -> bool
    {
        // Always suspend to allow the caller to await the task.
        return false;
    }
    auto await_suspend(std::coroutine_handle<> continuation) -> bool
    {
        /*
        cppreference on await_suspend return value:
        if await_suspend returns bool:
            - the value true returns control to the caller/resumer of the current coroutine
            - the value false resumes the current coroutine.
        */

        // Get the promise of the coroutine that is being awaited.
        auto &promise = m_coroutine_handle.promise();
        // Set the continuation to a noop coroutine to avoid resuming the caller too early.
        promise.m_continuation = std::noop_coroutine();
        // Resume the task coroutine. Since we set the continuation to a noop, the caller won't be resumed yet.
        m_coroutine_handle.resume();
        if (m_coroutine_handle.done())
        {
            // If the task is already done, we can resume the caller immediately.
            return false;
        }
        // Set the actual continuation to the caller's coroutine handle.
        promise.m_continuation = continuation;
        // Indicate that we have a continuation to resume later.
        return true;
    }
    auto await_resume() -> T
    {
        // The coroutine body has completed at this point.
        if (m_coroutine_handle.promise().m_exception)
        {
            // If an exception was thrown in the body it would have been stored in the promise.
            std::rethrow_exception(m_coroutine_handle.promise().m_exception);
        }
        if constexpr (std::is_void_v<T>)
        {
            return;
        }
        else
        {
            assert(m_coroutine_handle.promise().m_value.has_value() && "Task did not set a return value");
            return std::move(*(m_coroutine_handle.promise().m_value));
        }
    }

  private:
    coroutine_handle m_coroutine_handle;
};

template <typename T> struct task
{
    using sender_concept =
        stdexec::sender_t; // Mark task as a sender. This allows us to chain tasks with stdexec algorithms.
    using promise_type = task_promise<T>;                         // The promise type associated with this task.
    using coroutine_handle = std::coroutine_handle<promise_type>; // The coroutine handle type for this task.

    task_awaiter<T> operator co_await() &&
    {
        if (!m_handle)
        {
            throw std::runtime_error("Attempting to co_await a moved-from task");
        }
        // Move the handle out of the task to ensure it's only awaited once.
        return task_awaiter(std::exchange(m_handle, {}));
    }

    task(task &&other) noexcept : m_handle(std::exchange(other.m_handle, {}))
    {
    }

    task &operator=(task &&other) noexcept
    {
        if (this != &other)
        {
            if (m_handle)
            {
                m_handle.destroy();
            }
            m_handle = std::exchange(other.m_handle, {});
        }
        return *this;
    }

    task(const task &) = delete;
    task &operator=(const task &) = delete;

    ~task()
    {
        if (m_handle)
        {
            m_handle.destroy();
        }
    }

  private:
    friend promise_type;
    explicit task(coroutine_handle handle) : m_handle(handle)
    {
    }
    coroutine_handle m_handle{};
};

} // namespace kev