progress_thread.hpp

/*
 * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
 * SPDX-License-Identifier: Apache-2.0
 */
#pragma once
 
#include <chrono>
#include <condition_variable>
#include <cstdint>
#include <mutex>
#include <unordered_map>
 
#include <rapidsmpf/communicator/communicator.hpp>
#include <rapidsmpf/pausable_thread_loop.hpp>
#include <rapidsmpf/statistics.hpp>
 
namespace rapidsmpf {
 
class ProgressThread {
  public:
    enum ProgressState : bool {
        InProgress,
        Done,
    };
 
    using FunctionIndex = std::uint64_t;
 
    using ProgressThreadAddress = std::uintptr_t;
 
    struct FunctionID {
        ProgressThreadAddress thread_address{
            ProgressThreadAddress(0)
        };  
        FunctionIndex function_index{0};  
 
        FunctionID() = default;
 
        constexpr FunctionID(ProgressThreadAddress thread_addr, FunctionIndex index)
            : thread_address(thread_addr), function_index(index) {}
 
        [[nodiscard]] constexpr bool is_valid() const {
            return thread_address != ProgressThreadAddress(0);
        }
    };
 
    using Function = std::function<ProgressState()>;
 
    class FunctionState {
      public:
        explicit FunctionState(Function&& function);
 
        void operator()();
 
        Function function;  
        bool is_done{false};  
    };
 
    ProgressThread(
        Communicator::Logger& logger,
        std::shared_ptr<Statistics> statistics = Statistics::disabled(),
        Duration sleep = std::chrono::microseconds{1}
    );
 
    ~ProgressThread();
 
    void stop();
 
    FunctionID add_function(Function&& function);
 
    void remove_function(FunctionID function_id);
 
    void pause();
 
    void resume();
 
    bool is_running() const;
 
  private:
    void event_loop();
 
    detail::PausableThreadLoop thread_;
    Communicator::Logger& logger_;
    std::shared_ptr<Statistics> statistics_;
    bool is_thread_initialized_{false};
    bool active_{false};
    mutable std::mutex mutex_;
    std::condition_variable cv_;
    FunctionIndex next_function_id_{0};
    std::unordered_map<FunctionIndex, FunctionState> functions_;
};
 
}  // namespace rapidsmpf