memory_reserve_or_wait.hpp

 
#pragma once
 
#include <optional>
#include <set>
 
#include <coro/task.hpp>
 
#include <rapidsmpf/config.hpp>
#include <rapidsmpf/memory/buffer_resource.hpp>
#include <rapidsmpf/streaming/core/actor.hpp>
#include <rapidsmpf/streaming/core/coro_executor.hpp>
#include <rapidsmpf/streaming/core/coro_utils.hpp>
#include <rapidsmpf/utils/misc.hpp>
 
namespace rapidsmpf::streaming {
 
// Forward declaration
class Context;
 
class MemoryReserveOrWait {
  public:
    static constexpr std::int64_t missing_net_memory_delta = 0;
 
    MemoryReserveOrWait(
        config::Options options,
        MemoryType mem_type,
        std::shared_ptr<CoroThreadPoolExecutor> executor,
        std::shared_ptr<BufferResource> br
    );
 
    ~MemoryReserveOrWait() noexcept;
 
    Actor shutdown();
 
    coro::task<MemoryReservation> reserve_or_wait(
        std::size_t size, std::int64_t net_memory_delta
    );
 
    coro::task<std::pair<MemoryReservation, std::size_t>> reserve_or_wait_or_overbook(
        std::size_t size, std::int64_t net_memory_delta
    );
 
    coro::task<MemoryReservation> reserve_or_wait_or_fail(
        std::size_t size, std::int64_t net_memory_delta
    );
 
    [[nodiscard]] std::size_t size() const noexcept;
 
    [[nodiscard]] std::size_t periodic_memory_check_counter() const noexcept;
 
    [[nodiscard]] std::shared_ptr<CoroThreadPoolExecutor> const&
    executor() const noexcept;
 
    [[nodiscard]] std::shared_ptr<BufferResource> const& br() const noexcept;
 
    [[nodiscard]] Duration timeout() const noexcept;
 
  private:
    struct Request {
        std::size_t size;
 
        std::int64_t net_memory_delta;
 
        std::uint64_t sequence_number;
 
        coro::queue<MemoryReservation>& queue;
 
        friend bool operator<(Request const& a, Request const& b) {
            return std::tie(a.size, a.sequence_number)
                   < std::tie(b.size, b.sequence_number);
        }
    };
 
    coro::task<void> periodic_memory_check();
 
    mutable std::mutex mutex_;
    std::uint64_t sequence_counter{0};
    MemoryType const mem_type_;
    std::shared_ptr<CoroThreadPoolExecutor> executor_;
    std::shared_ptr<BufferResource> br_;
    Duration const timeout_;
    std::set<Request> reservation_requests_;
    std::atomic<std::uint64_t> periodic_memory_check_counter_{0};
    std::optional<coro::task<void>> periodic_memory_check_task_;
};
 
coro::task<MemoryReservation> reserve_memory(
    std::shared_ptr<Context> ctx,
    std::size_t size,
    std::int64_t net_memory_delta,
    MemoryType mem_type = MemoryType::DEVICE,
    std::optional<AllowOverbooking> allow_overbooking = std::nullopt
);
 
}  // namespace rapidsmpf::streaming