spill_manager.hpp

 
#pragma once
 
#include <map>
#include <mutex>
#include <optional>
 
#include <rapidsmpf/pausable_thread_loop.hpp>
#include <rapidsmpf/utils.hpp>
 
namespace rapidsmpf {
 
class BufferResource;
 
class SpillManager {
  public:
    using SpillFunction = std::function<std::size_t(std::size_t)>;
 
    using SpillFunctionID = std::size_t;
 
    SpillManager(
        BufferResource* br, std::optional<Duration> periodic_spill_check = std::nullopt
    );
 
    ~SpillManager();
 
    SpillFunctionID add_spill_function(SpillFunction spill_function, int priority);
 
    void remove_spill_function(SpillFunctionID fid);
 
    std::size_t spill(std::size_t amount);
 
    std::size_t spill_to_make_headroom(std::int64_t headroom = 0);
 
  private:
    mutable std::mutex mutex_;
    BufferResource* br_;
    std::size_t spill_function_id_counter_{0};
    std::map<SpillFunctionID, SpillFunction> spill_functions_;
    std::multimap<int, SpillFunctionID, std::greater<>> spill_function_priorities_;
    std::optional<detail::PausableThreadLoop> periodic_spill_thread_;
};
 
 
}  // namespace rapidsmpf