chunk.hpp

 
#pragma once
 
#include <memory>
#include <sstream>
#include <vector>
 
#include <rapidsmpf/buffer/buffer.hpp>
#include <rapidsmpf/buffer/packed_data.hpp>
#include <rapidsmpf/communicator/communicator.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 size_t metadata_message_header_size(size_t n_messages) {
        return sizeof(ChunkID) + sizeof(size_t)
               + n_messages
                     * (sizeof(PartID) + sizeof(size_t) + sizeof(uint32_t)
                        + sizeof(uint64_t));
    }
 
    [[nodiscard]] constexpr ChunkID chunk_id() const {
        return chunk_id_;
    }
 
    [[nodiscard]] constexpr size_t n_messages() const {
        return part_ids_.size();
    }
 
    [[nodiscard]] constexpr PartID part_id(size_t i) const {
        return part_ids_.at(i);
    }
 
    [[nodiscard]] constexpr size_t expected_num_chunks(size_t i) const {
        return expected_num_chunks_.at(i);
    }
 
    [[nodiscard]] constexpr bool is_control_message(size_t i) const {
        // We use `expected_num_chunks > 0` to flag a message as a "control message".
        return expected_num_chunks(i) > 0;
    }
 
    Chunk get_data(ChunkID new_chunk_id, size_t i, BufferResource* br);
 
    [[nodiscard]] constexpr uint32_t metadata_size(size_t i) const {
        return i == 0 ? meta_offsets_.at(0)
                      : meta_offsets_.at(i) - meta_offsets_.at(i - 1);
    }
 
    [[nodiscard]] constexpr size_t data_size(size_t i) const {
        return i == 0 ? data_offsets_.at(0)
                      : data_offsets_.at(i) - data_offsets_.at(i - 1);
    }
 
    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]] constexpr size_t concat_data_size() const {
        return data_offsets_.at(n_messages() - 1);
    }
 
    [[nodiscard]] constexpr size_t concat_metadata_size() const {
        return meta_offsets_.at(n_messages() - 1);
    }
 
    [[nodiscard]] std::unique_ptr<std::vector<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, size_t expected_num_chunks
    );
 
    static Chunk deserialize(std::vector<uint8_t> const& msg, bool validate = true);
 
    static bool validate_format(std::vector<uint8_t> const& serialized_buf);
 
    [[nodiscard]] bool is_ready() const {
        // data_offsets_[-1] contains the size of the data buffer. If it is 0, the chunk
        // has no data messages, so it is ready. Else, the chunk is ready if the data
        // buffer is non-null and the data buffer is ready.
        return !data_offsets_.empty()
               && (data_offsets_[n_messages() - 1] == 0
                   || (data_ && data_->is_latest_write_done()));
    }
 
    [[nodiscard]] std::string str() const;
 
    [[nodiscard]] std::unique_ptr<std::vector<uint8_t>> serialize() const;
 
    static Chunk concat(
        std::vector<Chunk>&& chunks,
        ChunkID chunk_id,
        BufferResource* br,
        std::optional<MemoryType> preferred_mem_type = std::nullopt
    );
 
  private:
    // constructor
    Chunk(
        ChunkID chunk_id,
        std::vector<PartID> part_ids,
        std::vector<size_t> expected_num_chunks,
        std::vector<uint32_t> meta_offsets,
        std::vector<uint64_t> data_offsets,
        std::unique_ptr<std::vector<uint8_t>> metadata = nullptr,
        std::unique_ptr<Buffer> data = nullptr
    );
 
    ChunkID chunk_id_;  
    std::vector<PartID> part_ids_;  
    std::vector<size_t> expected_num_chunks_;  
    std::vector<uint32_t>
        meta_offsets_;  
    std::vector<uint64_t>
        data_offsets_;  
 
    std::unique_ptr<std::vector<uint8_t>> metadata_;
 
    std::unique_ptr<Buffer> data_;
};
 
class ReadyForDataMessage {
  public:
    ChunkID cid;  
 
    static constexpr size_t byte_size = sizeof(ChunkID);
 
    [[nodiscard]] std::unique_ptr<std::vector<uint8_t>> pack() {
        auto msg = std::make_unique<std::vector<uint8_t>>(sizeof(ChunkID));
        std::memcpy(msg->data(), &cid, sizeof(cid));
        return msg;
    }
 
    [[nodiscard]] static ReadyForDataMessage unpack(
        std::unique_ptr<std::vector<uint8_t>> const& msg
    ) {
        ChunkID cid;
        std::memcpy(&cid, msg->data(), sizeof(cid));
        return ReadyForDataMessage{cid};
    }
 
    [[nodiscard]] std::string str() const {
        std::stringstream ss;
        ss << "ReadyForDataMessage(cid=" << cid << ")";
        return ss.str();
    }
};
 
inline std::ostream& operator<<(std::ostream& os, Chunk const& obj) {
    os << obj.str();
    return os;
}
 
inline std::ostream& operator<<(std::ostream& os, ReadyForDataMessage const& obj) {
    os << obj.str();
    return os;
}
 
}  // namespace detail
}  // namespace rapidsmpf::shuffler