Add init/teardown functions: umfInit() and umfTearDown()

- umfInit() increments the usage reference counter
  and initializes the global state of libumf
  if the usage reference counter was equal 0.
  Returns 0 on success or -1 on failure.

- umfTearDown() decrements the usage reference counter
  and destroys the global state of libumf
  if the usage reference counter is equal 0.

Fixes: #514

Signed-off-by: Lukasz Dorau <[email protected]>

Author: ldorau

include/umf.h

@@ -15,4 +15,13 @@
 #include <umf/mempolicy.h>
 #include <umf/memspace.h>
 
+// Increment the usage reference counter and initialize the global state of libumf
+// if the usage reference counter was equal 0.
+// Returns 0 on success or -1 on failure.
+int umfInit(void);
+
+// Decrement the usage reference counter and destroy the global state of libumf
+// if the usage reference counter is equal 0.
+void umfTearDown(void);
+
 #endif /* UMF_UNIFIED_MEMORY_FRAMEWORK_H */

src/CMakeLists.txt

@@ -53,6 +53,7 @@ set(BA_SOURCES
 
 set(UMF_SOURCES
     ${BA_SOURCES}
+    libumf.c
     ipc.c
     memory_pool.c
     memory_provider.c

src/base_alloc/base_alloc_global.c

@@ -62,10 +62,6 @@ static void umf_ba_create_global(void) {
 
     size_t smallestSize = BASE_ALLOC.ac_sizes[0];
     BASE_ALLOC.smallest_ac_size_log2 = log2Utils(smallestSize);
-
-#if defined(_WIN32) && !defined(UMF_SHARED_LIBRARY)
-    atexit(umf_ba_destroy_global);
-#endif
 }
 
 // returns index of the allocation class for a given size

src/base_alloc/base_alloc_linux.c

@@ -18,14 +18,6 @@
 static UTIL_ONCE_FLAG Page_size_is_initialized = UTIL_ONCE_FLAG_INIT;
 static size_t Page_size;
 
-// The highest possible priority (101) is used, because the constructor should be called
-// as the first one and the destructor as the last one in order to avoid use-after-free.
-void __attribute__((constructor(101))) umf_ba_constructor(void) {}
-
-void __attribute__((destructor(101))) umf_ba_destructor(void) {
-    umf_ba_destroy_global();
-}
-
 void *ba_os_alloc(size_t size) {
     return mmap(NULL, size, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS,
                 -1, 0);

src/libumf.c

@@ -0,0 +1,52 @@
+/*
+ *
+ * Copyright (C) 2024 Intel Corporation
+ *
+ * Under the Apache License v2.0 with LLVM Exceptions. See LICENSE.TXT.
+ * SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+ *
+ */
+
+#include <stddef.h>
+
+#include "base_alloc_global.h"
+#include "memspace_internal.h"
+#include "provider_tracking.h"
+#include "topology.h"
+#include "utils_log.h"
+
+umf_memory_tracker_handle_t TRACKER = NULL;
+
+static unsigned long long umfRefCount = 0;
+
+// Increment the usage reference counter and initialize the global state of libumf
+// if the usage reference counter was equal 0.
+// Returns 0 on success or -1 on failure.
+int umfInit(void) {
+    if (util_fetch_and_add64(&umfRefCount, 1) == 0) {
+        util_log_init();
+        TRACKER = umfMemoryTrackerCreate();
+    }
+
+    return (TRACKER) ? 0 : -1;
+}
+
+// Decrement the usage reference counter and destroy the global state of libumf
+// if the usage reference counter is equal 0.
+void umfTearDown(void) {
+    if (util_fetch_and_add64(&umfRefCount, -1) == 1) {
+#ifndef _WIN32
+        umfMemspaceHostAllDestroy();
+        umfMemspaceHighestCapacityDestroy();
+        umfMemspaceHighestBandwidthDestroy();
+#endif
+        umfDestroyTopology();
+
+        // make sure TRACKER is not used after being destroyed
+        umf_memory_tracker_handle_t t = TRACKER;
+        TRACKER = NULL;
+        umfMemoryTrackerDestroy(t);
+
+        umf_ba_destroy_global();
+    }
+}

src/libumf.def.in

@@ -10,6 +10,8 @@ VERSION 1.0
 
 EXPORTS
     DllMain
+    umfInit
+    umfTearDown
     umfCloseIPCHandle
     umfFree
     umfGetIPCHandle

src/libumf.map

@@ -4,6 +4,8 @@
 
 UMF_1.0 {
     global:
+        umfInit;
+        umfTearDown;
         umfCloseIPCHandle;
         umfFree;
         umfGetIPCHandle;

src/libumf_linux.c

@@ -7,31 +7,11 @@
  *
  */
 
-#include <stddef.h>
+#include <umf.h>
 
-#include "base_alloc_global.h"
-#include "memspace_internal.h"
-#include "provider_tracking.h"
-#include "topology.h"
-#include "utils_log.h"
+void __attribute__((constructor)) umfCreate(void) { (void)umfInit(); }
 
-umf_memory_tracker_handle_t TRACKER = NULL;
-
-void __attribute__((constructor)) umfCreate(void) {
-    util_log_init();
-    TRACKER = umfMemoryTrackerCreate();
-}
-
-void __attribute__((destructor)) umfDestroy(void) {
-    umf_memory_tracker_handle_t t = TRACKER;
-    // make sure TRACKER is not used after being destroyed
-    TRACKER = NULL;
-    umfMemoryTrackerDestroy(t);
-    umfMemspaceHostAllDestroy();
-    umfMemspaceHighestCapacityDestroy();
-    umfMemspaceHighestBandwidthDestroy();
-    umfDestroyTopology();
-}
+void __attribute__((destructor)) umfDestroy(void) { umfTearDown(); }
 
 void libumfInit(void) {
     // do nothing, additional initialization not needed

src/libumf_windows.c

@@ -7,49 +7,39 @@
  *
  */
 
-#include <stdlib.h>
 #include <windows.h>
 
-#include "base_alloc_global.h"
-#include "provider_tracking.h"
-#include "utils_log.h"
+#include <umf.h>
 
-umf_memory_tracker_handle_t TRACKER = NULL;
+#if defined(UMF_SHARED_LIBRARY) /* SHARED LIBRARY */
 
-static void umfCreate(void) {
-    util_log_init();
-    TRACKER = umfMemoryTrackerCreate();
-}
-
-static void umfDestroy(void) {
-    umfMemoryTrackerDestroy(TRACKER);
-    umf_ba_destroy_global();
-}
-
-#if defined(UMF_SHARED_LIBRARY)
 BOOL APIENTRY DllMain(HINSTANCE hinstDLL, DWORD fdwReason, LPVOID lpvReserved) {
-    if (fdwReason == DLL_PROCESS_DETACH) {
-        umfDestroy();
-    } else if (fdwReason == DLL_PROCESS_ATTACH) {
-        umfCreate();
+    int ret = 0;
+    if (fdwReason == DLL_PROCESS_ATTACH) {
+        ret = umfInit();
+    } else if (fdwReason == DLL_PROCESS_DETACH) {
+        umfTearDown();
     }
-    return TRUE;
+    return (ret == 0) ? TRUE : FALSE;
 }
 
 void libumfInit(void) {
     // do nothing, additional initialization not needed
 }
-#else
+
+#else /* STATIC LIBRARY */
+
 INIT_ONCE init_once_flag = INIT_ONCE_STATIC_INIT;
 
 BOOL CALLBACK initOnceCb(PINIT_ONCE InitOnce, PVOID Parameter,
                          PVOID *lpContext) {
-    umfCreate();
-    atexit(umfDestroy);
-    return TRACKER ? TRUE : FALSE;
+    int ret = umfInit();
+    atexit(umfTearDown);
+    return (ret == 0) ? TRUE : FALSE;
 }
 
 void libumfInit(void) {
     InitOnceExecuteOnce(&init_once_flag, initOnceCb, NULL, NULL);
 }
+
 #endif

src/memspaces/memspace_highest_bandwidth.c

@@ -90,10 +90,6 @@ static void umfMemspaceHighestBandwidthInit(void) {
             ret);
         assert(ret == UMF_RESULT_ERROR_NOT_SUPPORTED);
     }
-
-#if defined(_WIN32) && !defined(UMF_SHARED_LIBRARY)
-    atexit(umfMemspaceHighestBandwidthDestroy);
-#endif
 }
 
 umf_memspace_handle_t umfMemspaceHighestBandwidthGet(void) {

src/memspaces/memspace_highest_capacity.c

@@ -62,10 +62,6 @@ static void umfMemspaceHighestCapacityInit(void) {
         umfMemspaceHighestCapacityCreate(&UMF_MEMSPACE_HIGHEST_CAPACITY);
     assert(ret == UMF_RESULT_SUCCESS);
     (void)ret;
-
-#if defined(_WIN32) && !defined(UMF_SHARED_LIBRARY)
-    atexit(umfMemspaceHostAllDestroy);
-#endif
 }
 
 umf_memspace_handle_t umfMemspaceHighestCapacityGet(void) {

src/topology.c

@@ -34,10 +34,6 @@ static void umfCreateTopology(void) {
         hwloc_topology_destroy(topology);
         topology = NULL;
     }
-
-#if defined(_WIN32) && !defined(UMF_SHARED_LIBRARY)
-    atexit(umfDestroyTopology);
-#endif
 }
 
 hwloc_topology_t umfGetTopology(void) {