From 93b8be2126c921dd0e7b5814b144213e671d94eb Mon Sep 17 00:00:00 2001
From: Lukasz Dorau <lukasz.dorau@intel.com>
Date: Tue, 12 Nov 2024 09:22:10 +0100
Subject: [PATCH] Do not destroy tracker in umfTearDown() under the proxy
 library

Signed-off-by: Lukasz Dorau <lukasz.dorau@intel.com>
---
 src/base_alloc/base_alloc_global.c |  2 ++
 src/libumf.c                       | 20 ++++++++++++++++++++
 src/utils/utils_common.h           |  8 ++++++++
 3 files changed, 30 insertions(+)

diff --git a/src/base_alloc/base_alloc_global.c b/src/base_alloc/base_alloc_global.c
index ec6bc9fcb..196a282ba 100644
--- a/src/base_alloc/base_alloc_global.c
+++ b/src/base_alloc/base_alloc_global.c
@@ -67,6 +67,8 @@ static void umf_ba_create_global(void) {
 
     size_t smallestSize = BASE_ALLOC.ac_sizes[0];
     BASE_ALLOC.smallest_ac_size_log2 = log2Utils(smallestSize);
+
+    LOG_DEBUG("UMF base allocator created");
 }
 
 // returns index of the allocation class for a given size
diff --git a/src/libumf.c b/src/libumf.c
index 2fcda07a0..33cc98107 100644
--- a/src/libumf.c
+++ b/src/libumf.c
@@ -12,6 +12,7 @@
 #include "base_alloc_global.h"
 #include "memspace_internal.h"
 #include "provider_tracking.h"
+#include "utils_common.h"
 #include "utils_log.h"
 #if !defined(UMF_NO_HWLOC)
 #include "topology.h"
@@ -25,6 +26,11 @@ int umfInit(void) {
     if (utils_fetch_and_add64(&umfRefCount, 1) == 0) {
         utils_log_init();
         TRACKER = umfMemoryTrackerCreate();
+        LOG_DEBUG("UMF tracker created");
+    }
+
+    if (TRACKER) {
+        LOG_DEBUG("UMF library initialized");
     }
 
     return (TRACKER) ? 0 : -1;
@@ -39,12 +45,26 @@ void umfTearDown(void) {
         umfMemspaceLowestLatencyDestroy();
         umfDestroyTopology();
 #endif
+
+        if (utils_is_running_in_proxy_lib_with_size_threshold()) {
+            // We cannot destroy the TRACKER nor the base allocator
+            // when we are running in the proxy library with a size threshold,
+            // because it could lead to calling the system free() with an invalid pointer
+            // and a segfault as a result.
+            goto fini_umfTearDown;
+        }
+
         // make sure TRACKER is not used after being destroyed
         umf_memory_tracker_handle_t t = TRACKER;
         TRACKER = NULL;
         umfMemoryTrackerDestroy(t);
+        LOG_DEBUG("UMF tracker destroyed");
 
         umf_ba_destroy_global();
+        LOG_DEBUG("UMF base allocator destroyed");
+
+    fini_umfTearDown:
+        LOG_DEBUG("UMF library finalized");
     }
 }
 
diff --git a/src/utils/utils_common.h b/src/utils/utils_common.h
index a168873bf..d5917d4fa 100644
--- a/src/utils/utils_common.h
+++ b/src/utils/utils_common.h
@@ -70,6 +70,14 @@ static inline int utils_is_running_in_proxy_lib(void) {
     return utils_env_var_get_str("LD_PRELOAD", "libumf_proxy.so") ? 1 : 0;
 }
 
+// check if we are running in the proxy library with a size threshold
+static inline int utils_is_running_in_proxy_lib_with_size_threshold(void) {
+    return (utils_env_var_get_str("LD_PRELOAD", "libumf_proxy.so") &&
+            utils_env_var_get_str("UMF_PROXY", "size.threshold="))
+               ? 1
+               : 0;
+}
+
 // utils_parse_var - Parses var for a prefix,
 //                   optionally identifying a following argument.
 // Parameters: