UCP/EP: Cleanup proto reqs on failure

src/ucp/core/ucp_ep.c

@@ -173,6 +173,7 @@ ucs_status_t ucp_ep_create_base(ucp_worker_h worker, const char *peer_name,
     memset(&ucp_ep_ext_gen(ep)->ep_match, 0,
            sizeof(ucp_ep_ext_gen(ep)->ep_match));
 
+    ucs_hlist_head_init(&ucp_ep_ext_gen(ep)->proto_reqs);
     ucp_stream_ep_init(ep);
     ucp_am_ep_init(ep);
 
@@ -605,6 +606,8 @@ static ucs_status_t ucp_ep_create_to_sock_addr(ucp_worker_h worker,
         goto err_cleanup_lanes;
     }
 
+    ucp_ep_flush_state_reset(ep);
+
     *ep_p = ep;
     return UCS_OK;
 
@@ -963,6 +966,7 @@ void ucp_ep_disconnected(ucp_ep_h ep, int force)
 
     ucp_stream_ep_cleanup(ep);
     ucp_am_ep_cleanup(ep);
+    ucp_ep_reqs_purge(ep, UCS_ERR_CANCELED);
 
     ucp_ep_update_flags(ep, 0, UCP_EP_FLAG_USED);
 
@@ -2667,3 +2671,72 @@ void ucp_ep_do_keepalive(ucp_ep_h ep, ucp_lane_map_t *lane_map)
         *lane_map &= ~UCS_BIT(lane);
     }
 }
+
+static void ucp_ep_req_purge(ucp_ep_h ucp_ep, ucp_request_t *req,
+                             ucs_status_t status, int recursive)
+{
+    ucp_trace_req(req, "purged with status %s (%d) on ep %p",
+                  ucs_status_string(status), status, ucp_ep);
+
+    if (req->id != UCP_REQUEST_ID_INVALID) {
+        ucp_request_id_release(req);
+    }
+
+    if (req->flags & (UCP_REQUEST_FLAG_SEND_AM | UCP_REQUEST_FLAG_SEND_TAG)) {
+        ucs_assert(req->super_req == NULL);
+        ucs_assert(req->send.ep == ucp_ep);
+        ucp_request_complete_and_dereg_send(req, status);
+    } else if (req->flags & UCP_REQUEST_FLAG_RECV_AM) {
+        ucs_assert(req->super_req == NULL);
+        ucs_assert(recursive); /* Mustn't be directly contained in an EP list
+                                * of tracking requests */
+        ucp_request_complete_am_recv(req, status);
+    } else if (req->flags & UCP_REQUEST_FLAG_RECV_TAG) {
+        ucs_assert(req->super_req == NULL);
+        ucs_assert(recursive); /* Mustn't be directly contained in an EP list
+                                * of tracking requests */
+        ucp_request_complete_tag_recv(req, status);
+    } else if (req->flags & UCP_REQUEST_FLAG_RNDV_FRAG) {
+        ucs_assert(req->super_req != NULL);
+        ucs_assert(req->send.ep == ucp_ep);
+        ucs_assert(recursive); /* Mustn't be directly contained in an EP list
+                                * of tracking requests */
+
+        /* It means that purging started from a request responsible for sending
+         * RTR, so a request is responsible for copying data from staging buffer
+         * and it uses a receive part of a request */
+        req->super_req->recv.remaining -= req->recv.length;
+        if (req->super_req->recv.remaining == 0) {
+            ucp_ep_req_purge(ucp_ep, req->super_req, status, 1);
+        }
+
+        ucp_request_put(req);
+    } else {
+        ucs_assert(req->super_req != NULL);
+        ucs_assert(req->send.ep == ucp_ep);
+
+        ucp_ep_req_purge(ucp_ep, req->super_req, status, 1);
+        ucp_request_put(req);
+    }
+}
+
+void ucp_ep_reqs_purge(ucp_ep_h ucp_ep, ucs_status_t status)
+{
+    ucs_hlist_head_t *proto_reqs = &ucp_ep_ext_gen(ucp_ep)->proto_reqs;
+    ucp_request_t *req;
+
+    while (!ucs_hlist_is_empty(proto_reqs)) {
+        req = ucs_hlist_head_elem(proto_reqs, ucp_request_t, send.list);
+        ucp_ep_req_purge(ucp_ep, req, status, 0);
+    }
+
+    if (/* Flush state is already valid (i.e. EP doesn't exist on matching
+         * context) and not invalidated yet, also remote EP ID is already set */
+        !(ucp_ep->flags &
+          (UCP_EP_FLAG_ON_MATCH_CTX | UCP_EP_FLAG_CLOSE_REQ_VALID))) {
+        ucs_hlist_for_each_extract(req, &ucp_ep_flush_state(ucp_ep)->reqs,
+                                   send.list) {
+            ucp_ep_flush_request_ff(req, status);
+        }
+    }
+}

