make umfInit threadsafe

Author: lplewa

src/libumf.c

@@ -19,6 +19,7 @@
 #include "provider_level_zero_internal.h"
 #include "provider_tracking.h"
 #include "utils_common.h"
+#include "utils_concurrency.h"
 #include "utils_log.h"
 #if !defined(UMF_NO_HWLOC)
 #include "topology.h"
@@ -34,36 +35,49 @@ static umf_ctl_node_t CTL_NODE(umf)[] = {CTL_CHILD(provider), CTL_CHILD(pool),
 void initialize_global_ctl(void) { CTL_REGISTER_MODULE(NULL, umf); }
 
 umf_result_t umfInit(void) {
-    if (utils_fetch_and_add_u64(&umfRefCount, 1) == 0) {
-        utils_log_init();
-        umf_result_t umf_result = umfMemoryTrackerCreate(&TRACKER);
-        if (umf_result != UMF_RESULT_SUCCESS) {
-            LOG_ERR("Failed to create memory tracker");
-            return umf_result;
+    uint64_t refCount;
+    uint64_t one = 1;
+    do {
+        utils_atomic_load_acquire_u64(&umfRefCount, &refCount);
+        if (refCount == 0 &&
+            utils_compare_exchange_u64(&umfRefCount, &refCount, &one)) {
+            utils_log_init();
+            umf_result_t umf_result = umfMemoryTrackerCreate(&TRACKER);
+            if (!TRACKER) {
+                LOG_ERR("Failed to create memory tracker");
+                utils_atomic_decrement_size_t(&umfRefCount);
+                return umf_result;
+            }
+
+            LOG_DEBUG("UMF tracker created");
+
+            umf_result = umfIpcCacheGlobalInit();
+            if (umf_result != UMF_RESULT_SUCCESS) {
+                LOG_ERR("Failed to initialize IPC cache");
+                umfMemoryTrackerDestroy(TRACKER);
+                utils_atomic_decrement_size_t(&umfRefCount);
+                return umf_result;
+            }
+
+            LOG_DEBUG("UMF IPC cache initialized");
+            initialize_global_ctl();
+
+        } else {
+            do {
+                utils_atomic_load_acquire_u64(&umfRefCount, &refCount);
+            } while (refCount == 1);
         }
+        utils_atomic_increment_u64(&umfRefCount);
 
-        LOG_DEBUG("UMF tracker created");
+    } while (refCount < 2);
 
-        umf_result = umfIpcCacheGlobalInit();
-        if (umf_result != UMF_RESULT_SUCCESS) {
-            LOG_ERR("Failed to initialize IPC cache");
-            umfMemoryTrackerDestroy(TRACKER);
-            return umf_result;
-        }
-
-        LOG_DEBUG("UMF IPC cache initialized");
-        initialize_global_ctl();
-    }
-
-    if (TRACKER) {
-        LOG_DEBUG("UMF library initialized");
-    }
+    LOG_DEBUG("UMF library initialized");
 
     return UMF_RESULT_SUCCESS;
 }
 
 umf_result_t umfTearDown(void) {
-    if (utils_fetch_and_sub_u64(&umfRefCount, 1) == 1) {
+    if (utils_fetch_and_sub_u64(&umfRefCount, 1) == 2) {
 #if !defined(_WIN32) && !defined(UMF_NO_HWLOC)
         umfMemspaceHostAllDestroy();
         umfMemspaceHighestCapacityDestroy();
@@ -95,6 +109,7 @@ umf_result_t umfTearDown(void) {
         fini_cu_global_state();
         fini_tbb_global_state();
         LOG_DEBUG("UMF library finalized");
+        utils_atomic_decrement_u64(&umfRefCount);
     }
     return UMF_RESULT_SUCCESS;
 }