Merge pull request #10298 from hoopoepg/topic/fixed-race-in-datatype-init

PML/UCX/DATATYPE: fixed potential race

Author: awlauria

ompi/mca/pml/ucx/pml_ucx_datatype.c

@@ -214,11 +214,19 @@ pml_ucx_datatype_t *mca_pml_ucx_init_nbx_datatype(ompi_datatype_t *datatype,
 
 ucp_datatype_t mca_pml_ucx_init_datatype(ompi_datatype_t *datatype)
 {
+    static opal_thread_internal_mutex_t lock = OPAL_THREAD_INTERNAL_MUTEX_INITIALIZER;
     size_t size = 0; /* init to suppress compiler warning */
     ucp_datatype_t ucp_datatype;
     ucs_status_t status;
     int ret;
 
+    opal_thread_internal_mutex_lock(&lock);
+
+    if (datatype->pml_data != PML_UCX_DATATYPE_INVALID) {
+        /* datatype is already initialized in concurrent thread */
+        goto out;
+    }
+
     if (mca_pml_ucx_datatype_is_contig(datatype)) {
         ompi_datatype_type_size(datatype, &size);
         ucp_datatype = ucp_dt_make_contig(size);
@@ -271,7 +279,10 @@ ucp_datatype_t mca_pml_ucx_init_datatype(ompi_datatype_t *datatype)
     datatype->pml_data = ucp_datatype;
 #endif
 
-    return ucp_datatype;
+out:
+    opal_thread_internal_mutex_unlock(&lock);
+
+    return mca_pml_ucx_from_ompi_datatype(datatype);
 }
 
 static void mca_pml_ucx_convertor_construct(mca_pml_ucx_convertor_t *convertor)

ompi/mca/pml/ucx/pml_ucx_datatype.h

@@ -44,23 +44,25 @@ OBJ_CLASS_DECLARATION(mca_pml_ucx_convertor_t);
 
 
 __opal_attribute_always_inline__
-static inline ucp_datatype_t mca_pml_ucx_get_datatype(ompi_datatype_t *datatype)
+static inline ucp_datatype_t
+mca_pml_ucx_from_ompi_datatype(ompi_datatype_t *datatype)
 {
 #ifdef HAVE_UCP_REQUEST_PARAM_T
-    pml_ucx_datatype_t *ucp_type = (pml_ucx_datatype_t*)datatype->pml_data;
-
-    if (OPAL_LIKELY(ucp_type != PML_UCX_DATATYPE_INVALID)) {
-        return ucp_type->datatype;
-    }
+    return ((pml_ucx_datatype_t*)datatype->pml_data)->datatype;
 #else
-    ucp_datatype_t ucp_type = datatype->pml_data;
+    return (ucp_datatype_t)datatype->pml_data;
+#endif
+}
 
-    if (OPAL_LIKELY(ucp_type != PML_UCX_DATATYPE_INVALID)) {
-        return ucp_type;
+
+__opal_attribute_always_inline__
+static inline ucp_datatype_t mca_pml_ucx_get_datatype(ompi_datatype_t *datatype)
+{
+    if (OPAL_UNLIKELY(datatype->pml_data == PML_UCX_DATATYPE_INVALID)) {
+        return mca_pml_ucx_init_datatype(datatype);
     }
-#endif
 
-    return mca_pml_ucx_init_datatype(datatype);
+    return mca_pml_ucx_from_ompi_datatype(datatype);
 }
 
 #ifdef HAVE_UCP_REQUEST_PARAM_T