endpoint.h

#pragma once
 
#include <netdb.h>
 
#include <memory>
#include <string>
#include <vector>
 
#include <ucp/api/ucp.h>
 
#include <ucxx/address.h>
#include <ucxx/component.h>
#include <ucxx/endpoint_builder.h>
#include <ucxx/exception.h>
#include <ucxx/inflight_requests.h>
#include <ucxx/listener.h>
#include <ucxx/remote_key_builder.h>
#include <ucxx/request.h>
#include <ucxx/request_am_builder.h>
#include <ucxx/request_endpoint_close_builder.h>
#include <ucxx/request_flush_builder.h>
#include <ucxx/request_mem_builder.h>
#include <ucxx/request_stream_builder.h>
#include <ucxx/request_tag_builder.h>
#include <ucxx/request_tag_multi_builder.h>
#include <ucxx/typedefs.h>
#include <ucxx/utils/sockaddr.h>
#include <ucxx/worker.h>
 
namespace ucxx {
 
class Component;
class Request;
class RequestEndpointClose;
 
class RequestEndpointCloseBuilder;
 
namespace detail {
void registerInflightRequest(std::shared_ptr<Component> component, std::shared_ptr<Request> req);
}  // namespace detail

struct EpParamsDeleter {
  void operator()(ucp_ep_params_t* ptr);
};

class Endpoint : public Component {
 private:
  ucp_ep_h _handle{nullptr};          
  ucp_ep_h _originalHandle{nullptr};  
  bool _endpointErrorHandling{true};  
  std::unique_ptr<InflightRequests> _inflightRequests{
    std::make_unique<InflightRequests>()};  
  std::mutex _mutex{};  
  ucs_status_t _status{UCS_INPROGRESS};  
  std::atomic<bool> _closing{false};     
  EndpointCloseCallbackUserFunction _closeCallback{nullptr};  
  EndpointCloseCallbackUserData _closeCallbackArg{
    nullptr};  

  Endpoint(std::shared_ptr<Component> workerOrListener, bool endpointErrorHandling);

  void create(ucp_ep_params_t* params);

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

  [[nodiscard]] std::shared_ptr<RequestEndpointClose> closeRequest(
    const bool force,
    const bool enablePythonFuture,
    EndpointCloseCallbackUserFunction callbackFunction,
    EndpointCloseCallbackUserData callbackData);

  friend void endpointErrorCallback(void* arg, ucp_ep_h ep, ucs_status_t status);

  friend class RequestEndpointCloseBuilder;

  friend void detail::registerInflightRequest(std::shared_ptr<Component> component,
                                              std::shared_ptr<Request> req);
 
 public:
  Endpoint()                           = delete;
  Endpoint(const Endpoint&)            = delete;
  Endpoint& operator=(Endpoint const&) = delete;
  Endpoint(Endpoint&& o)               = delete;
  Endpoint& operator=(Endpoint&& o)    = delete;
 
  ~Endpoint();
 
  // Allow internal construction without exposing factory functions.
  friend class detail::ConstructorFactory;

  [[nodiscard]] ucp_ep_h getHandle();

  [[nodiscard]] bool isAlive() const;

  void raiseOnError();

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

  size_t cancelInflightRequests();

  [[nodiscard]] size_t getCancelingSize() const;

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

  void setCloseCallback(EndpointCloseCallbackUserFunction closeCallback,
                        EndpointCloseCallbackUserData closeCallbackArg);

  [[nodiscard]] RequestAmBuilder amSendBuilder(const void* const buffer,
                                               const size_t length,
                                               const ucs_memory_type_t memoryType);

  [[nodiscard]] RequestAmBuilder amSendBuilder(const void* const buffer,
                                               const size_t length,
                                               const AmSendParams& params);

  [[nodiscard]] RequestAmBuilder amSendBuilder(std::vector<ucp_dt_iov_t> iov,
                                               const AmSendParams& params);

  [[nodiscard]] RequestAmBuilder amRecvBuilder();

  [[nodiscard]] RequestMemBuilder memPutBuilder(const void* const buffer,
                                                size_t length,
                                                uint64_t remoteAddr,
                                                ucp_rkey_h rkey);

  [[nodiscard]] RequestMemBuilder memPutBuilder(const void* const buffer,
                                                size_t length,
                                                std::shared_ptr<ucxx::RemoteKey> remoteKey);

  [[nodiscard]] RequestMemBuilder memGetBuilder(void* buffer,
                                                size_t length,
                                                uint64_t remoteAddr,
                                                ucp_rkey_h rkey);

  [[nodiscard]] RequestMemBuilder memGetBuilder(void* buffer,
                                                size_t length,
                                                std::shared_ptr<ucxx::RemoteKey> remoteKey);

  [[nodiscard]] RequestStreamBuilder streamSendBuilder(const void* const buffer, size_t length);

  [[nodiscard]] RequestStreamBuilder streamRecvBuilder(void* buffer, size_t length);

  [[nodiscard]] RequestTagBuilder tagSendBuilder(const void* const buffer, size_t length, Tag tag);

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

  [[nodiscard]] RequestTagMultiBuilder tagMultiSendBuilder(const std::vector<const void*>& buffer,
                                                           const std::vector<size_t>& size,
                                                           const std::vector<int>& isCUDA,
                                                           const Tag tag);

  [[nodiscard]] RequestTagMultiBuilder tagMultiRecvBuilder(const Tag tag, const TagMask tagMask);

  [[nodiscard]] RequestFlushBuilder flushBuilder();

  [[nodiscard]] std::shared_ptr<Worker> getWorker();

  [[nodiscard]] RemoteKeyBuilder remoteKeyBuilder(SerializedRemoteKey serializedRemoteKey);

  [[nodiscard]] RequestEndpointCloseBuilder closeBuilder();

  void closeBlocking(uint64_t period = 0, uint64_t maxAttempts = 1);
};
 
}  // namespace ucxx