context.hpp

 
#pragma once
 
#include <memory>
#include <thread>
 
#include <coro/coro.hpp>
 
#include <rapidsmpf/communicator/communicator.hpp>
#include <rapidsmpf/config.hpp>
#include <rapidsmpf/error.hpp>
#include <rapidsmpf/statistics.hpp>
#include <rapidsmpf/streaming/core/channel.hpp>
#include <rapidsmpf/streaming/core/coro_executor.hpp>
#include <rapidsmpf/streaming/core/memory_reserve_or_wait.hpp>
#include <rapidsmpf/streaming/core/queue.hpp>
 
namespace rapidsmpf::streaming {
 
class Context {
  public:
    Context(
        config::Options options,
        std::shared_ptr<Communicator::Logger> logger,
        std::shared_ptr<CoroThreadPoolExecutor> executor,
        std::shared_ptr<BufferResource> br
    );
 
    Context(
        config::Options options,
        std::shared_ptr<Communicator::Logger> logger,
        std::shared_ptr<BufferResource> br
    );
 
    static std::shared_ptr<Context> from_options(
        RmmResourceAdaptor* mr,
        std::shared_ptr<Communicator::Logger> logger,
        config::Options options
    );
 
    // No copy constructor and assignment operator.
    Context(Context const&) = delete;
    Context& operator=(Context const&) = delete;
 
    // No move constructor and assignment operator.
    Context(Context&&) = delete;
    Context& operator=(Context&&) = delete;
 
    ~Context() noexcept;
 
    void shutdown() noexcept;
 
    [[nodiscard]] config::Options options() const noexcept;
 
    [[nodiscard]] std::shared_ptr<Communicator::Logger> const& logger() const noexcept;
 
    [[nodiscard]] std::shared_ptr<CoroThreadPoolExecutor> const&
    executor() const noexcept;
 
    [[nodiscard]] std::shared_ptr<BufferResource> const& br() const noexcept;
 
    [[nodiscard]] std::shared_ptr<MemoryReserveOrWait> const& memory(
        MemoryType mem_type
    ) const noexcept;
 
    [[nodiscard]] std::shared_ptr<Statistics> statistics() const noexcept;
 
    [[nodiscard]] std::shared_ptr<Channel> create_channel() const noexcept;
 
    [[nodiscard]] std::shared_ptr<SpillableMessages> const&
    spillable_messages() const noexcept;
 
    [[nodiscard]] std::shared_ptr<BoundedQueue> create_bounded_queue(
        std::size_t buffer_size
    ) const noexcept;
 
    [[nodiscard]] std::size_t uid() const noexcept;
 
  private:
    std::size_t const uid_;
    std::atomic<bool> is_shutdown_{false};
    std::thread::id creator_thread_id_;
    config::Options options_;
    std::shared_ptr<Communicator::Logger> logger_;
    std::shared_ptr<CoroThreadPoolExecutor> executor_;
    std::shared_ptr<BufferResource> br_;
    std::array<std::shared_ptr<MemoryReserveOrWait>, MEMORY_TYPES.size()> memory_ = {};
    std::shared_ptr<SpillableMessages> spillable_messages_;
    SpillManager::SpillFunctionID spill_function_id_{};
};
 
}  // namespace rapidsmpf::streaming