coro_executor.hpp

 
#pragma once
#include <memory>
#include <thread>
 
#include <coro/coro.hpp>
 
#include <rapidsmpf/config.hpp>
#include <rapidsmpf/error.hpp>
#include <rapidsmpf/statistics.hpp>
 
namespace rapidsmpf::streaming {
 
class CoroThreadPoolExecutor {
  public:
    CoroThreadPoolExecutor(
        std::uint32_t num_streaming_threads,
        std::shared_ptr<Statistics> statistics = Statistics::disabled()
    );
 
    CoroThreadPoolExecutor(
        config::Options options,
        std::shared_ptr<Statistics> statistics = Statistics::disabled()
    );
 
    ~CoroThreadPoolExecutor() noexcept;
 
    CoroThreadPoolExecutor(CoroThreadPoolExecutor&&) = delete;
    CoroThreadPoolExecutor(CoroThreadPoolExecutor const&) = delete;
    CoroThreadPoolExecutor& operator=(CoroThreadPoolExecutor& o) = delete;
    CoroThreadPoolExecutor& operator=(CoroThreadPoolExecutor&& o) = delete;
 
    void shutdown() noexcept;
 
    [[nodiscard]] std::uint32_t num_streaming_threads() const noexcept {
        return executor_->thread_count();
    }
 
    [[nodiscard]] std::unique_ptr<coro::thread_pool>& get() noexcept;
 
    [[nodiscard]] auto schedule() {
        return executor_->schedule();
    }
 
    [[nodiscard]] auto schedule(auto task) {
        return executor_->schedule(std::move(task));
    }
 
    [[nodiscard]] auto yield() {
        return executor_->yield();
    }
 
    auto spawn_detached(auto task) noexcept {
        return executor_->spawn_detached(std::move(task));
    }
 
    auto spawn_joinable(auto task) noexcept {
        return executor_->spawn_joinable(std::move(task));
    }
 
  private:
    std::atomic<bool> is_shutdown_{false};
    std::unique_ptr<coro::thread_pool> executor_;
    std::shared_ptr<Statistics> statistics_;
    std::thread::id creator_thread_id_;
};
 
}  // namespace rapidsmpf::streaming