request_am.h

#include <functional>
#include <map>
#include <memory>
#include <mutex>
#include <queue>
#include <string>
#include <unordered_map>
#include <vector>
 
#include <ucp/api/ucp.h>
 
#include <ucxx/typedefs.h>
 
namespace ucxx {
 
class Buffer;
class InflightRequests;
class RequestAm;
class Request;
class Worker;
 
namespace internal {
 
class AmData;

class RecvAmMessage {
 public:
  internal::AmData* _amData{nullptr};  
  ucp_ep_h _ep{nullptr};               
  std::shared_ptr<RequestAm> _request{
    nullptr};  
  std::shared_ptr<Buffer> _buffer{nullptr};  
 
  RecvAmMessage()                                = delete;
  RecvAmMessage(const RecvAmMessage&)            = delete;
  RecvAmMessage& operator=(RecvAmMessage const&) = delete;
  RecvAmMessage(RecvAmMessage&& o)               = delete;
  RecvAmMessage& operator=(RecvAmMessage&& o)    = delete;

  RecvAmMessage(internal::AmData* amData,
                ucp_ep_h ep,
                std::shared_ptr<RequestAm> request,
                std::shared_ptr<Buffer> buffer,
                AmReceiverCallbackType receiverCallback = AmReceiverCallbackType(),
                std::vector<std::byte> userHeader       = {});

  void callback(void* request, ucs_status_t status);
};
 
typedef std::unordered_map<ucp_ep_h, std::queue<std::shared_ptr<RequestAm>>> AmPoolType;
typedef std::map<std::shared_ptr<RequestAm>,
                 std::shared_ptr<RecvAmMessage>,
                 std::owner_less<std::shared_ptr<RequestAm>>>
  RecvAmMessageMapType;
 
typedef std::unordered_map<AmReceiverCallbackIdType, AmReceiverCallbackType>
  AmReceiverCallbackMapType;
typedef std::
  unordered_map<AmReceiverCallbackOwnerType, AmReceiverCallbackMapType, AmReceiverCallbackOwnerHash>
    AmReceiverCallbackOwnerMapType;

class AmData {
 public:
  std::weak_ptr<Worker> _worker{};  
  std::string _ownerString{};       
  AmPoolType _recvPool{};  
  AmPoolType _recvWait{};  
  RecvAmMessageMapType
    _recvAmMessageMap{};  
  AmReceiverCallbackOwnerMapType
    _receiverCallbacks{};  
  std::mutex _mutex{};     
  std::function<void(std::shared_ptr<Request>)>
    _registerInflightRequest{};  
  std::unordered_map<ucs_memory_type_t, AmAllocatorType>
    _allocators{};  
};
 
}  // namespace internal
 
}  // namespace ucxx