task.h

Go to the documentation of this file.

// BSD 3-Clause License
// Copyright (c) 2024, 🍀☀🌕🌥 🌊
// See the LICENSE file in the project root for full license information.
 
#pragma once
 
#include <coroutine>
#include <exception>
#include <optional>
#include <utility>
#include <type_traits>
#include <variant>
#include <atomic>
 
namespace kcenon::container::async
{
    template<typename T>
    class task;
 
    namespace detail
    {
        struct promise_base
        {
            std::exception_ptr exception_;
            std::coroutine_handle<> continuation_;
            std::atomic<bool> completed_{false};
 
            [[nodiscard]] std::suspend_never initial_suspend() noexcept
            {
                return {};
            }
 
            void unhandled_exception() noexcept
            {
                exception_ = std::current_exception();
            }
 
            struct final_awaiter
            {
                [[nodiscard]] bool await_ready() const noexcept
                {
                    return false;
                }
 
                template<typename Promise>
                std::coroutine_handle<> await_suspend(
                    std::coroutine_handle<Promise> handle) noexcept
                {
                    auto& promise = handle.promise();
                    // Read continuation BEFORE marking as complete to avoid
                    // data race with destructor. Once completed_ is true,
                    // the task owner may destroy the coroutine frame.
                    auto continuation = promise.continuation_;
                    // Mark as completed with release semantics to ensure
                    // all prior writes are visible to threads checking done()
                    promise.completed_.store(true, std::memory_order_release);
                    if (continuation)
                    {
                        return continuation;
                    }
                    return std::noop_coroutine();
                }
 
                void await_resume() noexcept {}
            };
 
            [[nodiscard]] final_awaiter final_suspend() noexcept
            {
                return {};
            }
        };
 
        template<typename T>
        struct promise_type : promise_base
        {
            std::optional<T> result_;
 
            [[nodiscard]] task<T> get_return_object() noexcept;
 
            template<typename U>
            void return_value(U&& value)
                requires std::convertible_to<U, T>
            {
                result_.emplace(std::forward<U>(value));
            }
 
            [[nodiscard]] T& result() &
            {
                if (exception_)
                {
                    std::rethrow_exception(exception_);
                }
                return *result_;
            }
 
            [[nodiscard]] T&& result() &&
            {
                if (exception_)
                {
                    std::rethrow_exception(exception_);
                }
                return std::move(*result_);
            }
        };
 
        template<>
        struct promise_type<void> : promise_base
        {
            [[nodiscard]] task<void> get_return_object() noexcept;
 
            void return_void() noexcept {}
 
            void result()
            {
                if (exception_)
                {
                    std::rethrow_exception(exception_);
                }
            }
        };
 
    } // namespace detail
 
    template<typename T>
    class task
    {
    public:
        using promise_type = detail::promise_type<T>;
        using handle_type = std::coroutine_handle<promise_type>;
 
        task() noexcept : handle_(nullptr) {}
 
        explicit task(handle_type handle) noexcept : handle_(handle) {}
 
        task(task&& other) noexcept : handle_(std::exchange(other.handle_, nullptr)) {}
 
        task& operator=(task&& other) noexcept
        {
            if (this != &other)
            {
                if (handle_)
                {
                    handle_.destroy();
                }
                handle_ = std::exchange(other.handle_, nullptr);
            }
            return *this;
        }
 
        ~task()
        {
            if (handle_)
            {
                handle_.destroy();
            }
        }
 
        // Non-copyable
        task(const task&) = delete;
        task& operator=(const task&) = delete;
 
        [[nodiscard]] bool valid() const noexcept
        {
            return handle_ != nullptr;
        }
 
        [[nodiscard]] explicit operator bool() const noexcept
        {
            return valid();
        }
 
        [[nodiscard]] bool done() const noexcept
        {
            return handle_ && handle_.promise().completed_.load(std::memory_order_acquire);
        }
 
        [[nodiscard]] decltype(auto) get() &
        {
            return handle_.promise().result();
        }
 
        [[nodiscard]] decltype(auto) get() &&
        {
            return std::move(handle_.promise()).result();
        }
 
        struct awaiter
        {
            handle_type handle_;
 
            [[nodiscard]] bool await_ready() const noexcept
            {
                return handle_.done();
            }
 
            std::coroutine_handle<> await_suspend(
                std::coroutine_handle<> awaiting) noexcept
            {
                handle_.promise().continuation_ = awaiting;
                return handle_;
            }
 
            [[nodiscard]] decltype(auto) await_resume()
            {
                return handle_.promise().result();
            }
        };
 
        [[nodiscard]] awaiter operator co_await() noexcept
        {
            return awaiter{handle_};
        }
 
        void resume() const
        {
            if (handle_ && !handle_.done())
            {
                handle_.resume();
            }
        }
 
    private:
        handle_type handle_;
    };
 
    // Implementation of get_return_object for value type
    template<typename T>
    task<T> detail::promise_type<T>::get_return_object() noexcept
    {
        return task<T>{std::coroutine_handle<detail::promise_type<T>>::from_promise(*this)};
    }
 
    // Implementation of get_return_object for void type
    inline task<void> detail::promise_type<void>::get_return_object() noexcept
    {
        return task<void>{std::coroutine_handle<detail::promise_type<void>>::from_promise(*this)};
    }
 
    template<typename T>
    task<std::decay_t<T>> make_ready_task(T&& value)
    {
        co_return std::forward<T>(value);
    }
 
    inline task<void> make_ready_task()
    {
        co_return;
    }
 
    template<typename T>
    task<T> make_exceptional_task(std::exception_ptr ex)
    {
        std::rethrow_exception(ex);
        // Unreachable, but needed for coroutine
        co_return T{};
    }
 
    template<>
    inline task<void> make_exceptional_task<void>(std::exception_ptr ex)
    {
        std::rethrow_exception(ex);
        co_return;
    }
 
} // namespace kcenon::container::async