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.
  It must be called just after dlopen()
  and it is not required in other scenarios.

- umfTearDown() decrements the usage reference counter
  and destroys the global state of libumf
  if the usage reference counter is equal 0.
  It must be called just before dlclose()
  and it is not required in other scenarios.

Fixes: #514

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

Author: ldorau

include/umf.h

@@ -15,4 +15,17 @@
 #include <umf/mempolicy.h>
 #include <umf/memspace.h>
 
+///
+/// @brief  Increment the usage reference counter and initialize the global state of libumf
+///         if the usage reference counter was equal to 0.
+///         It must be called just after dlopen() and it is not required in other scenarios.
+/// @return 0 on success or -1 on failure.
+int umfInit(void);
+
+///
+/// @brief Decrement the usage reference counter and destroy the global state of libumf
+///        if the usage reference counter is equal to 0.
+///        It must be called just before dlclose() and it is not required in other scenarios.
+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

@@ -67,10 +67,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,46 @@
+/*
+ *
+ * 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;
+
+int umfInit(void) {
+    if (util_fetch_and_add64(&umfRefCount, 1) == 0) {
+        util_log_init();
+        TRACKER = umfMemoryTrackerCreate();
+    }
+
+    return (TRACKER) ? 0 : -1;
+}
+
+void umfTearDown(void) {
+    if (util_fetch_and_add64(&umfRefCount, -1) == 1) {
+#ifndef _WIN32
+        umfMemspaceHostAllDestroy();
+        umfMemspaceHighestCapacityDestroy();
+        umfMemspaceHighestBandwidthDestroy();
+        umfDestroyTopology();
+#endif
+        // 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,38 @@
  *
  */
 
-#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();
+    if (fdwReason == DLL_PROCESS_ATTACH) {
+        (void)umfInit();
+    } else if (fdwReason == DLL_PROCESS_DETACH) {
+        umfTearDown();
     }
     return TRUE;
 }
 
 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

@@ -95,10 +95,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

@@ -67,10 +67,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

@@ -39,10 +39,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) {