message.hpp

 
#pragma once
 
#include <any>
#include <cstdint>
#include <functional>
#include <memory>
#include <stdexcept>
#include <typeinfo>
#include <utility>
 
#include <rapidsmpf/buffer/buffer.hpp>
#include <rapidsmpf/buffer/content_description.hpp>
#include <rapidsmpf/buffer/resource.hpp>
#include <rapidsmpf/error.hpp>
 
namespace rapidsmpf::streaming {
 
 
class Message {
  public:
    using CopyCallback =
        std::function<Message(Message const&, MemoryReservation& reservation)>;
 
    Message() = default;
 
    template <typename T>
    Message(
        std::uint64_t sequence_number,
        std::unique_ptr<T> payload,
        ContentDescription content_description,
        CopyCallback copy_cb = nullptr
    )
        : sequence_number_(sequence_number),
          content_description_{content_description},
          copy_cb_{std::move(copy_cb)} {
        RAPIDSMPF_EXPECTS(
            payload != nullptr, "nullptr not allowed", std::invalid_argument
        );
        // Conceptually, a `std::unique_ptr` would be sufficient for exclusive ownership.
        // However, since `std::any` requires its contents to be copyable, we store the
        // payload in a `std::shared_ptr` instead. This shared ownership is internal only
        // and never exposed to the user.
        payload_ = std::shared_ptr<T>(std::move(payload));
    }
 
    // In tandem with coro::queue the move assignment of std::any breaks GCC's
    // uninitialized variable tracking and we get a warning that std::any::_M_manager' may
    // be used uninitialized [-Wmaybe-uninitialized] This is a bug in GCC 14.x which we
    // workaround by suppressing the warning for the move ctors/assignment. Fixed in
    // GCC 15.2.
#if defined(__GNUC__) && __GNUC__ == 14
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wmaybe-uninitialized"
#endif
    Message(Message&& other) noexcept = default;
 
    Message& operator=(Message&& other) noexcept = default;
#if defined(__GNUC__) && __GNUC__ == 14
#pragma GCC diagnostic pop
#endif
    Message(Message const&) = delete;
    Message& operator=(Message const&) = delete;
 
    void reset() noexcept {
        return payload_.reset();
    }
 
    [[nodiscard]] bool empty() const noexcept {
        return !payload_.has_value();
    }
 
    [[nodiscard]] constexpr std::uint64_t sequence_number() const noexcept {
        return sequence_number_;
    }
 
    template <typename T>
    [[nodiscard]] bool holds() const noexcept {
        return payload_.type() == typeid(std::shared_ptr<T>);
    }
 
    template <typename T>
    [[nodiscard]] T const& get() const {
        return *get_ptr<T>();
    }
 
    template <typename T>
    [[nodiscard]] T release() {
        std::shared_ptr<T> ret = get_ptr<T>();
        reset();
        return std::move(*ret);
    }
 
    [[nodiscard]] constexpr ContentDescription const&
    content_description() const noexcept {
        return content_description_;
    }
 
    [[nodiscard]] constexpr CopyCallback const& copy_cb() const noexcept {
        return copy_cb_;
    }
 
    [[nodiscard]] constexpr size_t copy_cost() const noexcept {
        return content_description().content_size();
    }
 
    [[nodiscard]] Message copy(MemoryReservation& reservation) const {
        RAPIDSMPF_EXPECTS(
            copy_cb(), "message doesn't support `copy`", std::invalid_argument
        );
        return copy_cb()(*this, reservation);
    }
 
  private:
    template <typename T>
    [[nodiscard]] std::shared_ptr<T> get_ptr() const {
        RAPIDSMPF_EXPECTS(!empty(), "message is empty", std::invalid_argument);
        RAPIDSMPF_EXPECTS(holds<T>(), "wrong message type", std::invalid_argument);
        return std::any_cast<std::shared_ptr<T>>(payload_);
    }
 
  private:
    std::uint64_t sequence_number_{0};
    std::any payload_;
    ContentDescription content_description_;
    CopyCallback copy_cb_;
};
 
}  // namespace rapidsmpf::streaming