btl/uct: fix a race condition when setting up endpoints

This commit fixes a regression I introduced when changing how the endpoint flags
are set. There is no guarantee when the remote completion message comes in that
the local endpoint has been connected (only the remote endpoint). Instead of
always setting the endpoint as ready here the code now only sets that the remote
is connected and then sets the endpoint as ready only if both the local and
remote endpoints are connected. When the local endpoint is connected it
similarly checks that the remote is connected before setting the endpoint as
ready.

This was verified using RMA-MT with 128 threads and it does fix the race.

Signed-off-by: Nathan Hjelm <[email protected]>

Author: hjelmn

opal/mca/btl/uct/btl_uct_endpoint.c

@@ -368,6 +368,9 @@ static int mca_btl_uct_endpoint_connect_endpoint(
         if (UCS_OK != ucs_status) {
             return OPAL_ERROR;
         }
+
+        mca_btl_uct_endpoint_set_flag(uct_btl, endpoint, tl_context->context_id, tl_endpoint,
+                                      MCA_BTL_UCT_ENDPOINT_FLAG_EP_CONNECTED);
     }
 
     opal_timer_t now = opal_timer_base_get_usec();

opal/mca/btl/uct/btl_uct_endpoint.h

@@ -14,6 +14,7 @@
  *                         reserved.
  * Copyright (c) 2020      Amazon.com, Inc. or its affiliates.
  *                         All Rights reserved.
+ * Copyright (c) 2025      Google, LLC. All rights reserved.
  * $COPYRIGHT$
  *
  * Additional copyrights may follow
@@ -106,5 +107,25 @@ static inline int mca_btl_uct_endpoint_check_am(mca_btl_uct_module_t *module,
                                       module->am_tl->tl_index);
 }
 
+// Requires that the endpoint lock is held.
+static inline void mca_btl_uct_endpoint_set_flag(mca_btl_uct_module_t *module, mca_btl_uct_endpoint_t *endpoint,
+                                                 int context_id, mca_btl_uct_tl_endpoint_t *tl_endpoint, int32_t flag) {
+    opal_atomic_wmb();
+    int32_t flag_value = opal_atomic_or_fetch_32(&tl_endpoint->flags, flag);
+    if ((flag_value & (MCA_BTL_UCT_ENDPOINT_FLAG_EP_CONNECTED | MCA_BTL_UCT_ENDPOINT_FLAG_CONN_REM_READY)) ==
+        (MCA_BTL_UCT_ENDPOINT_FLAG_EP_CONNECTED | MCA_BTL_UCT_ENDPOINT_FLAG_CONN_REM_READY)) {
+	opal_atomic_fetch_or_32(&tl_endpoint->flags, MCA_BTL_UCT_ENDPOINT_FLAG_CONN_READY);
+
+        opal_atomic_wmb();
+
+        mca_btl_uct_base_frag_t *frag;
+        OPAL_LIST_FOREACH (frag, &module->pending_frags, mca_btl_uct_base_frag_t) {
+            if (frag->context->context_id == context_id && endpoint == frag->endpoint) {
+                frag->ready = true;
+            }
+        }
+    }
+}
+
 END_C_DECLS
 #endif

opal/mca/btl/uct/btl_uct_tl.c

@@ -6,7 +6,7 @@
  *                         and Technology (RIST). All rights reserved.
  * Copyright (c) 2018      Triad National Security, LLC. All rights
  *                         reserved.
- * Copyright (c) 2019      Google, LLC. All rights reserved.
+ * Copyright (c) 2019-2025 Google, LLC. All rights reserved.
  * $COPYRIGHT$
  *
  * Additional copyrights may follow
@@ -199,9 +199,6 @@ int mca_btl_uct_process_connection_request(mca_btl_uct_module_t *module,
     int32_t ep_flags;
     int rc;
 
-    BTL_VERBOSE(("got connection request for endpoint %p. type = %d. context id = %d",
-                 (void *) endpoint, req->type, req->context_id));
-
     if (NULL == endpoint) {
         BTL_ERROR(("could not create endpoint for connection request"));
         return UCS_ERR_UNREACHABLE;
@@ -211,6 +208,9 @@ int mca_btl_uct_process_connection_request(mca_btl_uct_module_t *module,
 
     ep_flags = opal_atomic_fetch_or_32(&tl_endpoint->flags, MCA_BTL_UCT_ENDPOINT_FLAG_CONN_REC);
 
+    BTL_VERBOSE(("got connection request for endpoint %p. type = %d. context id = %d. ep_flags = %x",
+                 (void *) endpoint, req->type, req->context_id, ep_flags));
+
     if (!(ep_flags & MCA_BTL_UCT_ENDPOINT_FLAG_CONN_REC)) {
         /* create any necessary resources */
         rc = mca_btl_uct_endpoint_connect(module, endpoint, req->context_id, req->ep_addr,
@@ -225,22 +225,13 @@ int mca_btl_uct_process_connection_request(mca_btl_uct_module_t *module,
      * message. this might be overkill but there is little documentation at the UCT level on when
      * an endpoint can be used. */
     if (req->type == 1) {
-        /* remote side is ready */
-        mca_btl_uct_base_frag_t *frag;
-
+        /* remote side is connected */
         /* to avoid a race with send adding pending frags grab the lock here */
         OPAL_THREAD_SCOPED_LOCK(&endpoint->ep_lock, {
             BTL_VERBOSE(("connection ready. sending %" PRIsize_t " frags",
                          opal_list_get_size(&module->pending_frags)));
-            (void) opal_atomic_or_fetch_32(&tl_endpoint->flags,
-                                           MCA_BTL_UCT_ENDPOINT_FLAG_CONN_READY);
-            opal_atomic_wmb();
-
-            OPAL_LIST_FOREACH (frag, &module->pending_frags, mca_btl_uct_base_frag_t) {
-                if (frag->context->context_id == req->context_id && endpoint == frag->endpoint) {
-                    frag->ready = true;
-                }
-            }
+            mca_btl_uct_endpoint_set_flag(module, endpoint, req->context_id, tl_endpoint,
+                                          MCA_BTL_UCT_ENDPOINT_FLAG_CONN_REM_READY);
         });
     }
 

opal/mca/btl/uct/btl_uct_types.h

@@ -27,8 +27,10 @@ struct mca_btl_uct_base_frag_t;
 #    define MCA_BTL_UCT_ENDPOINT_FLAG_CONN_REC 0x1
 /** remote endpoint read */
 #    define MCA_BTL_UCT_ENDPOINT_FLAG_CONN_REM_READY 0x2
+/** local UCT endpoint connected */
+#    define MCA_BTL_UCT_ENDPOINT_FLAG_EP_CONNECTED 0x4
 /** connection was established */
-#    define MCA_BTL_UCT_ENDPOINT_FLAG_CONN_READY 0x4
+#    define MCA_BTL_UCT_ENDPOINT_FLAG_CONN_READY 0x8
 
 /* AM tags */
 /** BTL fragment */