[POC] add free handler to the proxy lib

include/umf/proxy_lib_handlers.h

@@ -0,0 +1,54 @@
+/*
+ * 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
+
+// TODO - improve and cleanup comments
+
+// malloc API handlers
+// NOTE - in malloc/aligned_malloc pre handlers the default pool could be
+// changed along with the requested size and alignment
+// - in free pre handler user can update the ptr
+typedef void (*umf_proxy_lib_handler_malloc_pre_t)(void *user_data,
+                                                   size_t *size);
+typedef void (*umf_proxy_lib_handler_aligned_malloc_pre_t)(
+    void *user_data, umf_memory_pool_handle_t *pool, size_t *size,
+    size_t *alignment);
+typedef void (*umf_proxy_lib_handler_free_pre_t)(void *user_data, void **ptr,
+                                                 umf_memory_pool_handle_t pool);
+
+void umfSetProxyLibHandlerMallocPre(umf_proxy_lib_handler_malloc_pre_t handler,
+                                    void *user_data);
+void umfSetProxyLibHandlerAlignedMallocPre(
+    umf_proxy_lib_handler_aligned_malloc_pre_t handler, void *user_data);
+void umfSetProxyLibHandlerFreePre(umf_proxy_lib_handler_free_pre_t handler,
+                                  void *user_data);
+
+// NOTE - in the malloc/aligned_malloc post handlers the pointer to allocated
+// data could be changed by the user
+typedef void (*umf_proxy_lib_handler_malloc_post_t)(
+    void *user_data, void **ptr, umf_memory_pool_handle_t pool);
+typedef void (*umf_proxy_lib_handler_aligned_malloc_post_t)(
+    void *user_data, void **ptr, umf_memory_pool_handle_t pool);
+
+void umfSetProxyLibHandlerMallocPost(
+    umf_proxy_lib_handler_malloc_post_t handler, void *user_data);
+void umfSetProxyLibHandlerAlignedMallocPost(
+    umf_proxy_lib_handler_aligned_malloc_post_t handler, void *user_data);
+
+#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,22 @@ 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_malloc_pre_t Handler_malloc_pre = NULL;
+static umf_proxy_lib_handler_aligned_malloc_pre_t Handler_aligned_malloc_pre =
+    NULL;
+static umf_proxy_lib_handler_free_pre_t Handler_free_pre = NULL;
+
+static umf_proxy_lib_handler_malloc_post_t Handler_malloc_post = NULL;
+static umf_proxy_lib_handler_aligned_malloc_post_t Handler_aligned_malloc_post =
+    NULL;
+
+static void *Handler_malloc_pre_user_data = NULL;
+static void *Handler_aligned_malloc_pre_user_data = NULL;
+static void *Handler_free_pre_user_data = NULL;
+static void *Handler_malloc_post_user_data = NULL;
+static void *Handler_aligned_malloc_post_user_data = NULL;
+
 /*****************************************************************************/
 /*** The constructor and destructor of the proxy library *********************/
 /*****************************************************************************/
