add free handle POC

Author: bratpiorka

include/umf/proxy_lib_handlers.h

@@ -0,0 +1,27 @@
+/*
+ * 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
+*/
+
+#ifndef UMF_PROXY_LIB_HANDLERS_H
+#define UMF_PROXY_LIB_HANDLERS_H 1
+
+#include <umf/memory_pool.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// malloc API handlers
+typedef void (*umf_proxy_lib_handler_free_pre_t)(
+    void *ptr, umf_memory_pool_handle_t pool);
+
+void umfSetProxyLibHandlerFreePre(umf_proxy_lib_handler_free_pre_t handler);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif /* UMF_PROXY_LIB_HANDLERS_H */

src/proxy_lib/proxy_lib.c

@@ -50,9 +50,9 @@
 #include <umf/memory_pool.h>
 #include <umf/memory_provider.h>
 #include <umf/providers/provider_os_memory.h>
+#include <umf/proxy_lib_handlers.h>
 
 #include "base_alloc_linear.h"
-#include "proxy_lib.h"
 #include "utils_common.h"
 #include "utils_load_library.h"
 #include "utils_log.h"
@@ -135,6 +135,9 @@ static __TLS int was_called_from_umfPool = 0;
 // TODO remove this WA when the issue is fixed.
 static __TLS int was_called_from_malloc_usable_size = 0;
 
+// malloc API handlers
+static umf_proxy_lib_handler_free_pre_t Handler_free_pre = NULL;
+
 /*****************************************************************************/
 /*** The constructor and destructor of the proxy library *********************/
 /*****************************************************************************/
@@ -391,11 +394,19 @@ void free(void *ptr) {
         return;
     }
 
+    // NOTE: for system allocations made during UMF and Proxy Lib
+    // initialisation, we never call free handlers, as they should handle
+    // only user-made allocations
     if (ba_leak_free(ptr) == 0) {
         return;
     }
 
-    if (Proxy_pool && (umfPoolByPtr(ptr) == Proxy_pool)) {
+    umf_memory_pool_handle_t pool = umfPoolByPtr(ptr);
+    if (Proxy_pool && (pool == Proxy_pool)) {
+        if (Handler_free_pre) {
+            Handler_free_pre(ptr, pool);
+        }
+
         if (umfPoolFree(Proxy_pool, ptr) != UMF_RESULT_SUCCESS) {
             LOG_ERR("umfPoolFree() failed");
         }
@@ -404,6 +415,9 @@ void free(void *ptr) {
 
 #ifndef _WIN32
     if (Size_threshold_value) {
+        if (Handler_free_pre) {
+            Handler_free_pre(ptr, NULL);
+        }
         System_free(ptr);
         return;
     }
@@ -555,3 +569,9 @@ void *_aligned_offset_recalloc(void *ptr, size_t num, size_t size,
 }
 
 #endif
+
+// malloc API handlers
+
+void umfSetProxyLibHandlerFreePre(umf_proxy_lib_handler_free_pre_t handler) {
+    Handler_free_pre = handler;
+}

src/proxy_lib/proxy_lib.def

@@ -21,3 +21,5 @@ EXPORTS
 	_aligned_offset_malloc
 	_aligned_offset_realloc
 	_aligned_offset_recalloc
+; handlers
+	umfSetProxyLibHandlerFreePre

src/proxy_lib/proxy_lib.h

@@ -4,14 +4,16 @@
 
 # linker VERSION script
 
-{
+PROXY_LIB_0.10 {
     global:
         aligned_alloc;
         calloc;
         free;
         malloc;
         malloc_usable_size;
         realloc;
+    # handlers
+        umfSetProxyLibHandlerFreePre;
     local:
         *;
 };

src/proxy_lib/proxy_lib.map

@@ -7,7 +7,10 @@
  *
  */
 
-#include "proxy_lib.h"
+#include <umf/proxy_lib_handlers.h>
+
+void proxy_lib_create_common(void);
+void proxy_lib_destroy_common(void);
 
 // The priority 102 is used, because the constructor should be called as the second one
 // (just after the first constructor of the base allocator with priority 101)

src/proxy_lib/proxy_lib_linux.c

@@ -9,7 +9,10 @@
 
 #include <Windows.h>
 
-#include "proxy_lib.h"
+#include <umf/proxy_lib_handlers.h>
+
+void proxy_lib_create_common(void);
+void proxy_lib_destroy_common(void);
 
 static void proxy_lib_create(void) { proxy_lib_create_common(); }
 

src/proxy_lib/proxy_lib_windows.c

@@ -447,14 +447,19 @@ add_umf_test(
 if(UMF_PROXY_LIB_ENABLED AND UMF_BUILD_SHARED_LIBRARY)
     add_umf_test(
         NAME proxy_lib_basic
-        SRCS ${BA_SOURCES_FOR_TEST} test_proxy_lib.cpp
+        SRCS ${BA_SOURCES_FOR_TEST} proxy_lib/proxy_lib.cpp
+        LIBS ${UMF_UTILS_FOR_TEST} umf_proxy)
+
+    add_umf_test(
+        NAME proxy_lib_handlers
+        SRCS ${BA_SOURCES_FOR_TEST} proxy_lib/proxy_lib_handlers.cpp
         LIBS ${UMF_UTILS_FOR_TEST} umf_proxy)
 
     # TODO enable this test on Windows
     if(LINUX)
         add_umf_test(
             NAME test_proxy_lib_size_threshold
-            SRCS ${BA_SOURCES_FOR_TEST} test_proxy_lib_size_threshold.cpp
+            SRCS ${BA_SOURCES_FOR_TEST} proxy_lib/proxy_lib_size_threshold.cpp
             LIBS ${UMF_UTILS_FOR_TEST} umf_proxy)
         set_property(TEST umf-test_proxy_lib_size_threshold
                      PROPERTY ENVIRONMENT UMF_PROXY="size.threshold=64")

test/CMakeLists.txt

@@ -0,0 +1,42 @@
+/*
+ * 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
+*/
+
+#if defined(__APPLE__)
+#include <malloc/malloc.h>
+#else
+#include <malloc.h>
+#endif
+
+#include <umf/proxy_lib_handlers.h>
+
+#include "base.hpp"
+#include "test_helpers.h"
+#include "utils_common.h"
+
+using umf_test::test;
+
+#define SIZE_64 64
+#define ALIGN_1024 1024
+
+static int free_cnt = 0;
+void handler_free_pre(void *ptr, umf_memory_pool_handle_t pool) {
+    (void)ptr;
+    (void)pool;
+    free_cnt += 1;
+}
+
+TEST_F(test, proxyLib_handlers_free) {
+
+    void *ptr = ::malloc(SIZE_64);
+    ASSERT_NE(ptr, nullptr);
+
+    umfSetProxyLibHandlerFreePre(handler_free_pre);
+    ASSERT_EQ(free_cnt, 0);
+
+    ::free(ptr);
+    ASSERT_EQ(free_cnt, 1);
+}