Merge pull request #530 from ldorau/Add_init_teardown_functions_umfInit_and_umfTearDown

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

Author: bratpiorka

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

@@ -10,6 +10,7 @@
 #include <assert.h>
 #include <stdio.h>
 #include <stdlib.h>
+#include <string.h>
 
 #include "base_alloc.h"
 #include "base_alloc_global.h"
@@ -45,6 +46,10 @@ void umf_ba_destroy_global(void) {
             BASE_ALLOC.ac[i] = NULL;
         }
     }
+
+    // portable version of "ba_is_initialized = UTIL_ONCE_FLAG_INIT;"
+    static UTIL_ONCE_FLAG is_initialized = UTIL_ONCE_FLAG_INIT;
+    memcpy(&ba_is_initialized, &is_initialized, sizeof(ba_is_initialized));
 }
 
 static void umf_ba_create_global(void) {
@@ -62,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

@@ -78,6 +78,11 @@ void umfMemspaceHighestBandwidthDestroy(void) {
     if (UMF_MEMSPACE_HIGHEST_BANDWIDTH) {
         umfMemspaceDestroy(UMF_MEMSPACE_HIGHEST_BANDWIDTH);
         UMF_MEMSPACE_HIGHEST_BANDWIDTH = NULL;
+
+        // portable version of "UMF_MEMSPACE_HBW_INITIALIZED = UTIL_ONCE_FLAG_INIT;"
+        static UTIL_ONCE_FLAG is_initialized = UTIL_ONCE_FLAG_INIT;
+        memcpy(&UMF_MEMSPACE_HBW_INITIALIZED, &is_initialized,
+               sizeof(UMF_MEMSPACE_HBW_INITIALIZED));
     }
 }
 
@@ -90,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

@@ -54,6 +54,11 @@ void umfMemspaceHighestCapacityDestroy(void) {
     if (UMF_MEMSPACE_HIGHEST_CAPACITY) {
         umfMemspaceDestroy(UMF_MEMSPACE_HIGHEST_CAPACITY);
         UMF_MEMSPACE_HIGHEST_CAPACITY = NULL;
+
+        // portable version of "UMF_MEMSPACE_HIGHEST_CAPACITY_INITIALIZED = UTIL_ONCE_FLAG_INIT;"
+        static UTIL_ONCE_FLAG is_initialized = UTIL_ONCE_FLAG_INIT;
+        memcpy(&UMF_MEMSPACE_HIGHEST_CAPACITY_INITIALIZED, &is_initialized,
+               sizeof(UMF_MEMSPACE_HIGHEST_CAPACITY_INITIALIZED));
     }
 }
 
@@ -62,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/memspaces/memspace_host_all.c

@@ -73,6 +73,11 @@ void umfMemspaceHostAllDestroy(void) {
     if (UMF_MEMSPACE_HOST_ALL) {
         umfMemspaceDestroy(UMF_MEMSPACE_HOST_ALL);
         UMF_MEMSPACE_HOST_ALL = NULL;
+
+        // portable version of "UMF_MEMSPACE_HOST_ALL_INITIALIZED = UTIL_ONCE_FLAG_INIT;"
+        static UTIL_ONCE_FLAG is_initialized = UTIL_ONCE_FLAG_INIT;
+        memcpy(&UMF_MEMSPACE_HOST_ALL_INITIALIZED, &is_initialized,
+               sizeof(UMF_MEMSPACE_HOST_ALL_INITIALIZED));
     }
 }
 

src/provider/provider_level_zero.c

@@ -124,6 +124,10 @@ void ze_memory_provider_finalize(void *provider) {
 
     util_init_once(&ze_is_initialized, init_ze_global_state);
     umf_ba_global_free(provider);
+
+    // portable version of "ze_is_initialized = UTIL_ONCE_FLAG_INIT;"
+    static UTIL_ONCE_FLAG is_initialized = UTIL_ONCE_FLAG_INIT;
+    memcpy(&ze_is_initialized, &is_initialized, sizeof(ze_is_initialized));
 }
 
 static umf_result_t ze_memory_provider_alloc(void *provider, size_t size,

src/topology.c

@@ -19,6 +19,11 @@ static UTIL_ONCE_FLAG topology_initialized = UTIL_ONCE_FLAG_INIT;
 void umfDestroyTopology(void) {
     if (topology) {
         hwloc_topology_destroy(topology);
+
+        // portable version of "topology_initialized = UTIL_ONCE_FLAG_INIT;"
+        static UTIL_ONCE_FLAG is_initialized = UTIL_ONCE_FLAG_INIT;
+        memcpy(&topology_initialized, &is_initialized,
+               sizeof(topology_initialized));
     }
 }
 
@@ -34,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) {

test/CMakeLists.txt

@@ -279,3 +279,14 @@ else()
     message(
         STATUS "IPC shared memory test is supported on Linux only - skipping")
 endif()
+
+if(LINUX
+   AND UMF_BUILD_SHARED_LIBRARY
+   AND UMF_POOL_SCALABLE_ENABLED)
+    add_umf_test(NAME init_teardown SRCS test_init_teardown.c)
+    # append LD_LIBRARY_PATH to the libumf
+    set_property(
+        TEST umf-init_teardown
+        PROPERTY ENVIRONMENT_MODIFICATION
+                 "LD_LIBRARY_PATH=path_list_append:${CMAKE_BINARY_DIR}/lib")
+endif()