@@ -353,20 +369,34 @@ static inline size_t ba_leak_pool_contains_pointer(void *ptr) {
 /*****************************************************************************/
 
 void *malloc(size_t size) {
+    if (Handler_malloc_pre) {
+        Handler_malloc_pre(Handler_malloc_pre_user_data, &size);
+    }
+
+    void *ptr = NULL;
 #ifndef _WIN32
     if (size < Size_threshold_value) {
-        return System_malloc(size);
+        ptr = System_malloc(size);
+        goto handler_post;
     }
 #endif /* _WIN32 */
 
-    if (!was_called_from_umfPool && Proxy_pool) {
+    umf_memory_pool_handle_t pool = Proxy_pool;
+    if (!was_called_from_umfPool && pool) {
         was_called_from_umfPool = 1;
-        void *ptr = umfPoolMalloc(Proxy_pool, size);
+        ptr = umfPoolMalloc(pool, size);
         was_called_from_umfPool = 0;
-        return ptr;
+        goto handler_post;
     }
 
-    return ba_leak_malloc(size);
+    ptr = ba_leak_malloc(size);
+
+handler_post:
+    if (Handler_malloc_post) {
+        Handler_malloc_post(Handler_malloc_post_user_data, &ptr, pool);
+    }
+
+    return ptr;
 }
 
 void *calloc(size_t nmemb, size_t size) {
@@ -387,15 +417,28 @@ void *calloc(size_t nmemb, size_t size) {
 }
 
 void free(void *ptr) {
+    umf_memory_pool_handle_t pool = NULL;
+    if (Handler_free_pre) {
+        pool = umfPoolByPtr(ptr);
+        Handler_free_pre(Handler_free_pre_user_data, &ptr, pool);
+    }
+
     if (ptr == NULL) {
         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)) {
+    if (Proxy_pool && pool == NULL) {
+        pool = umfPoolByPtr(ptr);
+    }
+
+    if (pool == Proxy_pool) {
         if (umfPoolFree(Proxy_pool, ptr) != UMF_RESULT_SUCCESS) {
             LOG_ERR("umfPoolFree() failed");
         }
@@ -448,20 +491,36 @@ void *realloc(void *ptr, size_t size) {
 }
 
 void *aligned_alloc(size_t alignment, size_t size) {
+    umf_memory_pool_handle_t pool = Proxy_pool;
+    if (Handler_aligned_malloc_pre) {
+        Handler_aligned_malloc_pre(Handler_aligned_malloc_pre_user_data, &pool,
+                                   &size, &alignment);
+    }
+
+    void *ptr = NULL;
 #ifndef _WIN32
     if (size < Size_threshold_value) {
-        return System_aligned_alloc(alignment, size);
+        ptr = System_aligned_alloc(alignment, size);
+        goto handler_post;
     }
 #endif /* _WIN32 */
 
     if (!was_called_from_umfPool && Proxy_pool) {
         was_called_from_umfPool = 1;
-        void *ptr = umfPoolAlignedMalloc(Proxy_pool, size, alignment);
+        ptr = umfPoolAlignedMalloc(Proxy_pool, size, alignment);
         was_called_from_umfPool = 0;
-        return ptr;
+        goto handler_post;
     }
 
-    return ba_leak_aligned_alloc(alignment, size);
+    ptr = ba_leak_aligned_alloc(alignment, size);
+
+handler_post:
+    if (Handler_aligned_malloc_post) {
+        Handler_aligned_malloc_post(Handler_aligned_malloc_post_user_data, &ptr,
+                                    pool);
+    }
+
+    return ptr;
 }
 
 #ifdef _WIN32
@@ -555,3 +614,35 @@ void *_aligned_offset_recalloc(void *ptr, size_t num, size_t size,
 }
 
 #endif
+
+// malloc API handlers
+
+void umfSetProxyLibHandlerMallocPre(umf_proxy_lib_handler_malloc_pre_t handler,
+                                    void *user_data) {
+    Handler_malloc_pre = handler;
+    Handler_malloc_pre_user_data = user_data;
+}
+
+void umfSetProxyLibHandlerAlignedMallocPre(
+    umf_proxy_lib_handler_aligned_malloc_pre_t handler, void *user_data) {
+    Handler_aligned_malloc_pre = handler;
+    Handler_aligned_malloc_pre_user_data = user_data;
+}
+
+void umfSetProxyLibHandlerFreePre(umf_proxy_lib_handler_free_pre_t handler,
+                                  void *user_data) {
+    Handler_free_pre = handler;
+    Handler_free_pre_user_data = user_data;
+}
+
+void umfSetProxyLibHandlerMallocPost(
+    umf_proxy_lib_handler_malloc_post_t handler, void *user_data) {
+    Handler_malloc_post = handler;
+    Handler_malloc_post_user_data = user_data;
+}
+
+void umfSetProxyLibHandlerAlignedMallocPost(
+    umf_proxy_lib_handler_aligned_malloc_post_t handler, void *user_data) {
+    Handler_aligned_malloc_post = handler;
+    Handler_aligned_malloc_post_user_data = user_data;
+}

src/proxy_lib/proxy_lib.def

@@ -21,3 +21,9 @@ EXPORTS
 	_aligned_offset_malloc
 	_aligned_offset_realloc
 	_aligned_offset_recalloc
+; handlers
+	umfSetProxyLibHandlerMallocPre
+	umfSetProxyLibHandlerAlignedMallocPre
+	umfSetProxyLibHandlerFreePre
+	umfSetProxyLibHandlerMallocPost
+	umfSetProxyLibHandlerAlignedMallocPost

src/proxy_lib/proxy_lib.map

@@ -12,6 +12,12 @@
         malloc;
         malloc_usable_size;
         realloc;
+    # handlers
+        umfSetProxyLibHandlerMallocPre;
+        umfSetProxyLibHandlerAlignedMallocPre;
+        umfSetProxyLibHandlerFreePre;
+        umfSetProxyLibHandlerMallocPost;
+        umfSetProxyLibHandlerAlignedMallocPost;
     local:
         *;
 };

src/proxy_lib/proxy_lib_linux.c

@@ -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_windows.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(); }
 

test/CMakeLists.txt

@@ -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")