finish_counter.hpp

 
#pragma once
 
#include <mutex>
#include <vector>
 
#include <rapidsmpf/communicator/communicator.hpp>
#include <rapidsmpf/shuffler/chunk.hpp>
 
namespace rapidsmpf::shuffler {
 
namespace detail {
 
class FinishCounter {
  public:
    FinishCounter(Rank nranks, PartID n_local_partitions);
 
    ~FinishCounter() = default;
 
    void move_goalpost(Rank src_rank, ChunkID nchunks);
 
    void add_finished_chunk();
 
    [[nodiscard]] bool all_finished() const;
 
    [[nodiscard]] std::string str() const;
 
  private:
    Rank const nranks_;
 
    Rank n_ranks_with_goalpost_{0};  
    ChunkID total_chunk_goal_{0};  
    ChunkID total_finished_chunks_{0};  
    std::vector<bool> rank_reported_;  
 
    mutable std::mutex mutex_;
};
 
}  // namespace detail
 
inline std::ostream& operator<<(std::ostream& os, detail::FinishCounter const& obj) {
    os << obj.str();
    return os;
}
 
}  // namespace rapidsmpf::shuffler