core.hpp

 
#pragma once
 
#include <cstdint>
#include <functional>
#include <memory>
#include <string>
#include <unordered_map>
#include <vector>
 
#include <rapidsmpf/communicator/communicator.hpp>
#include <rapidsmpf/memory/buffer.hpp>
#include <rapidsmpf/memory/buffer_resource.hpp>
#include <rapidsmpf/statistics.hpp>
 
namespace rapidsmpf::communicator {
 
class MetadataPayloadExchange {
  public:
    class Message {
      public:
        Message(
            Rank peer_rank,
            std::vector<std::uint8_t>&& metadata,
            std::unique_ptr<Buffer> data = nullptr
        );
 
        [[nodiscard]] constexpr Rank peer_rank() const noexcept {
            return peer_rank_;
        }
 
        [[nodiscard]] constexpr std::vector<std::uint8_t> const&
        metadata() const noexcept {
            return metadata_;
        }
 
        [[nodiscard]] std::vector<std::uint8_t> release_metadata() noexcept;
 
        [[nodiscard]] Buffer const* data() const;
 
        [[nodiscard]] std::unique_ptr<Buffer> release_data() noexcept;
 
        void set_data(std::unique_ptr<Buffer> buffer);
 
      private:
        Rank peer_rank_;
        std::vector<std::uint8_t> metadata_;
        std::unique_ptr<Buffer> data_;
    };
 
    virtual ~MetadataPayloadExchange() = default;
 
    virtual void send(std::unique_ptr<Message> message) = 0;
 
    virtual void send(std::vector<std::unique_ptr<Message>>&& messages) = 0;
 
    virtual void progress() = 0;
 
    [[nodiscard]] virtual std::vector<std::unique_ptr<Message>> recv() = 0;
 
    [[nodiscard]] virtual bool is_idle() const = 0;
};
 
 
}  // namespace rapidsmpf::communicator