context.hpp

 
#pragma once
 
#include <memory>
 
#include <rapidsmpf/communicator/communicator.hpp>
#include <rapidsmpf/config.hpp>
#include <rapidsmpf/error.hpp>
#include <rapidsmpf/progress_thread.hpp>
#include <rapidsmpf/statistics.hpp>
#include <rapidsmpf/streaming/core/channel.hpp>
#include <rapidsmpf/streaming/core/queue.hpp>
 
#include <coro/coro.hpp>
 
namespace rapidsmpf::streaming {
 
class Context {
  public:
    Context(
        config::Options options,
        std::shared_ptr<Communicator> comm,
        std::shared_ptr<ProgressThread> progress_thread,
        std::unique_ptr<coro::thread_pool> executor,
        std::shared_ptr<BufferResource> br,
        std::shared_ptr<Statistics> statistics
    );
 
    Context(
        config::Options options,
        std::shared_ptr<Communicator> comm,
        std::shared_ptr<BufferResource> br,
        std::shared_ptr<Statistics> statistics = Statistics::disabled()
    );
 
    ~Context() noexcept;
 
    [[nodiscard]] config::Options get_options() const noexcept;
 
    [[nodiscard]] std::shared_ptr<Communicator> comm() const noexcept;
 
    [[nodiscard]] std::shared_ptr<ProgressThread> progress_thread() const noexcept;
 
    [[nodiscard]] std::unique_ptr<coro::thread_pool>& executor() noexcept;
 
    [[nodiscard]] BufferResource* br() const noexcept;
 
    [[nodiscard]] std::shared_ptr<Statistics> statistics() const noexcept;
 
    [[nodiscard]] std::shared_ptr<Channel> create_channel() const noexcept;
 
    [[nodiscard]] std::shared_ptr<SpillableMessages> spillable_messages() const {
        return spillable_messages_;
    }
 
    [[nodiscard]] std::shared_ptr<BoundedQueue> create_bounded_queue(
        std::size_t buffer_size
    ) const noexcept;
 
  private:
    config::Options options_;
    std::shared_ptr<Communicator> comm_;
    std::shared_ptr<ProgressThread> progress_thread_;
    std::unique_ptr<coro::thread_pool> executor_;
    std::shared_ptr<BufferResource> br_;
    std::shared_ptr<Statistics> statistics_;
    std::shared_ptr<SpillableMessages> spillable_messages_;
    SpillManager::SpillFunctionID spill_function_id_{};
};
 
}  // namespace rapidsmpf::streaming