moving to a single callback

Signed-off-by: niranda perera <niranda.perera@gmail.com>

Author: nirandaperera

cpp/include/rapidsmpf/shuffler/finish_counter.hpp

@@ -4,13 +4,13 @@
  */
 #pragma once
 
+#include <array>
 #include <chrono>
 #include <condition_variable>
 #include <functional>
 #include <mutex>
 #include <optional>
 #include <unordered_map>
-#include <unordered_set>
 #include <vector>
 
 #include <rapidsmpf/communicator/communicator.hpp>
@@ -44,13 +44,43 @@ namespace detail {
  */
 class FinishCounter {
   public:
+    /**
+     * @brief Callback function type called when a partition is finished.
+     *
+     * The callback receives the partition ID of the finished partition.
+     *
+     * @warning A callback must be fast and non-blocking and should not call any of the
+     * `wait*` methods. And be very careful if acquiring locks. Ideally it should be used
+     * to signal a separate thread to do the actual processing (eg. WaitHand).
+     *
+     * @note When a callback is registered, it will be identified by the
+     * FinishedCbId returned. So, if a callback needs to be preemptively canceled,
+     * the corresponding identifier needs to be provided.
+     *
+     * @note Every callback will be called as and when each partition is finished. If
+     * there were finished partitions before the callback was registered, the callback
+     * will be called for them immediately by the caller thread. Else, the callback will
+     * be called by the progress thread (Therefore, it will be called
+     * `n_local_partitions_` times in total).
+     *
+     * @note Caller needs to be careful when using both callbacks and wait* methods
+     * together.
+     */
+    using FinishedCallback = std::function<void(PartID)>;
+
     /**
      * @brief Construct a finish counter.
      *
      * @param nranks The total number of ranks participating in the shuffle.
      * @param local_partitions The partition IDs local to the current rank.
+     * @param finished_callback The callback to notify when a partition is finished
+     * (optional).
      */
-    FinishCounter(Rank nranks, std::vector<PartID> const& local_partitions);
+    FinishCounter(
+        Rank nranks,
+        std::vector<PartID> const& local_partitions,
+        FinishedCallback&& finished_callback = nullptr
+    );
 
     ~FinishCounter();
 
@@ -88,77 +118,6 @@ class FinishCounter {
      */
     [[nodiscard]] bool all_finished() const;
 
-    /**
-     * @brief Returns whether a partition is finished (non-blocking).
-     *
-     * @param pid The partition ID to check.
-     * @return True if the partition is finished, otherwise False.
-     */
-    [[nodiscard]] bool is_finished(PartID pid) const;
-
-    /**
-     * @brief Callback function type called when a partition is finished.
-     *
-     * The callback receives the partition ID of the finished partition.
-     *
-     * @warning A callback must be fast and non-blocking and should not call any of the
-     * `wait*` methods. And be very careful if acquiring locks. Ideally it should be used
-     * to signal a separate thread to do the actual processing (eg. WaitHand).
-     *
-     * @note When a callback is registered, it will be identified by the
-     * FinishedCbId returned. So, if a callback needs to be preemptively canceled,
-     * the corresponding identifier needs to be provided.
-     *
-     * @note Every callback will be called as and when each partition is finished. If
-     * there were finished partitions before the callback was registered, the callback
-     * will be called for them immediately by the caller thread. Else, the callback will
-     * be called by the progress thread (Therefore, it will be called
-     * `n_local_partitions_` times in total).
-     *
-     * @note Caller needs to be careful when using both callbacks and wait* methods
-     * together.
-     */
-    using FinishedCallback = std::function<void(PartID)>;
-
-    /**
-     * @brief Type used to identify callbacks.
-     */
-    using FinishedCbId = size_t;
-
-    /**
-     * @brief Register a callback to be notified when any partition is finished.
-     *
-     * This function registers a callback that will be called when a partition is finished
-     * (and for all currently finished partitions). The callback receives partition IDs as
-     * they complete. If all partitions are already finished, the callback is executed
-     * immediately for all partitions and invalid_cb_id is returned.
-     *
-     * @param cb The callback to invoke when partitions are finished.
-     *
-     * @return A unique callback ID that can be used to cancel the callback, or
-     * invalid_cb_id if the callback was executed immediately.
-     */
-    FinishedCbId register_finished_callback(FinishedCallback&& cb);
-
-    /**
-     * @brief Special constant indicating an invalid or immediately-executed callback ID.
-     *
-     * This value is returned by register_finished_callback when the callback is executed
-     * immediately (e.g., when all partitions are already finished).
-     */
-    static constexpr FinishedCbId invalid_cb_id =
-        std::numeric_limits<FinishedCbId>::max();
-
-    /**
-     * @brief Cancel a previously registered callback.
-     *
-     * This function removes a callback registered with register_finished_callback using
-     * its ID. It is safe to call this with invalid_cb_id or an already-cancelled ID.
-     *
-     * @param callback_id callback ID.
-     */
-    void remove_finished_callback(FinishedCbId callback_id);
-
     /**
      * @brief Returns the partition ID of a finished partition that hasn't been waited on
      * (blocking). Optionally a timeout (in ms) can be provided.
@@ -250,30 +209,32 @@ class FinishCounter {
     // when all ranks has reported their goal that the goalpost is final.
     std::unordered_map<PartID, PartitionInfo> goalposts_;
 
-    std::vector<PartID> finished_partitions_{};  ///< partition IDs of finished partitions
+    // mutex to control access between the progress thread and the caller thread on shared
+    // resources
+    mutable std::mutex mutex_;  // TODO: use a shared_mutex lock?
 
-    std::mutex finished_cbs_mutex_;  ///< mutex to protect the finished_cbs_ and
-                                     ///< next_finished_cb_id_
+    ///@brief Handler to implement the wait* methods using callbacks
+    struct WaitHandler {
+        std::unordered_set<PartID>
+            to_wait{};  ///< finished partitions available to wait on
+        bool active{true};
+        std::condition_variable cv;
+        std::mutex mutex;
 
-    struct CallbackContainer {
-        FinishedCbId cb_id;  ///< callback ID to identify the callback
+        ~WaitHandler();
 
-        // index of the next partition that the callback is interested in. cb will
-        // called from next_pid_idx to end of finished_partitions_
-        size_t next_pid_idx;
+        // Callback to listen on the finished partitions
+        void on_finished_cb(PartID pid);
 
-        FinishedCallback cb;
-    };
+        PartID wait_any(std::optional<std::chrono::milliseconds> timeout);
 
-    std::vector<CallbackContainer> finished_cbs_{};
-    FinishedCbId next_finished_cb_id_{0};  ///< next callback ID to assign
+        void wait_on(PartID pid, std::optional<std::chrono::milliseconds> timeout);
+    };
 
-    // mutex to control access between the progress thread and the caller thread on shared
-    // resources
-    mutable std::mutex mutex_;  // TODO: use a shared_mutex lock?
+    WaitHandler wait_handler_{};
 
-    class WaitHandler;  ///< Handler to implement the wait* methods using callbacks
-    std::unique_ptr<WaitHandler> wait_handler_;
+    FinishedCallback finished_callback_ =
+        nullptr;  ///< callback to notify when a partition is finished
 };
 
 }  // namespace detail

cpp/include/rapidsmpf/shuffler/shuffler.hpp

@@ -77,6 +77,9 @@ class Shuffler {
         PartitionOwner partition_owner
     );
 
+    /// @copydoc detail::FinishCounter::FinishedCallback
+    using FinishedCallback = detail::FinishCounter::FinishedCallback;
+
     /**
      * @brief Construct a new shuffler for a single shuffle.
      *
@@ -87,6 +90,7 @@ class Shuffler {
      * @param total_num_partitions Total number of partitions in the shuffle.
      * @param stream The CUDA stream for memory operations.
      * @param br Buffer resource used to allocate temporary and the shuffle result.
+     * @param finished_callback Callback to notify when a partition is finished.
      * @param statistics The statistics instance to use (disabled by default).
      * @param partition_owner Function to determine partition ownership.
      */
@@ -97,10 +101,46 @@ class Shuffler {
         PartID total_num_partitions,
         rmm::cuda_stream_view stream,
         BufferResource* br,
+        FinishedCallback&& finished_callback,
         std::shared_ptr<Statistics> statistics = Statistics::disabled(),
         PartitionOwner partition_owner = round_robin
     );
 
+    /**
+     * @brief Construct a new shuffler for a single shuffle.
+     *
+     * @param comm The communicator to use.
+     * @param progress_thread The progress thread to use.
+     * @param op_id The operation ID of the shuffle. This ID is unique for this operation,
+     * and should not be reused until all nodes has called `Shuffler::shutdown()`.
+     * @param total_num_partitions Total number of partitions in the shuffle.
+     * @param stream The CUDA stream for memory operations.
+     * @param br Buffer resource used to allocate temporary and the shuffle result.
+     * @param statistics The statistics instance to use (disabled by default).
+     * @param partition_owner Function to determine partition ownership.
+     */
+    Shuffler(
+        std::shared_ptr<Communicator> comm,
+        std::shared_ptr<ProgressThread> progress_thread,
+        OpID op_id,
+        PartID total_num_partitions,
+        rmm::cuda_stream_view stream,
+        BufferResource* br,
+        std::shared_ptr<Statistics> statistics = Statistics::disabled(),
+        PartitionOwner partition_owner = round_robin
+    )
+        : Shuffler(
+              comm,
+              progress_thread,
+              op_id,
+              total_num_partitions,
+              stream,
+              br,
+              nullptr,
+              statistics,
+              partition_owner
+          ) {}
+
     ~Shuffler();
 
     /**
@@ -170,18 +210,6 @@ class Shuffler {
      */
     [[nodiscard]] bool is_finished(PartID pid) const;
 
-    /// @copydoc detail::FinishCounter::FinishedCallback
-    using FinishedCallback = detail::FinishCounter::FinishedCallback;
-
-    /// @copydoc detail::FinishCounter::FinishedCbId
-    using FinishedCbId = detail::FinishCounter::FinishedCbId;
-
-    /// @copydoc detail::FinishCounter::register_finished_callback
-    FinishedCbId register_finished_callback(FinishedCallback&& cb);
-
-    /// @copydoc detail::FinishCounter::remove_finished_callback
-    void remove_finished_callback(FinishedCbId callback_id);
-
     /**
      * @brief Wait for any partition to finish.
      *