worker.h

#pragma once
 
#include <atomic>
#include <functional>
#include <memory>
#include <mutex>
#include <queue>
#include <string>
#include <thread>
#include <utility>
#include <variant>
 
#include <ucp/api/ucp.h>
 
#include <ucxx/address_builder.h>
#include <ucxx/buffer.h>
#include <ucxx/component.h>
#include <ucxx/context.h>
#include <ucxx/delayed_submission.h>
#include <ucxx/endpoint_builder.h>
#include <ucxx/future.h>
#include <ucxx/inflight_requests.h>
#include <ucxx/listener_builder.h>
#include <ucxx/notifier.h>
#include <ucxx/request_flush_builder.h>
#include <ucxx/request_tag_builder.h>
#include <ucxx/typedefs.h>
#include <ucxx/worker_builder.h>
#include <ucxx/worker_progress_thread.h>
 
namespace ucxx {
 
class Component;
class Request;
 
class WorkerBuilder;
 
namespace detail {
void registerInflightRequest(std::shared_ptr<Component> component, std::shared_ptr<Request> req);
}  // namespace detail
 
class Address;
class Buffer;
class Endpoint;
class Listener;
class RequestAm;
 
namespace internal {
class AmData;
}  // namespace internal

class Worker : public Component {
 private:
  ucp_worker_h _handle{nullptr};        
  int _epollFileDescriptor{-1};         
  int _workerFileDescriptor{-1};        
  std::mutex _inflightRequestsMutex{};  
  std::unique_ptr<InflightRequests> _inflightRequests{
    std::make_unique<InflightRequests>()};  
  std::mutex
    _inflightRequestsToCancelMutex{};  
  std::unique_ptr<InflightRequests> _inflightRequestsToCancel{
    std::make_unique<InflightRequests>()};  
  WorkerProgressThread _progressThread{};   
  std::atomic<bool> _progressThreadRunning{false};      
  std::atomic<bool> _progressThreadPollingMode{false};  
  std::atomic<std::thread::id> _progressThreadId{std::thread::id()};  
  std::function<void(void*)> _progressThreadStartCallback{
    nullptr};  
  void* _progressThreadStartCallbackArg{
    nullptr};  
  std::shared_ptr<DelayedSubmissionCollection> _delayedSubmissionCollection{
    nullptr};  
 
  friend class detail::ConstructorFactory;
 
 protected:
  bool _enableFuture{
    false};  
  bool _enableRequestAttributes{
    false};  
  std::mutex _futuresPoolMutex{};  
  std::queue<std::shared_ptr<Future>>
    _futuresPool{};  
  std::shared_ptr<Notifier> _notifier{nullptr};  
  std::shared_ptr<internal::AmData>
    _amData;  
  BufferType _cudaBufferType{BufferType::Invalid};  
 
 private:
  void drainWorkerTagRecv();

  [[nodiscard]] std::shared_ptr<RequestAm> getAmRecv(
    ucp_ep_h ep, std::function<std::shared_ptr<RequestAm>()> createAmRecvRequestFunction);

  void stopProgressThreadNoWarn();

  [[nodiscard]] std::shared_ptr<Request> registerInflightRequest(std::shared_ptr<Request> request);

  bool progressPending();

  void setCudaBufferType(BufferType bufferType);
 
 protected:
  explicit Worker(std::shared_ptr<Context> context,
                  const bool enableDelayedSubmission = false,
                  const bool enableFuture            = false);
 
 public:
  Worker()                         = delete;
  Worker(const Worker&)            = delete;
  Worker& operator=(Worker const&) = delete;
  Worker(Worker&& o)               = delete;
  Worker& operator=(Worker&& o)    = delete;

  friend class WorkerBuilder;

  friend void detail::registerInflightRequest(std::shared_ptr<Component> component,
                                              std::shared_ptr<Request> req);

  virtual ~Worker();

  [[nodiscard]] ucp_worker_h getHandle();

  [[nodiscard]] std::string getInfo();

  void initBlockingProgressMode();

  int getEpollFileDescriptor();

  bool arm();

  bool progressWorkerEvent(const int epollTimeout = -1);

  void signal();

  bool waitProgress();

  bool progressOnce();

  bool progress();

  void registerDelayedSubmission(std::shared_ptr<Request> request,
                                 DelayedSubmissionCallbackType callback);

  [[nodiscard]] bool registerGenericPre(DelayedSubmissionCallbackType callback,
                                        uint64_t period = 0);

  [[nodiscard]] bool registerGenericPost(DelayedSubmissionCallbackType callback,
                                         uint64_t period = 0);

  [[nodiscard]] bool isDelayedRequestSubmissionEnabled() const;

  [[nodiscard]] bool isFutureEnabled() const;

  [[nodiscard]] bool isRequestAttributesEnabled() const noexcept;

  [[nodiscard]] BufferType getCudaBufferType() const;

  virtual void populateFuturesPool();

  virtual void clearFuturesPool();

  [[nodiscard]] virtual std::shared_ptr<Future> getFuture();

  [[nodiscard]] virtual RequestNotifierWaitState waitRequestNotifier(uint64_t periodNs);

  virtual void runRequestNotifier();

  virtual void stopRequestNotifierThread();

  void setProgressThreadStartCallback(std::function<void(void*)> callback, void* callbackArg);

  void startProgressThread(const bool pollingMode = false, const int epollTimeout = 1);

  void stopProgressThread();

  [[nodiscard]] bool isProgressThreadRunning();

  [[nodiscard]] std::thread::id getProgressThreadId();

  size_t cancelInflightRequests(uint64_t period = 0, uint64_t maxAttempts = 1);

  void scheduleRequestCancel(TrackedRequests trackedRequests);

  void removeInflightRequest(std::shared_ptr<Request> request);

  [[nodiscard]] std::shared_ptr<TagProbeInfo> tagProbe(const Tag tag,
                                                       const TagMask tagMask = TagMaskFull,
                                                       const bool remove     = false) const;

  [[nodiscard]] RequestTagBuilder tagRecvBuilder(void* buffer,
                                                 size_t length,
                                                 Tag tag,
                                                 TagMask tagMask);

  [[nodiscard]] RequestTagBuilder tagRecvWithHandleBuilder(void* buffer,
                                                           std::shared_ptr<TagProbeInfo> probeInfo);

  [[nodiscard]] RequestTagBuilder tagRecvWithHandleBuilder(void* buffer,
                                                           size_t length,
                                                           std::shared_ptr<TagProbeInfo> probeInfo);

  [[nodiscard]] AddressBuilder addressBuilder();

  [[nodiscard]] EndpointBuilder endpointBuilder(std::string ipAddress, uint16_t port);

  [[nodiscard]] EndpointBuilder endpointBuilder(std::shared_ptr<Address> address);

  [[nodiscard]] ListenerBuilder listenerBuilder(uint16_t port,
                                                ucp_listener_conn_callback_t callback,
                                                void* callbackArgs);

  void registerAmAllocator(ucs_memory_type_t memoryType, AmAllocatorType allocator);

  void registerAmReceiverCallback(AmReceiverCallbackInfo info, AmReceiverCallbackType callback);

  [[nodiscard]] bool amProbe(const ucp_ep_h endpointHandle) const;

  [[nodiscard]] RequestFlushBuilder flushBuilder();

  struct Attributes {
    ucs_thread_mode_t threadMode{UCS_THREAD_MODE_MULTI};
    size_t maxAmHeader{0};
    std::string name{};
    size_t maxDebugString{0};
  };

  [[nodiscard]] Attributes queryAttributes() const;
};
 
}  // namespace ucxx