spillable_messages.hpp

 
#pragma once
 
#include <cstdint>
#include <map>
#include <mutex>
#include <optional>
#include <unordered_map>
 
#include <rapidsmpf/memory/buffer_resource.hpp>
#include <rapidsmpf/memory/content_description.hpp>
#include <rapidsmpf/streaming/core/message.hpp>
 
namespace rapidsmpf::streaming {
 
class SpillableMessages {
  public:
    using MessageId = std::uint64_t;
 
    SpillableMessages() = default;
    SpillableMessages(SpillableMessages const&) = delete;
    SpillableMessages& operator=(SpillableMessages const&) = delete;
    SpillableMessages(SpillableMessages&&) noexcept = delete;
    SpillableMessages& operator=(SpillableMessages&&) noexcept = delete;
 
    [[nodiscard]] MessageId insert(Message&& message);
 
    [[nodiscard]] Message extract(MessageId mid);
 
    [[nodiscard]] Message copy(MessageId mid, MemoryReservation& reservation);
 
    [[nodiscard]] std::size_t spill(MessageId mid, BufferResource* br) const;
 
    std::map<MessageId, ContentDescription> get_content_descriptions() const;
 
    ContentDescription get_content_description(MessageId mid) const;
 
  private:
    struct Item {
        mutable std::mutex mutex;
        std::optional<Message> message;
 
        Item() noexcept = default;
 
        Item(Message&& message) : message(std::move(message)) {}
    };
 
    // Never lock the global mutex and an item's mutex at the same time!
    mutable std::mutex global_mutex_;
    MessageId counter_{0};
    std::unordered_map<MessageId, std::shared_ptr<Item>> items_;
    mutable std::map<MessageId, ContentDescription> content_descriptions_;
};
 
}  // namespace rapidsmpf::streaming