tag_probe.h

 
#pragma once
 
#include <atomic>
#include <memory>
#include <optional>
 
#include <ucp/api/ucp.h>
#include <ucxx/typedefs.h>
 
namespace ucxx {
 
// Forward declaration for friend access
class RequestTag;
 
class TagRecvInfo {
 public:
  const Tag senderTag;  
  const size_t length;  
 
  TagRecvInfo() : senderTag(Tag(0)), length(0) {}
 
  explicit TagRecvInfo(const ucp_tag_recv_info_t& info);
};
 
class TagProbeInfo {
 public:
  TagProbeInfo(const TagProbeInfo&)            = delete;
  TagProbeInfo& operator=(TagProbeInfo const&) = delete;
  TagProbeInfo(TagProbeInfo&& o)               = delete;
  TagProbeInfo& operator=(TagProbeInfo&& o)    = delete;
 
  ~TagProbeInfo();
 
  bool isMatched() const;
 
  const TagRecvInfo& getInfo() const;
 
  ucp_tag_message_h getHandle() const;
 
  ucp_tag_message_h releaseHandle() const;
 
 private:
  const bool _matched{false};  
  const std::optional<TagRecvInfo> _info{
    std::nullopt};  
  const ucp_tag_message_h _handle{
    nullptr};  
  mutable std::atomic<bool> _consumed{false};  
 
  void consume() const;
 
  TagProbeInfo() = default;
 
  TagProbeInfo(const ucp_tag_recv_info_t& info, ucp_tag_message_h handle);
 
  friend std::shared_ptr<TagProbeInfo> createTagProbeInfo();
 
  friend std::shared_ptr<TagProbeInfo> createTagProbeInfo(const ucp_tag_recv_info_t& info,
                                                          ucp_tag_message_h handle);
 
  // Allow RequestTag to consume the handle
  friend class RequestTag;
};
 
}  // namespace ucxx