src/ucp/core/ucp_ep.h

@@ -446,6 +446,8 @@ typedef struct {
         ucp_ep_flush_state_t      flush_state;   /* Remote completion status */
     };
     ucp_ep_ext_control_t          *control_ext;  /* Control data path extension */
+    /* List of requests which are waiting for remote completion */
+    ucs_hlist_head_t              proto_reqs;
 } ucp_ep_ext_gen_t;
 
 
@@ -660,4 +662,13 @@ ucs_status_t ucp_ep_do_uct_ep_keepalive(ucp_ep_h ucp_ep, uct_ep_h uct_ep,
  */
 void ucp_ep_do_keepalive(ucp_ep_h ep, ucp_lane_map_t *lane_map);
 
+/**
+ * @brief Purge flush and protocol requests scheduled on a given UCP endpoint.
+ *
+ * @param [in]     ucp_ep           Endpoint object on which requests should be
+ *                                  purged.
+ * @param [in]     status           Completion status.
+ */
+void ucp_ep_reqs_purge(ucp_ep_h ucp_ep, ucs_status_t status);
+
 #endif

src/ucp/core/ucp_request.h

@@ -135,9 +135,11 @@ struct ucp_request {
                 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 */
+            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 */
 
@@ -178,9 +180,8 @@ struct ucp_request {
                 } msg_proto;
 
                 struct {
-                    ucs_ptr_map_key_t sreq_id;     /* Send request ID */
-                    uint64_t          remote_addr; /* Remote address */
-                    ucp_rkey_h        rkey;        /* Remote memory key */
+                    uint64_t   remote_addr; /* Remote address */
+                    ucp_rkey_h rkey; /* Remote memory key */
                 } rma;
 
                 struct {
@@ -241,16 +242,14 @@ struct ucp_request {
                 } rndv_rtr;
 
                 struct {
-                    ucs_hlist_link_t       list_elem; /* Element in the per-EP list of UCP
-                                                         flush requests */
-                    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 */
+                    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 {

src/ucp/core/ucp_request.inl

@@ -256,6 +256,31 @@ ucp_request_can_complete_stream_recv(ucp_request_t *req)
     return 1;
 }
 
+
+static UCS_F_ALWAYS_INLINE ucp_request_t*
+ucp_request_mem_alloc(const char *name)
+{
+    ucp_request_t *req = (ucp_request_t*)ucs_malloc(sizeof(*req), name);
+
+    if (ucs_unlikely(req == NULL)) {
+        return NULL;
+    }
+
+    ucs_trace_req("allocated request %p (%s)", req, name);
+    ucp_request_id_reset(req);
+    UCS_PROFILE_REQUEST_NEW(req, "ucp_request", 0);
+
+    return req;
+}
+
+static UCS_F_ALWAYS_INLINE void ucp_request_mem_free(ucp_request_t *req)
+{
+    UCS_PROFILE_REQUEST_FREE(req);
+    ucp_request_id_check(req, ==, UCP_REQUEST_ID_INVALID);
+    ucs_trace_req("freed request %p", req);
+    ucs_free(req);
+}
+
 /*
  * @return Whether completed.
  *         *req_status if filled with the completion status if completed.
@@ -816,6 +841,13 @@ static UCS_F_ALWAYS_INLINE void ucp_request_id_alloc(ucp_request_t *req)
 {
     ucp_request_id_check(req, ==, UCP_REQUEST_ID_INVALID);
     ucp_ep_ptr_id_alloc(req->send.ep, req, &req->id);
+
+    /* TODO: combine checks for err_mode and PTR MAP indirect flag */
+    if (ucs_unlikely(ucp_ep_config(req->send.ep)->key.err_mode ==
+                     UCP_ERR_HANDLING_MODE_PEER)) {
+        ucs_hlist_add_tail(&ucp_ep_ext_gen(req->send.ep)->proto_reqs,
+                           &req->send.list);
+    }
 }
 
 static UCS_F_ALWAYS_INLINE ucs_ptr_map_key_t
@@ -825,6 +857,18 @@ ucp_request_get_id(const ucp_request_t *req)
     return req->id;
 }
 
+static UCS_F_ALWAYS_INLINE void ucp_request_ep_list_del(ucp_request_t *req)
+{
+    ucp_ep_h ep = req->send.ep;
+
+    ucs_assert(ep != NULL);
+    /* TODO: combine checks for err_mode and PTR MAP indirect flag */
+    if (ucs_unlikely(ucp_ep_config(ep)->key.err_mode ==
+                     UCP_ERR_HANDLING_MODE_PEER)) {
+        ucs_hlist_del(&ucp_ep_ext_gen(ep)->proto_reqs, &req->send.list);
+    }
+}
+
 static UCS_F_ALWAYS_INLINE void
 ucp_request_id_release(ucp_request_t *req)
 {
@@ -835,6 +879,7 @@ ucp_request_id_release(ucp_request_t *req)
     status = ucs_ptr_map_del(&worker->ptr_map, req->id);
     ucs_assertv(status == UCS_OK, "req %p: failed to release id", req);
     ucp_request_id_reset(req);
+    ucp_request_ep_list_del(req);
 }
 
 static UCS_F_ALWAYS_INLINE ucs_status_t
@@ -844,19 +889,21 @@ ucp_request_get_by_id(ucp_worker_h worker, ucs_ptr_map_key_t id,
     ucs_status_t status;
     void *ptr;
 
+    ucs_assert(id != UCP_REQUEST_ID_INVALID);
+
     status = ucs_ptr_map_get(&worker->ptr_map, id, extract, &ptr);
     if (ucs_unlikely(status != UCS_OK)) {
         return status;
     }
 
     *req_p = (ucp_request_t*)ptr;
     ucp_request_id_check(*req_p, ==, id);
-
     if (extract) {
         /* If request ID was released, then need to reset the request ID to use
          * the value for checking whether the request ID should be put to PTR
          * map or not in case of error handling */
         ucp_request_id_reset(*req_p);
+        ucp_request_ep_list_del(*req_p);
     }
 
     return status;
@@ -895,6 +942,16 @@ ucp_request_invoke_uct_completion_success(ucp_request_t *req)
     return UCS_OK;
 }
 
+/* The function can be used to complete any UCP send request */
+static UCS_F_ALWAYS_INLINE void
+ucp_request_complete_and_dereg_send(ucp_request_t *sreq, ucs_status_t status)
+{
+    ucs_assert(!sreq->send.ep->worker->context->config.ext.proto_enable);
+    ucp_request_send_generic_dt_finish(sreq);
+    ucp_request_send_buffer_dereg(sreq);
+    ucp_request_complete_send(sreq, status);
+}
+
 
 #define UCP_REQUEST_GET_BY_ID(_req_p, _worker, _req_id, _extract, \
                               _action, _fmt_str, ...) \

src/ucp/core/ucp_worker.c

@@ -438,7 +438,9 @@ static unsigned ucp_worker_iface_err_handle_progress(void *arg)
     ucs_assert(ucp_ep->flags & UCP_EP_FLAG_FAILED);
 
     ucp_ep_discard_lanes(ucp_ep, status);
+    ucp_ep_reqs_purge(ucp_ep, status);
     ucp_stream_ep_cleanup(ucp_ep);
+
     if (ucp_ep->flags & UCP_EP_FLAG_USED) {
         if (ucp_ep->flags & UCP_EP_FLAG_CLOSE_REQ_VALID) {
             ucs_assert(ucp_ep->flags & UCP_EP_FLAG_CLOSED);

src/ucp/rma/flush.c

@@ -162,9 +162,10 @@ static void ucp_ep_flush_progress(ucp_request_t *req)
                 ucs_trace_req("flush request %p remote completions done", req);
             } else {
                 req->send.flush.cmpl_sn = flush_state->send_sn;
-                ucs_hlist_add_tail(&flush_state->reqs, &req->send.flush.list_elem);
-                ucs_trace_req("added flush request %p to ep remote completion queue"
-                              " with sn %d", req, req->send.flush.cmpl_sn);
+                ucs_hlist_add_tail(&flush_state->reqs, &req->send.list);
+                ucs_trace_req("added flush request %p to ep remote completion"
+                              " queue with sn %d",
+                              req, req->send.flush.cmpl_sn);
             }
         }
         req->send.flush.sw_started = 1;

src/ucp/rma/rma.inl

@@ -89,7 +89,7 @@ static inline void ucp_ep_rma_remote_request_completed(ucp_ep_t *ep)
     ucp_worker_flush_ops_count_dec(ep->worker);
     ++flush_state->cmpl_sn;
 
-    ucs_hlist_for_each_extract_if(req, &flush_state->reqs, send.flush.list_elem,
+    ucs_hlist_for_each_extract_if(req, &flush_state->reqs, send.list,
                                   UCS_CIRCULAR_COMPARE32(
                                           req->send.flush.cmpl_sn, <=,
                                           flush_state->cmpl_sn)) {

src/ucp/rndv/rndv.c

@@ -266,14 +266,6 @@ ucp_rndv_adjust_zcopy_length(size_t min_zcopy, size_t max_zcopy, size_t align,
     return result_length;
 }
 
-static void ucp_rndv_complete_send(ucp_request_t *sreq, ucs_status_t status)
-{
-    ucs_assert(!sreq->send.ep->worker->context->config.ext.proto_enable);
-    ucp_request_send_generic_dt_finish(sreq);
-    ucp_request_send_buffer_dereg(sreq);
-    ucp_request_complete_send(sreq, status);
-}
-
 void ucp_rndv_req_send_ack(ucp_request_t *ack_req, ucp_request_t *req,
                            ucs_ptr_map_key_t remote_req_id, ucs_status_t status,
                            ucp_am_id_t am_id, const char *ack_str)
@@ -1478,7 +1470,8 @@ UCS_PROFILE_FUNC(ucs_status_t, ucp_rndv_ats_handler,
     if (sreq->flags & UCP_REQUEST_FLAG_OFFLOADED) {
         ucp_tag_offload_cancel_rndv(sreq);
     }
-    ucp_rndv_complete_send(sreq, rep_hdr->status);
+
+    ucp_request_complete_and_dereg_send(sreq, rep_hdr->status);
     return UCS_OK;
 }
 
@@ -1495,7 +1488,7 @@ ucs_status_t ucp_rndv_rts_handle_status_from_pending(ucp_request_t *sreq,
         }
 
         ucp_request_id_release(sreq);
-        ucp_rndv_complete_send(sreq, status);
+        ucp_request_complete_and_dereg_send(sreq, status);
     }
 
     return UCS_OK;
@@ -1549,7 +1542,7 @@ UCS_PROFILE_FUNC(ucs_status_t, ucp_rndv_progress_am_bcopy, (self),
         return UCS_ERR_NO_RESOURCE;
     }
 
-    ucp_rndv_complete_send(sreq, status);
+    ucp_request_complete_and_dereg_send(sreq, status);
 
     return UCS_OK;
 }

src/ucp/wireup/wireup.c

@@ -172,7 +172,7 @@ ucs_status_t ucp_wireup_msg_progress(uct_pending_req_t *self)
 
 out_free_req:
     ucs_free(req->send.buffer);
-    ucs_free(req);
+    ucp_request_mem_free(req);
 out:
     UCS_ASYNC_UNBLOCK(&ep->worker->async);
     return status;
@@ -224,7 +224,7 @@ ucp_wireup_msg_send(ucp_ep_h ep, uint8_t type, const ucp_tl_bitmap_t *tl_bitmap,
     /* We cannot allocate from memory pool because it's not thread safe
      * and this function may be called from any thread
      */
-    req = ucs_malloc(sizeof(*req), "wireup_msg_req");
+    req = ucp_request_mem_alloc("wireup_msg_req");
     if (req == NULL) {
         ucs_error("failed to allocate request for sending WIREUP message");
         return UCS_ERR_NO_MEMORY;
@@ -240,7 +240,7 @@ ucp_wireup_msg_send(ucp_ep_h ep, uint8_t type, const ucp_tl_bitmap_t *tl_bitmap,
                                     &req->send.wireup, &req->send.buffer,
                                     &req->send.length);
     if (status != UCS_OK) {
-        ucs_free(req);
+        ucp_request_mem_free(req);
         return status;
     }
 
@@ -704,6 +704,7 @@ ucp_wireup_send_ep_removed(ucp_worker_h worker, const ucp_wireup_msg_t *msg,
     }
 
     ucp_ep_update_remote_id(reply_ep, msg->src_ep_id);
+    ucp_ep_flush_state_reset(reply_ep);
     status = ucp_wireup_msg_send(reply_ep, UCP_WIREUP_MSG_EP_REMOVED,
                                  &ucp_tl_bitmap_min, NULL);
     if (status != UCS_OK) {