make umfInit threadsafe

Author: lplewa

src/libumf.c

@@ -8,6 +8,7 @@
  */
 
 #include <stddef.h>
+#include <stdint.h>
 #include <string.h>
 
 #include "base_alloc_global.h"
@@ -19,6 +20,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,35 +36,46 @@ static umf_ctl_node_t CTL_NODE(umf)[] = {CTL_CHILD(provider), CTL_CHILD(pool),
 void initialize_global_ctl(void) { CTL_REGISTER_MODULE(NULL, umf); }
 
 int umfInit(void) {
-    if (utils_fetch_and_add_u64(&umfRefCount, 1) == 0) {
-        utils_log_init();
-        TRACKER = umfMemoryTrackerCreate();
-        if (!TRACKER) {
-            LOG_ERR("Failed to create memory tracker");
-            return -1;
+    uint64_t refCount;
+    do {
+        if (utils_fetch_and_add_u64(&umfRefCount, 1) == 0) {
+            utils_log_init();
+            TRACKER = umfMemoryTrackerCreate();
+            if (!TRACKER) {
+                LOG_ERR("Failed to create memory tracker");
+                utils_atomic_decrement_size_t(&umfRefCount);
+                return -1;
+            }
+
+            LOG_DEBUG("UMF tracker created");
+
+            umf_result_t 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 -1;
+            }
+
+            LOG_DEBUG("UMF IPC cache initialized");
+            initialize_global_ctl();
+
+            utils_atomic_increment_size_t(&umfRefCount);
         }
 
-        LOG_DEBUG("UMF tracker created");
+        do {
+            utils_atomic_load_acquire_u64(&umfRefCount, &refCount);
+        } while (refCount == 1);
 
-        umf_result_t umf_result = umfIpcCacheGlobalInit();
-        if (umf_result != UMF_RESULT_SUCCESS) {
-            LOG_ERR("Failed to initialize IPC cache");
-            return -1;
-        }
-
-        LOG_DEBUG("UMF IPC cache initialized");
-        initialize_global_ctl();
-    }
+    } while (refCount > 2);
 
-    if (TRACKER) {
-        LOG_DEBUG("UMF library initialized");
-    }
+    LOG_DEBUG("UMF library initialized");
 
     return 0;
 }
 
 void 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();
@@ -94,6 +107,7 @@ void umfTearDown(void) {
         fini_cu_global_state();
         fini_tbb_global_state();
         LOG_DEBUG("UMF library finalized");
+        utils_atomic_decrement_size_t(&umfRefCount);
     }
 }