postbox.hpp

 
#pragma once
 
#include <mutex>
#include <string>
#include <unordered_map>
#include <vector>
 
#include <rapidsmpf/error.hpp>
#include <rapidsmpf/shuffler/chunk.hpp>
 
namespace rapidsmpf::shuffler::detail {
 
class ChunksToSend {
  public:
    ChunksToSend() = default;
 
    void insert(std::unique_ptr<Chunk> c);
 
    [[nodiscard]] std::vector<Chunk> extract_ready();
 
    [[nodiscard]] bool empty() const;
 
    [[nodiscard]] std::string str() const;
 
  private:
    mutable std::mutex mutex_{};
    std::vector<std::unique_ptr<Chunk>> chunks_{};
};
 
inline std::ostream& operator<<(std::ostream& os, ChunksToSend const& obj) {
    os << obj.str();
    return os;
}
 
class ReceivedChunks {
  public:
    ReceivedChunks(std::size_t num_keys_hint = 0) {
        if (num_keys_hint > 0) {
            pigeonhole_.reserve(num_keys_hint);
        }
    }
 
    void insert(Chunk&& chunk);
 
    [[nodiscard]] bool is_empty(PartID pid) const;
 
    [[nodiscard]] std::vector<Chunk> extract(PartID pid);
 
    [[nodiscard]] bool empty() const;
 
    [[nodiscard]] std::string str() const;
 
    [[nodiscard]] std::size_t spill(BufferResource* br, std::size_t amount);
 
  private:
    // TODO: more fine-grained locking e.g. by locking each partition individually.
    mutable std::mutex mutex_;
    std::unordered_map<PartID, std::vector<Chunk>>
        pigeonhole_;  
};
 
inline std::ostream& operator<<(std::ostream& os, ReceivedChunks const& obj) {
    os << obj.str();
    return os;
}
 
}  // namespace rapidsmpf::shuffler::detail