text: /**
   * Copyright (C) Mellanox Technologies Ltd. 2001-2020.  ALL RIGHTS RESERVED.
   * Copyright (c) UT-Battelle, LLC. 2015-2017. ALL RIGHTS RESERVED.
   * Copyright (C) Los Alamos National Security, LLC. 2019 ALL RIGHTS RESERVED.
   *
   * See file LICENSE for terms.
   */
  
  
  #ifndef UCP_REQUEST_H_
  #define UCP_REQUEST_H_
  
  #include "ucp_context.h"
  #include "ucp_mm.h"
  
  #include <ucp/api/ucp.h>
  #include <ucp/dt/datatype_iter.h>
  #include <uct/api/uct.h>
  #include <ucs/datastruct/mpool.h>
  #include <ucs/datastruct/queue_types.h>
  #include <ucs/debug/assert.h>
  #include <ucp/dt/dt.h>
  #include <ucp/rma/rma.h>
  #include <ucp/wireup/wireup.h>
  #include <ucp/core/ucp_am.h>
  #include <ucp/core/ucp_worker.h>
  
  
  #define ucp_trace_req(_sreq, _message, ...) \
      ucs_trace_req("req %p: " _message, (_sreq), ## __VA_ARGS__)
  
  
  /**
   * Request flags
   */
  enum {
      UCP_REQUEST_FLAG_COMPLETED             = UCS_BIT(0),
      UCP_REQUEST_FLAG_RELEASED              = UCS_BIT(1),
      /* UCS_BIT(2) and UCS_BIT(3) are vacant flags */
      UCP_REQUEST_FLAG_SYNC_LOCAL_COMPLETED  = UCS_BIT(4),
      UCP_REQUEST_FLAG_SYNC_REMOTE_COMPLETED = UCS_BIT(5),
      UCP_REQUEST_FLAG_CALLBACK              = UCS_BIT(6),
      UCP_REQUEST_FLAG_PROTO_INITIALIZED     = UCS_BIT(7),
      UCP_REQUEST_FLAG_SYNC                  = UCS_BIT(8),
      UCP_REQUEST_FLAG_OFFLOADED             = UCS_BIT(10),
      UCP_REQUEST_FLAG_BLOCK_OFFLOAD         = UCS_BIT(11),
      UCP_REQUEST_FLAG_STREAM_RECV_WAITALL   = UCS_BIT(12),
      UCP_REQUEST_FLAG_SEND_AM               = UCS_BIT(13),
      UCP_REQUEST_FLAG_SEND_TAG              = UCS_BIT(14),
      UCP_REQUEST_FLAG_RNDV_FRAG             = UCS_BIT(15),
      UCP_REQUEST_FLAG_RECV_AM               = UCS_BIT(16),
      UCP_REQUEST_FLAG_RECV_TAG              = UCS_BIT(17),
      UCP_REQUEST_FLAG_RKEY_INUSE            = UCS_BIT(18),
  #if UCS_ENABLE_ASSERT
      UCP_REQUEST_FLAG_STREAM_RECV           = UCS_BIT(19),
      UCP_REQUEST_DEBUG_FLAG_EXTERNAL        = UCS_BIT(20),
      UCP_REQUEST_FLAG_SUPER_VALID           = UCS_BIT(21)
  #else
      UCP_REQUEST_FLAG_STREAM_RECV           = 0,
      UCP_REQUEST_DEBUG_FLAG_EXTERNAL        = 0,
      UCP_REQUEST_FLAG_SUPER_VALID           = 0    
  #endif
  };
  
  
  /**
   * Protocols enumerator to work with send request state
   */
  enum {
      UCP_REQUEST_SEND_PROTO_BCOPY_AM = 0,
      UCP_REQUEST_SEND_PROTO_ZCOPY_AM,
      UCP_REQUEST_SEND_PROTO_RNDV_GET,
      UCP_REQUEST_SEND_PROTO_RNDV_PUT,
      UCP_REQUEST_SEND_PROTO_RMA
  };
  
  
  /**
   * Receive descriptor flags.
   */
  enum {
      UCP_RECV_DESC_FLAG_UCT_DESC         = UCS_BIT(0), /* Descriptor allocated by UCT */
      UCP_RECV_DESC_FLAG_EAGER            = UCS_BIT(1), /* Eager tag message */
      UCP_RECV_DESC_FLAG_EAGER_ONLY       = UCS_BIT(2), /* Eager tag message with single fragment */
      UCP_RECV_DESC_FLAG_EAGER_SYNC       = UCS_BIT(3), /* Eager tag message which requires reply */
      UCP_RECV_DESC_FLAG_EAGER_OFFLOAD    = UCS_BIT(4), /* Eager tag from offload */
      UCP_RECV_DESC_FLAG_EAGER_LAST       = UCS_BIT(5), /* Last fragment of eager tag message.
                                                           Used by tag offload protocol. */
      UCP_RECV_DESC_FLAG_RNDV             = UCS_BIT(6), /* Rendezvous request */
      UCP_RECV_DESC_FLAG_RECV_STARTED     = UCS_BIT(7), /* Receive operation on this descriptor
                                                           was initiated by ucp_am_recv_data_nbx */
      UCP_RECV_DESC_FLAG_MALLOC           = UCS_BIT(8), /* Descriptor was allocated with malloc
                                                           and must be freed, not returned to the
                                                           memory pool or UCT */
      UCP_RECV_DESC_FLAG_AM_CB_INPROGRESS = UCS_BIT(9), /* Descriptor should not be released,
                                                           because UCT AM callback is still in
                                                           the call stack and descriptor is not
                                                           initialized yet. */
      UCP_RECV_DESC_FLAG_RELEASED         = UCS_BIT(10) /* Indicates that the descriptor was
                                                           released and cannot be used. */
  };
  
  
  /**
   * Receive descriptor list pointers
   */
  enum {
      UCP_RDESC_HASH_LIST = 0,
      UCP_RDESC_ALL_LIST  = 1
  };
  
  
  /**
   * Request in progress.
   */
  struct ucp_request {
      /* Operation status */
      ucs_status_t      status;
      /* Request flags */
      uint32_t          flags;
      /* Local request ID taken from PTR MAP */
      ucs_ptr_map_key_t id;
  
      union {
          void                      *user_data; /* Completion user data */
          ucp_request_t             *super_req; /* Super request that is used
                                                   by protocols */
      };
  
      union {
  
          /* "send" part - used for tag_send, am_send, stream_send, put, get, and atomic
           * operations */
          struct {
              ucp_ep_h                ep;
              union {
                  void                   *buffer; /* Send buffer */
                  ucp_request_callback_t flushed_cb; /* Called when flushed */
              };
              ucp_datatype_t          datatype; /* Send type */
              size_t                  length; /* Total length, in bytes */
              ucp_send_nbx_callback_t cb; /* Completion callback */
              ucs_hlist_link_t        list; /* Element in the per-EP list of UCP
                                               flush/proto requests */
  
              const ucp_proto_config_t *proto_config; /* Selected protocol for the request */
  
              /* This structure holds all mutable fields, and everything else
               * except common send/recv fields 'status' and 'flags' is immutable
               * TODO: rework RMA case where length is used instead of dt.offset */
              struct {
                  union {
                      ucp_datatype_iter_t  dt_iter;  /* Send buffer state */
                      ucp_dt_state_t       dt;       /* Position in the send buffer */
                  };
                  uct_completion_t         uct_comp; /* UCT completion used by flush */
              } state;
  
              union {
                  ucp_wireup_msg_t  wireup;
  
                  struct {
                      /* Used to identify matching parts of a large message */
                      uint64_t message_id;
                      union {
                          ucp_tag_t tag;
  
                          struct {
                              union {
                                  /* Can be union, because once header is packed to
                                   * reg_desc, it is not accessed anymore. */
                                  void           *header;
                                  ucp_mem_desc_t *reg_desc; /* pointer to pre-registered buffer,
                                                               used for sending header with
                                                               zcopy protocol */
                              };
                              uint32_t       header_length;
                              uint16_t       am_id;
                              uint16_t       flags;
                          } am;
                      };
                  } msg_proto;
  
                  struct {
                      uint64_t   remote_addr; /* Remote address */
                      ucp_rkey_h rkey; /* Remote memory key */
                  } rma;
  
                  struct {
                      /* Remote request ID received from a peer */
                      ucs_ptr_map_key_t      remote_req_id;
                      uint8_t                am_id;
                      ucs_status_t           status;
                      ucp_tag_t              sender_tag; /* Sender tag, which is
                                                            sent back in sync ack */
                      ucp_request_callback_t comp_cb;    /* Called to complete the
                                                            request */
                  } proto;
  
                  struct {
                      uct_pending_req_t    *req;
                      ucp_wireup_ep_t      *wireup_ep;
                  } proxy;
  
                  struct {
                      uint64_t          remote_address;  /* address of the sender/receiver's data
                                                            buffer for the GET/PUT operation */
                      /* Remote request ID received from a peer */
                      ucs_ptr_map_key_t remote_req_id;
                      ucp_rkey_h        rkey;            /* key for remote send/receive buffer for
                                                            the GET/PUT operation */
                      union {
                          struct {
                              ucp_lane_map_t lanes_map_all; /* actual lanes map */
                              uint8_t        lanes_count; /* actual lanes count */
                              uint8_t        rkey_index[UCP_MAX_LANES];
                          };
                          struct {
                              ucs_ptr_map_key_t rreq_id; /* id of receive request */
                          } rtr;
                      };
                  } rndv;
  
                  struct {
                      /* Remote request ID received from a peer */
                      ucs_ptr_map_key_t remote_req_id;
                  } rndv_data;
  
                  struct {
                      /* Element in queue for segmented RKEY ptr */
                      ucs_queue_elem_t  queue_elem;
                      /* Remote request ID received from a peer */
                      ucs_ptr_map_key_t remote_req_id;
                      /* Key for remote send buffer */
                      ucp_rkey_h        rkey;
                  } rkey_ptr;
  
                  struct {
                      /* The length of the data that should be fetched from sender
                       * side */
                      size_t            length;
                      /* Offset in the receiver's buffer */
                      size_t            offset;
                  } rndv_rtr;
  
                  struct {
                      unsigned           uct_flags; /* Flags to pass to @ref uct_ep_flush */
                      uct_worker_cb_id_t prog_id; /* Progress callback ID */
                      uint32_t           cmpl_sn; /* Sequence number of the remote completion
                                                     this request is waiting for */
                      uint8_t            sw_started;
                      uint8_t            sw_done;
                      uint8_t            num_lanes; /* How many lanes are being flushed */
                      ucp_lane_map_t     started_lanes; /* Which lanes need were flushed */
                  } flush;
  
                  struct {
                      /* UCT EP that should be flushed and destroyed */
                      uct_ep_h           uct_ep;
                      /* Flags that should be passed into @ref uct_ep_flush */
                      unsigned           ep_flush_flags;
                      /* Progress ID, if it's UCS_CALLBACKQ_ID_NULL, no operations
                       * are in-progress */
                      uct_worker_cb_id_t cb_id;
                  } discard_uct_ep;
  
                  struct {
                      uint64_t              remote_addr; /* Remote address */
                      ucp_rkey_h            rkey;        /* Remote memory key */
                      uint64_t              value;       /* Atomic argument */
                      uct_atomic_op_t       uct_op;      /* Requested UCT AMO */
                  } amo;
  
                  struct {
                      ucs_queue_elem_t  queue;     /* Elem in outgoing ssend reqs queue */
                      ucp_tag_t         ssend_tag; /* Tag in offload sync send */
                      void              *rndv_op;  /* Handler of issued rndv send. Need to cancel
                                                      the operation if it is completed by SW. */
                  } tag_offload;
  
                  struct {
                      /* Remote request ID received from a peer */
                      ucs_ptr_map_key_t remote_req_id;
                  } get_reply;
  
                  struct {
                      /* Remote request ID received from a peer */
                      ucs_ptr_map_key_t  remote_req_id;
                      /* Atomic reply data */
                      ucp_atomic_reply_t data;
                  } atomic_reply;
              };
  
              union {
                  ucp_lane_index_t  am_bw_index;     /* AM BW lane index */
                  ucp_lane_map_t    lanes_map_avail; /* Used lanes map */
              };
              uint8_t               mem_type;        /* Memory type, values are
                                                      * ucs_memory_type_t */
              ucp_lane_index_t      pending_lane;    /* Lane on which request was moved
                                                      * to pending state */
              ucp_lane_index_t      lane;            /* Lane on which this request is being sent */
              ucp_lane_index_t      multi_lane_idx;  /* Index of the lane with multi-send */
              uct_pending_req_t     uct;             /* UCT pending request */
              ucp_mem_desc_t        *mdesc;
          } send;
  
          /* "receive" part - used for tag_recv, am_recv and stream_recv operations */
          struct {
              ucs_queue_elem_t      queue;    /* Expected queue element */
              void                  *buffer;  /* Buffer to receive data to */
              ucp_datatype_t        datatype; /* Receive type */
              size_t                length;   /* Total length, in bytes */
              ucs_memory_type_t     mem_type; /* Memory type */
              ucp_dt_state_t        state;
              ucp_worker_t          *worker;
              uct_tag_context_t     uct_ctx;  /* Transport offload context */
              ssize_t               remaining;  /* How much more data
                                                 * to be received */
  
              /* Remote request ID received from a peer */
              ucs_ptr_map_key_t     remote_req_id;
  
              union {
                  struct {
                      ucp_tag_t                   tag;        /* Expected tag */
                      ucp_tag_t                   tag_mask;   /* Expected tag mask */
                      uint64_t                    sn;         /* Tag match sequence */
                      ucp_tag_recv_nbx_callback_t cb;         /* Completion callback */
                      ucp_tag_recv_info_t         info;       /* Completion info to fill */
  
                      /* Can use union, because rdesc is used in expected flow,
                       * while non_contig_buf is used in unexpected flow only. */
                      union {
                          ucp_mem_desc_t      *rdesc;   /* Offload bounce buffer */
                          void                *non_contig_buf; /* Used for assembling
                                                                  multi-fragment
                                                                  non-contig unexpected
                                                                  message in tag offload flow. */
                      };
                      ucp_worker_iface_t      *wiface;    /* Cached iface this request
                                                             is received on. Used in
                                                             tag offload expected callbacks*/
                  } tag;
  
                  struct {
                      size_t                  offset;   /* offset in recv buffer */
                  } frag;
  
                  struct {
                      ucp_stream_recv_nbx_callback_t cb;     /* Completion callback */
                      size_t                         offset; /* Receive data offset */
                      size_t                         length; /* Completion info to fill */
                  } stream;
  
                   struct {
                      ucp_am_recv_data_nbx_callback_t cb;    /* Completion callback */
                      ucp_recv_desc_t                 *desc; /* RTS desc */
                  } am;
              };
          } recv;
  
          struct {
              ucp_worker_h            worker;     /* Worker to flush */
              ucp_send_nbx_callback_t cb;         /* Completion callback */
              uct_worker_cb_id_t      prog_id;    /* Progress callback ID */
              int                     comp_count; /* Countdown to request completion */
              ucp_ep_ext_gen_t        *next_ep;   /* Next endpoint to flush */
          } flush_worker;
      };
  };
  
  
  /**
   * Unexpected receive descriptor. If it is initialized in the headroom of UCT
   * descriptor, the layout looks like the following:
   *
   *
   * headroom                                    data
   * |-------------------------------------------|-------------------------|
   * | unused | ucp_recv_desc |      priv_length |                         |
   * |        |               |                  |                         |
   * |-------------------------------------------|-------------------------|
   *
   * Some protocols (i. e. tag offload) may need some space right before the
   * incoming data to add specific headers needed for further message processing.
   * Note: priv_length value should be in [0, UCP_WORKER_HEADROOM_PRIV_SIZE] range.
   */
  struct ucp_recv_desc {
      union {
          ucs_list_link_t     tag_list[2];     /* Hash list TAG-element */
          ucs_queue_elem_t    stream_queue;    /* Queue STREAM-element */
          ucs_queue_elem_t    tag_frag_queue;  /* Tag fragments queue */
          ucp_am_first_desc_t am_first;        /* AM first fragment data needed
                                                  for assembling the message */
          ucs_queue_elem_t    am_mid_queue;    /* AM middle fragments queue */
      };
      uint32_t                length;          /* Received length */
      union {
          uint32_t            payload_offset;  /* Offset from end of the descriptor
                                                * to AM data */
          uint32_t            am_malloc_offset; /* Offset from rdesc, holding
                                                   assembled multi-fragment active
                                                   message, to the originally
                                                   malloc'd buffer pointer */
      };
      uint16_t                flags;           /* Flags */
      int16_t                 uct_desc_offset; /* Offset which needs to be
                                                  substructed from rdesc when
                                                  releasing it back to UCT */
  };
  
  
  /**
   * Defines protocol functions for ucp_request_send_start() function.
   * TODO will be removed when switching to new protocols implementation.
   */
  struct ucp_request_send_proto {
      uct_pending_callback_t     contig_short;     /**< Progress short data */
      uct_pending_callback_t     bcopy_single;     /**< Progress bcopy single fragment */
      uct_pending_callback_t     bcopy_multi;      /**< Progress bcopy multi-fragment */
      uct_pending_callback_t     zcopy_single;     /**< Progress zcopy single fragment */
      uct_pending_callback_t     zcopy_multi;      /**< Progress zcopy multi-fragment */
      uct_completion_callback_t  zcopy_completion; /**< Callback for UCT zcopy completion */
      size_t                     only_hdr_size;    /**< Header size for single / short */
  };
  
  
  extern ucs_mpool_ops_t ucp_request_mpool_ops;
  extern ucs_mpool_ops_t ucp_rndv_get_mpool_ops;
  extern const ucp_request_param_t ucp_request_null_param;
  
  
  int ucp_request_pending_add(ucp_request_t *req, unsigned pending_flags);
  
  ucs_status_t ucp_request_memory_reg(ucp_context_t *context, ucp_md_map_t md_map,
                                      void *buffer, size_t length, ucp_datatype_t datatype,
                                      ucp_dt_state_t *state, ucs_memory_type_t mem_type,
                                      ucp_request_t *req_dbg, unsigned uct_flags);
  
  void ucp_request_memory_dereg(ucp_context_t *context, ucp_datatype_t datatype,
                                ucp_dt_state_t *state, ucp_request_t *req_dbg);
  
  void ucp_request_dt_invalidate(ucp_request_t *req, ucs_status_t status);
  
  ucs_status_t ucp_request_send_start(ucp_request_t *req, ssize_t max_short,
                                      size_t zcopy_thresh, size_t zcopy_max,
                                      size_t dt_count, size_t priv_iov_count,
                                      size_t length,
                                      const ucp_ep_msg_config_t* msg_config,
                                      const ucp_request_send_proto_t *proto);
  
  /* Fast-forward to data end */
  void ucp_request_send_state_ff(ucp_request_t *req, ucs_status_t status);
  
  ucs_status_t ucp_request_recv_msg_truncated(ucp_request_t *req, size_t length,
                                              size_t offset);
  
  #endif
  