delayed_submission.h

 
#pragma once
 
#include <functional>
#include <memory>
#include <mutex>
#include <optional>
#include <stdexcept>
#include <string>
#include <utility>
#include <variant>
#include <vector>
 
#include <ucp/api/ucp.h>
#include <ucs/memory/memory_type.h>
 
#include <ucxx/log.h>
#include <ucxx/request_data.h>
 
namespace ucxx {
 
typedef std::function<void()> DelayedSubmissionCallbackType;
 
template <typename T>
class BaseDelayedSubmissionCollection {
 protected:
  std::string _name{"undefined"};  
  bool _enabled{true};  
  std::vector<T> _collection{};  
  std::mutex _mutex{};           
 
  virtual void scheduleLog(T item) = 0;
 
  virtual void processItem(T item) = 0;
 
 public:
  explicit BaseDelayedSubmissionCollection(const std::string name, const bool enabled)
    : _name{name}, _enabled{enabled}
  {
  }
 
  BaseDelayedSubmissionCollection()                                                  = delete;
  BaseDelayedSubmissionCollection(const BaseDelayedSubmissionCollection&)            = delete;
  BaseDelayedSubmissionCollection& operator=(BaseDelayedSubmissionCollection const&) = delete;
  BaseDelayedSubmissionCollection(BaseDelayedSubmissionCollection&& o)               = delete;
  BaseDelayedSubmissionCollection& operator=(BaseDelayedSubmissionCollection&& o)    = delete;
 
  virtual void schedule(T item)
  {
    if (!_enabled) throw std::runtime_error("Resource is disabled.");
 
    {
      std::lock_guard<std::mutex> lock(_mutex);
      _collection.push_back(item);
    }
    scheduleLog(item);
  }
 
  void process()
  {
    decltype(_collection) itemsToProcess;
    {
      std::lock_guard<std::mutex> lock(_mutex);
      // Move _collection to a local copy in order to to hold the lock for as
      // short as possible
      itemsToProcess = std::move(_collection);
    }
 
    if (itemsToProcess.size() > 0) {
      ucxx_trace_req("Submitting %lu %s callbacks", itemsToProcess.size(), _name.c_str());
      for (auto& item : itemsToProcess)
        processItem(item);
    }
  }
};
 
class RequestDelayedSubmissionCollection
  : public BaseDelayedSubmissionCollection<
      std::pair<std::shared_ptr<Request>, DelayedSubmissionCallbackType>> {
 protected:
  void scheduleLog(
    std::pair<std::shared_ptr<Request>, DelayedSubmissionCallbackType> item) override;
 
  void processItem(
    std::pair<std::shared_ptr<Request>, DelayedSubmissionCallbackType> item) override;
 
 public:
  explicit RequestDelayedSubmissionCollection(const std::string name, const bool enabled);
};
 
class GenericDelayedSubmissionCollection
  : public BaseDelayedSubmissionCollection<DelayedSubmissionCallbackType> {
 protected:
  void scheduleLog(DelayedSubmissionCallbackType item) override;
 
  void processItem(DelayedSubmissionCallbackType callback) override;
 
 public:
  explicit GenericDelayedSubmissionCollection(const std::string name);
};
 
class DelayedSubmissionCollection {
 private:
  GenericDelayedSubmissionCollection _genericPre{
    "generic pre"};  
  GenericDelayedSubmissionCollection _genericPost{
    "generic post"};  
  RequestDelayedSubmissionCollection _requests{
    "request", false};  
  bool _enableDelayedRequestSubmission{false};
 
 public:
  explicit DelayedSubmissionCollection(bool enableDelayedRequestSubmission = false);
 
  DelayedSubmissionCollection()                                              = delete;
  DelayedSubmissionCollection(const DelayedSubmissionCollection&)            = delete;
  DelayedSubmissionCollection& operator=(DelayedSubmissionCollection const&) = delete;
  DelayedSubmissionCollection(DelayedSubmissionCollection&& o)               = delete;
  DelayedSubmissionCollection& operator=(DelayedSubmissionCollection&& o)    = delete;
 
  void processPre();
 
  void processPost();
 
  void registerRequest(std::shared_ptr<Request> request, DelayedSubmissionCallbackType callback);
 
  void registerGenericPre(DelayedSubmissionCallbackType callback);
 
  void registerGenericPost(DelayedSubmissionCallbackType callback);
 
  bool isDelayedRequestSubmissionEnabled() const;
};
 
}  // namespace ucxx