request_tag_multi.h

#pragma once
 
#include <memory>
#include <mutex>
#include <string>
#include <vector>
 
#include <ucp/api/ucp.h>
 
#include <ucxx/buffer.h>
#include <ucxx/endpoint.h>
#include <ucxx/future.h>
#include <ucxx/request.h>
#include <ucxx/request_tag_multi_builder.h>
 
namespace ucxx {
 
class RequestTagMulti;

struct BufferRequest {
  std::shared_ptr<Request> request{nullptr};  
  std::shared_ptr<std::string> stringBuffer{nullptr};  
  std::shared_ptr<Buffer> buffer{nullptr};  
 
  BufferRequest();
  ~BufferRequest();
 
  BufferRequest(const BufferRequest&)            = delete;
  BufferRequest& operator=(BufferRequest const&) = delete;
  BufferRequest(BufferRequest&& o)               = delete;
  BufferRequest& operator=(BufferRequest&& o)    = delete;
};

typedef std::shared_ptr<BufferRequest> BufferRequestPtr;

class RequestTagMulti : public Request {
 private:
  size_t _totalRequests{0};  
  std::mutex
    _completedRequestsMutex{};   
  size_t _completedRequests{0};  
  ucs_status_t _finalStatus{
    UCS_OK};  
 
 public:
  std::vector<BufferRequestPtr> _bufferRequests{};  
  bool _isFilled{false};                            
 
 private:
  RequestTagMulti()                                  = delete;
  RequestTagMulti(const RequestTagMulti&)            = delete;
  RequestTagMulti& operator=(RequestTagMulti const&) = delete;
  RequestTagMulti(RequestTagMulti&& o)               = delete;
  RequestTagMulti& operator=(RequestTagMulti&& o)    = delete;

  RequestTagMulti(std::shared_ptr<Endpoint> endpoint,
                  const std::variant<data::TagMultiSend, data::TagMultiReceive> requestData,
                  std::string operationName,
                  const bool enablePythonFuture);
 
  void populateDelayedSubmissionImpl() override;

  void recvFrames();

  void recvHeader();

  void send();
 
 public:
  // Allow internal construction without exposing factory functions.
  friend class detail::ConstructorFactory;

  virtual ~RequestTagMulti();

  void markCompleted(ucs_status_t status);

  void recvCallback(ucs_status_t status);
 
  void cancel() override;
};

typedef std::shared_ptr<RequestTagMulti> RequestTagMultiPtr;
 
}  // namespace ucxx