async_task.hpp

Go to the documentation of this file.

 
#pragma once
 
#include "task_registry.hpp"
 
#include <atomic>
#include <concepts>
#include <coroutine>
#include <exception>
#include <optional>
#include <string>
#include <uni20/common/trace.hpp>
#include <utility>
 
namespace uni20::async
{
 
class IScheduler;
class TaskPromiseBase;
class AsyncTaskPromise;
class CudaTaskPromise;
class TaskFactory;
class AsyncTask;
class CudaTask;
 
enum class TaskDomain
{
  none,
  host,
  cuda,
};
 
struct TaskRoute
{
    TaskDomain domain = TaskDomain::none;
    std::optional<int> cuda_device{};
};
 
template <typename T>
concept TaskPromise = std::derived_from<T, TaskPromiseBase> && requires {
  { T::task_domain } -> std::convertible_to<TaskDomain>;
};
 
class TaskHandle {
  public:
    constexpr TaskHandle() noexcept = default;
 
    template <TaskPromise Promise> [[nodiscard]] static TaskHandle from(std::coroutine_handle<Promise> handle) noexcept;
 
    [[nodiscard]] explicit constexpr operator bool() const noexcept { return static_cast<bool>(coroutine_); }
 
    [[nodiscard]] constexpr std::coroutine_handle<> coroutine() const noexcept { return coroutine_; }
 
    [[nodiscard]] TaskPromiseBase& promise() const noexcept;
 
    [[nodiscard]] TaskDomain domain() const noexcept
    {
      DEBUG_PRECONDITION(domain_ != TaskDomain::none);
      return domain_;
    }
 
    [[nodiscard]] bool done() const noexcept { return !coroutine_ || coroutine_.done(); }
 
    friend constexpr bool operator==(TaskHandle lhs, TaskHandle rhs) noexcept
    {
      return lhs.coroutine_ == rhs.coroutine_;
    }
 
  private:
    TaskHandle(std::coroutine_handle<> coroutine, TaskPromiseBase* promise, TaskDomain domain) noexcept
        : coroutine_(coroutine), promise_(promise), domain_(domain)
    {
      DEBUG_CHECK_EQUAL(static_cast<bool>(coroutine_), promise_ != nullptr);
      DEBUG_CHECK_EQUAL(static_cast<bool>(coroutine_), domain_ != TaskDomain::none);
    }
 
    std::coroutine_handle<> coroutine_{};
    TaskPromiseBase* promise_ = nullptr;
    TaskDomain domain_ = TaskDomain::none;
 
    friend class TaskPromiseBase;
};
 
class BasicTask {
  public:
    BasicTask() noexcept = default;
 
    BasicTask(BasicTask const&) = delete;
    BasicTask& operator=(BasicTask const&) = delete;
 
    BasicTask(BasicTask&& other) noexcept
        : handle_(std::exchange(other.handle_, {})), await_exception_(std::exchange(other.await_exception_, {}))
    {}
 
    BasicTask& operator=(BasicTask&& other) noexcept;
 
    ~BasicTask() noexcept;
 
    [[nodiscard]] explicit operator bool() const noexcept { return static_cast<bool>(handle_); }
 
    [[nodiscard]] TaskHandle handle() const noexcept { return handle_; }
 
    [[nodiscard]] std::coroutine_handle<> coroutine_handle() const noexcept { return handle_.coroutine(); }
 
    BasicTask& debug_name(std::string const& label)
    {
      TaskRegistry::name_task(this->coroutine_handle(), label);
      return *this;
    }
 
    void resume();
 
    void abandon_leak();
 
    [[nodiscard]] TaskHandle release_handle();
 
    void set_cancel_on_resume() noexcept;
 
    void exception_on_resume(std::exception_ptr error) noexcept;
 
    void destroy() noexcept;
 
    [[nodiscard]] bool set_scheduler(IScheduler* scheduler);
 
    [[nodiscard]] std::optional<int> preferred_numa_node() const noexcept;
    void set_preferred_numa_node(std::optional<int> node) noexcept;
 
    static void reschedule(BasicTask task);
 
    [[nodiscard]] static BasicTask make_sole_owner(BasicTask&& task);
 
    [[nodiscard]] bool await_ready() const noexcept { return !handle_ || handle_.done(); }
 
    template <TaskPromise ParentPromise>
    std::coroutine_handle<> await_suspend(std::coroutine_handle<ParentPromise> parent);
 
    void await_resume();
 
  protected:
    struct construction_key
    {};
 
    explicit BasicTask(TaskHandle handle, construction_key) noexcept : handle_(handle) {}
 
  private:
    [[nodiscard]] bool can_destroy_coroutine(TaskHandle handle) const noexcept;
    [[nodiscard]] TaskHandle release_ownership();
    void release() noexcept;
    void destroy_owned_coroutine() noexcept;
 
    TaskHandle handle_{};
    std::exception_ptr await_exception_{};
 
    friend class TaskPromiseBase;
    friend class TaskFactory;
    friend struct AsyncTaskTestAccess;
};
 
class AsyncTask final : public BasicTask {
  public:
    using base_type = BasicTask;
    using promise_type = AsyncTaskPromise;
 
    AsyncTask() = default;
    AsyncTask(AsyncTask const&) = delete;
    AsyncTask& operator=(AsyncTask const&) = delete;
    AsyncTask(AsyncTask&&) noexcept = default;
    AsyncTask& operator=(AsyncTask&&) noexcept = default;
 
    AsyncTask& debug_name(std::string const& label)
    {
      this->BasicTask::debug_name(label);
      return *this;
    }
 
  private:
    explicit AsyncTask(TaskHandle handle) noexcept : BasicTask(handle, construction_key{}) {}
 
    friend class AsyncTaskPromise;
};
 
} // namespace uni20::async
 
#include "async_task_promise.hpp"