chunk.hpp

 
#pragma once
 
#include <memory>
#include <vector>
 
#include <rapidsmpf/communicator/communicator.hpp>
#include <rapidsmpf/memory/buffer.hpp>
#include <rapidsmpf/memory/buffer_resource.hpp>
#include <rapidsmpf/memory/packed_data.hpp>
 
namespace rapidsmpf::shuffler {
 
using PartID = std::uint32_t;
 
namespace detail {
 
using ChunkID = std::uint64_t;
 
class Chunk {
    // friend a method that creates a dummy chunk for testing
    friend Chunk make_dummy_chunk(ChunkID, PartID);
 
  public:
    Chunk(Chunk&& other) noexcept = default;
 
    Chunk& operator=(Chunk&& other) noexcept = default;
 
    // delete copy constructor
    Chunk(Chunk const&) = delete;
 
    // delete copy assignment operator
    Chunk& operator=(Chunk const&) = delete;
 
    static constexpr std::size_t metadata_message_header_size() {
        return sizeof(ChunkID) + sizeof(PartID) + sizeof(std::size_t)
               + sizeof(std::uint32_t) + sizeof(std::uint64_t);
    }
 
    [[nodiscard]] constexpr ChunkID chunk_id() const {
        return chunk_id_;
    }
 
    [[nodiscard]] constexpr PartID part_id() const {
        return part_id_;
    }
 
    [[nodiscard]] constexpr std::size_t expected_num_chunks() const {
        return expected_num_chunks_;
    }
 
    [[nodiscard]] constexpr bool is_control_message() const {
        // We use `expected_num_chunks > 0` to flag a message as a "control message".
        return expected_num_chunks() > 0;
    }
 
    [[nodiscard]] constexpr std::uint32_t metadata_size() const {
        return metadata_size_;
    }
 
    [[nodiscard]] constexpr std::size_t data_size() const {
        return data_size_;
    }
 
    void set_data_buffer(std::unique_ptr<Buffer> data) {
        RAPIDSMPF_EXPECTS(!data_, "buffer is already set");
        data_ = std::move(data);
    }
 
    [[nodiscard]] bool is_data_buffer_set() const {
        return data_ != nullptr;
    }
 
    [[nodiscard]] bool is_metadata_buffer_set() const {
        return metadata_ != nullptr && !metadata_->empty();
    }
 
    [[nodiscard]] MemoryType data_memory_type() const {
        RAPIDSMPF_EXPECTS(data_, "data buffer is not set");
        return data_->mem_type();
    }
 
    [[nodiscard]] std::unique_ptr<std::vector<std::uint8_t>> release_metadata_buffer() {
        return std::move(metadata_);
    }
 
    [[nodiscard]] std::unique_ptr<Buffer> release_data_buffer() {
        return std::move(data_);
    }
 
    static Chunk from_packed_data(
        ChunkID chunk_id, PartID part_id, PackedData&& packed_data
    );
 
    static Chunk from_finished_partition(
        ChunkID chunk_id, PartID part_id, std::size_t expected_num_chunks
    );
 
    static Chunk deserialize(
        std::vector<std::uint8_t> const& msg, BufferResource* br, bool validate = true
    );
 
    static bool validate_format(std::vector<std::uint8_t> const& serialized_buf);
 
    [[nodiscard]] bool is_ready() const {
        // data_size_ contains the size of the data buffer. If it is 0, the chunk
        // has no data, so it is ready. Else, the chunk is ready if the data
        // buffer is non-null and the data buffer is ready.
        return data_size_ == 0 || (data_ && data_->is_latest_write_done());
    }
 
    [[nodiscard]] std::string str() const;
 
    [[nodiscard]] std::unique_ptr<std::vector<std::uint8_t>> serialize() const;
 
  private:
    // constructor
    Chunk(
        ChunkID chunk_id,
        PartID part_id,
        std::size_t expected_num_chunks,
        std::uint32_t metadata_size,
        std::uint64_t data_size,
        std::unique_ptr<std::vector<std::uint8_t>> metadata = nullptr,
        std::unique_ptr<Buffer> data = nullptr
    );
 
    ChunkID chunk_id_;  
    PartID part_id_;  
    std::size_t
        expected_num_chunks_;  
    std::uint32_t metadata_size_;  
    std::uint64_t data_size_;  
 
    std::unique_ptr<std::vector<std::uint8_t>> metadata_;
 
    std::unique_ptr<Buffer> data_;
};
 
std::ostream& operator<<(std::ostream& os, Chunk const& obj);
 
}  // namespace detail
}  // namespace rapidsmpf::shuffler