UCP/WIREUP: Implement multi-lane/shm support for CM

src/ucp/core/ucp_ep.c

@@ -786,11 +786,10 @@ void ucp_ep_disconnected(ucp_ep_h ep, int force)
 
     ep->flags &= ~UCP_EP_FLAG_USED;
 
-    if ((ep->flags & (UCP_EP_FLAG_CONNECT_REQ_QUEUED |
+    /* in case of CM connection ep has to be disconnected */
+    if (!ucp_ep_has_cm_lane(ep) &&
+        (ep->flags & (UCP_EP_FLAG_CONNECT_REQ_QUEUED |
                       UCP_EP_FLAG_REMOTE_CONNECTED)) && !force) {
-        /* in case of CM connection ep has to be disconnected */
-        ucs_assert(!ucp_ep_has_cm_lane(ep));
-
         /* Endpoints which have remote connection are destroyed only when the
          * worker is destroyed, to enable remote endpoints keep sending
          * TODO negotiate disconnect.
@@ -2106,7 +2105,8 @@ ucp_wireup_ep_t * ucp_ep_get_cm_wireup_ep(ucp_ep_h ep)
         return NULL;
     }
 
-    return ucp_wireup_ep_test(ep->uct_eps[lane]) ?
+    return ((ep->uct_eps[lane] != NULL) &&
+            ucp_wireup_ep_test(ep->uct_eps[lane])) ?
            ucs_derived_of(ep->uct_eps[lane], ucp_wireup_ep_t) : NULL;
 }
 

src/ucp/core/ucp_ep.h

@@ -65,7 +65,9 @@ enum {
     UCP_EP_FLAG_SOCKADDR_PARTIAL_ADDR  = UCS_BIT(21),/* DEBUG: Partial worker address was sent
                                                                to the remote peer when starting
                                                                connection establishment on this EP */
-    UCP_EP_FLAG_FLUSH_STATE_VALID      = UCS_BIT(22) /* DEBUG: flush_state is valid */
+    UCP_EP_FLAG_FLUSH_STATE_VALID      = UCS_BIT(22),/* DEBUG: flush_state is valid */
+    UCP_EP_FLAG_DISCONNECTED_CM_LANE   = UCS_BIT(23) /* DEBUG: CM lane was disconnected, i.e.
+                                                               @uct_ep_disconnect was called for CM EP */
 };
 
 

src/ucp/core/ucp_ep.inl

@@ -147,11 +147,29 @@ static UCS_F_ALWAYS_INLINE ucp_ep_h ucp_ep_from_ext_proto(ucp_ep_ext_proto_t *ep
     return (ucp_ep_h)ucs_strided_elem_get(ep_ext, 2, 0);
 }
 
+static inline int
+ucp_ep_config_key_has_cm_lane(const ucp_ep_config_key_t *config_key)
+{
+    return config_key->cm_lane != UCP_NULL_LANE;
+}
+
+static inline int ucp_ep_has_cm_lane(ucp_ep_h ep)
+{
+    return (ep->cfg_index != UCP_WORKER_CFG_INDEX_NULL) &&
+           ucp_ep_config_key_has_cm_lane(&ucp_ep_config(ep)->key);
+}
+
+static UCS_F_ALWAYS_INLINE ucp_lane_index_t ucp_ep_get_cm_lane(ucp_ep_h ep)
+{
+    return ucp_ep_config(ep)->key.cm_lane;
+}
+
 static UCS_F_ALWAYS_INLINE ucp_ep_flush_state_t* ucp_ep_flush_state(ucp_ep_h ep)
 {
     ucs_assert(ep->flags & UCP_EP_FLAG_FLUSH_STATE_VALID);
     ucs_assert(!(ep->flags & UCP_EP_FLAG_ON_MATCH_CTX));
-    ucs_assert(!(ep->flags & UCP_EP_FLAG_LISTENER));
+    ucs_assert(!(ep->flags & UCP_EP_FLAG_LISTENER) ||
+               ucp_ep_has_cm_lane(ep));
     ucs_assert(!(ep->flags & UCP_EP_FLAG_CLOSE_REQ_VALID));
     return &ucp_ep_ext_gen(ep)->flush_state;
 }
@@ -236,23 +254,6 @@ ucp_ep_config_get_dst_md_cmpt(const ucp_ep_config_key_t *key,
     return key->dst_md_cmpts[idx];
 }
 
-static inline int
-ucp_ep_config_key_has_cm_lane(const ucp_ep_config_key_t *config_key)
-{
-    return config_key->cm_lane != UCP_NULL_LANE;
-}
-
-static inline int ucp_ep_has_cm_lane(ucp_ep_h ep)
-{
-    return (ep->cfg_index != UCP_WORKER_CFG_INDEX_NULL) &&
-           ucp_ep_config_key_has_cm_lane(&ucp_ep_config(ep)->key);
-}
-
-static UCS_F_ALWAYS_INLINE ucp_lane_index_t ucp_ep_get_cm_lane(ucp_ep_h ep)
-{
-    return ucp_ep_config(ep)->key.cm_lane;
-}
-
 static inline int
 ucp_ep_config_connect_p2p(ucp_worker_h worker,
                           const ucp_ep_config_key_t *ep_config_key,

src/ucp/core/ucp_proxy_ep.c

@@ -229,6 +229,10 @@ void ucp_proxy_ep_replace(ucp_proxy_ep_t *proxy_ep)
      * is pointed to by another proxy ep. if so, redirect that other proxy ep
      * to point to the underlying uct ep. */
     for (lane = 0; lane < ucp_ep_num_lanes(ucp_ep); ++lane) {
+        if (ucp_ep->uct_eps[lane] == NULL) {
+            continue;
+        }
+
         ucp_proxy_ep_replace_if_owned(ucp_ep->uct_eps[lane], &proxy_ep->super,
                                       tl_ep);
     }

src/ucp/core/ucp_request.h

@@ -208,6 +208,15 @@ struct ucp_request {
                     uct_worker_cb_id_t        prog_id;/* Slow-path callback */
                 } disconnect;
 
+                struct {
+                    uct_worker_h          uct_worker;     /* UCT worker where discard UCT EP operation
+                                                           * submitted on */
+                    uct_ep_h              uct_ep;         /* UCT EP that should be flushed and
+                                                             destroyed */
+                    unsigned              ep_flush_flags; /* Flags that should be passed into
+                                                             @ref uct_ep_flush */
+                } discard_uct_ep;
+
                 struct {
                     uint64_t              remote_addr; /* Remote address */
                     ucp_rkey_h            rkey;        /* Remote memory key */

src/ucp/core/ucp_worker.c

@@ -445,7 +445,8 @@ static unsigned ucp_worker_iface_err_handle_progress(void *arg)
         if (lane != failed_lane) {
             ucs_trace("ep %p: destroy uct_ep[%d]=%p", ucp_ep, lane,
                       ucp_ep->uct_eps[lane]);
-            uct_ep_destroy(ucp_ep->uct_eps[lane]);
+            ucp_worker_discard_uct_ep(ucp_ep->worker, ucp_ep->uct_eps[lane],
+                                      UCT_FLUSH_FLAG_CANCEL);
             ucp_ep->uct_eps[lane] = NULL;
         }
     }
@@ -2306,7 +2307,6 @@ void ucp_worker_release_address(ucp_worker_h worker, ucp_address_t *address)
     ucs_free(address);
 }
 
-
 void ucp_worker_print_info(ucp_worker_h worker, FILE *stream)
 {
     ucp_context_h context = worker->context;
@@ -2350,3 +2350,104 @@ void ucp_worker_print_info(ucp_worker_h worker, FILE *stream)
 
     UCP_WORKER_THREAD_CS_EXIT_CONDITIONAL(worker);
 }
+
+static unsigned ucp_worker_discard_uct_ep_destroy_progress(void *arg)
+{
+    uct_ep_h uct_ep = (uct_ep_h)arg;
+
+    ucs_debug("destroy uct_ep=%p", uct_ep);
+    uct_ep_destroy(uct_ep);
+
+    return 1;
+}
+
+static void
+ucp_worker_discard_uct_ep_flush_comp(uct_completion_t *self,
+                                     ucs_status_t status)
+{
+    uct_worker_cb_id_t cb_id    = UCS_CALLBACKQ_ID_NULL;
+    ucp_request_t *req          = ucs_container_of(self, ucp_request_t,
+                                                   send.state.uct_comp);
+    uct_ep_h uct_ep             = req->send.discard_uct_ep.uct_ep;
+
+    ucs_trace_req("flush completion req=%p status=%d", req, status);
+
+    uct_worker_progress_register_safe(req->send.discard_uct_ep.uct_worker,
+                                      ucp_worker_discard_uct_ep_destroy_progress,
+                                      uct_ep, UCS_CALLBACKQ_FLAG_ONESHOT, &cb_id);
+
+    ucp_request_put(req);
+}
+
+static unsigned ucp_worker_discard_uct_ep_progress(void *arg)
+{
+    ucp_request_t *req = (ucp_request_t*)arg;
+    uct_ep_h uct_ep    = req->send.discard_uct_ep.uct_ep;
+    ucs_status_t status;
+
+    req->send.state.uct_comp.func  = ucp_worker_discard_uct_ep_flush_comp;
+    req->send.state.uct_comp.count = 1;
+
+    for (;;) {
+        status = uct_ep_flush(uct_ep, req->send.discard_uct_ep.ep_flush_flags,
+                              &req->send.state.uct_comp);
+        if (status == UCS_ERR_NO_RESOURCE) {
+            status = uct_ep_pending_add(uct_ep, &req->send.uct, 0);
+            if (status == UCS_OK) {
+                break;
+            }
+        } else {
+            if (status != UCS_INPROGRESS) {
+                ucp_worker_discard_uct_ep_flush_comp(&req->send.state.uct_comp,
+                                                     status);
+            }
+            break;
+        }
+    }
+
+    return 1;
+}
+
+void ucp_worker_discard_uct_ep(ucp_worker_h worker, uct_ep_h uct_ep,
+                               unsigned ep_flush_flags)
+{
+    uct_worker_cb_id_t cb_id = UCS_CALLBACKQ_ID_NULL;
+    ucp_wireup_ep_t *wireup_ep;
+    ucp_request_t *req;
+    int is_owner;
+
+    ucs_assert(uct_ep != NULL);
+
+    if (ucp_wireup_ep_test(uct_ep)) {
+        wireup_ep = ucp_wireup_ep(uct_ep);
+        ucs_assert(wireup_ep != NULL);
+
+        is_owner = wireup_ep->super.is_owner;
+        uct_ep   = ucp_wireup_ep_extract_next_ep(uct_ep);
+
+        /* destroy WIREUP EP allocated for this UCT EP, since
+         * discard operation most likely won't have an access to
+         * UCP EP as it could be destroyed by the caller */
+        uct_ep_destroy(&wireup_ep->super.super);
+
+        if (!is_owner) {
+            /* do nothing, if this wireup EP is not an owner for UCT EP */
+            return;
+        }
+    }
+
+    req = ucp_request_get(worker);
+    if (ucs_unlikely(req == NULL)) {
+        ucs_error("unable to allocate reqquest");
+        return;
+    }
+
+    ucs_assert(!ucp_wireup_ep_test(uct_ep));
+    req->send.discard_uct_ep.uct_worker     = worker->uct;
+    req->send.discard_uct_ep.uct_ep         = uct_ep;
+    req->send.discard_uct_ep.ep_flush_flags = ep_flush_flags;
+    uct_worker_progress_register_safe(worker->uct,
+                                      ucp_worker_discard_uct_ep_progress,
+                                      req, UCS_CALLBACKQ_FLAG_ONESHOT,
+                                      &cb_id);
+}

src/ucp/core/ucp_worker.h

@@ -308,4 +308,7 @@ ucs_status_t ucp_worker_set_ep_failed(ucp_worker_h worker, ucp_ep_h ucp_ep,
                                       uct_ep_h uct_ep, ucp_lane_index_t lane,
                                       ucs_status_t status);
 
+void ucp_worker_discard_uct_ep(ucp_worker_h worker, uct_ep_h uct_ep,
+                               unsigned ep_flush_flags);
+
 #endif

src/ucp/wireup/select.c

@@ -833,7 +833,7 @@ ucp_wireup_add_cm_lane(const ucp_wireup_select_params_t *select_params,
     select_info.path_index = 0;  /**< Only one lane per CM device */
 
     /* server is not a proxy because it can create all lanes connected */
-    return ucp_wireup_add_lane_desc(&select_info, select_info.rsc_index,
+    return ucp_wireup_add_lane_desc(&select_info, UCP_MAX_MDS,
                                     UCP_LANE_TYPE_CM, 0, select_ctx);
 }