delayed_submission.h

#pragma once
 
#include <deque>
#include <functional>
#include <memory>
#include <mutex>
#include <optional>
#include <set>
#include <stdexcept>
#include <string>
#include <utility>
#include <variant>
 
#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;

typedef uint64_t ItemIdType;

template <typename T>
class BaseDelayedSubmissionCollection {
 protected:
  std::string _name{"undefined"};  
  bool _enabled{true};    
  ItemIdType _itemId{0};  
  std::optional<ItemIdType> _processing{
    std::nullopt};  
  std::deque<std::pair<ItemIdType, T>> _collection{};  
  std::set<ItemIdType> _canceled{};                    
  std::mutex _mutex{};  

  virtual void scheduleLog(ItemIdType id, T item) = 0;

  virtual void processItem(ItemIdType id, T item) = 0;
 
 public:
  explicit BaseDelayedSubmissionCollection(std::string name, const bool enabled)
    : _name{std::move(name)}, _enabled{enabled}
  {
  }
 
  BaseDelayedSubmissionCollection()                                                  = delete;
  BaseDelayedSubmissionCollection(const BaseDelayedSubmissionCollection&)            = delete;
  BaseDelayedSubmissionCollection& operator=(BaseDelayedSubmissionCollection const&) = delete;
  BaseDelayedSubmissionCollection(BaseDelayedSubmissionCollection&& o)               = delete;
  BaseDelayedSubmissionCollection& operator=(BaseDelayedSubmissionCollection&& o)    = delete;

  [[nodiscard]] virtual ItemIdType schedule(T item)
  {
    if (!_enabled) throw std::runtime_error("Resource is disabled.");
 
    ItemIdType id;
    {
      std::lock_guard<std::mutex> lock(_mutex);
      id = _itemId++;
      _collection.emplace_back(id, item);
    }
    scheduleLog(id, item);
 
    return id;
  }

  void process()
  {
    // Process only those that were already inserted to prevent from never
    // returning if `_collection` grows indefinitely.
    size_t toProcess = 0;
    {
      std::lock_guard<std::mutex> lock(_mutex);
      toProcess = _collection.size();
    }
 
    for (size_t i = 0; i < toProcess; ++i) {
      std::pair<ItemIdType, T> item;
      {
        std::lock_guard<std::mutex> lock(_mutex);
        item = std::move(_collection.front());
        _collection.pop_front();
        if (_canceled.erase(item.first)) continue;
        _processing = std::optional<ItemIdType>{item.first};
      }
 
      processItem(item.first, item.second);
    }
 
    {
      // Clear the value of `_processing` as no more requests will be processed.
      std::lock_guard<std::mutex> lock(_mutex);
      _processing = std::nullopt;
    }
  }

  void cancel(ItemIdType id)
  {
    std::lock_guard<std::mutex> lock(_mutex);
    if (_processing.has_value() && _processing.value() == id)
      throw std::runtime_error("Cannot cancel, item is being processed.");
 
    _canceled.insert(id);
    ucxx_trace_req("Canceled item: %lu", id);
  }
};

class RequestDelayedSubmissionCollection
  : public BaseDelayedSubmissionCollection<
      std::pair<std::shared_ptr<Request>, DelayedSubmissionCallbackType>> {
 protected:
  void scheduleLog(
    ItemIdType id,
    std::pair<std::shared_ptr<Request>, DelayedSubmissionCallbackType> item) override;
 
  void processItem(
    ItemIdType id,
    std::pair<std::shared_ptr<Request>, DelayedSubmissionCallbackType> item) override;
 
 public:
  explicit RequestDelayedSubmissionCollection(std::string name, const bool enabled);
};

class GenericDelayedSubmissionCollection
  : public BaseDelayedSubmissionCollection<DelayedSubmissionCallbackType> {
 protected:
  void scheduleLog(ItemIdType id, DelayedSubmissionCallbackType item) override;
 
  void processItem(ItemIdType id, DelayedSubmissionCallbackType callback) override;
 
 public:
  explicit GenericDelayedSubmissionCollection(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);

  [[nodiscard]] ItemIdType registerGenericPre(DelayedSubmissionCallbackType callback);

  [[nodiscard]] ItemIdType registerGenericPost(DelayedSubmissionCallbackType callback);

  void cancelGenericPre(ItemIdType id);

  void cancelGenericPost(ItemIdType id);

  [[nodiscard]] bool isDelayedRequestSubmissionEnabled() const;
};
 
}  // namespace ucxx