From bf1999f068cce6254efef342fac0553c011bb54c Mon Sep 17 00:00:00 2001 From: Rafal Rudnicki Date: Tue, 17 Dec 2024 12:02:34 +0100 Subject: [PATCH] add free handle POC --- include/umf/proxy_lib_handlers.h | 54 +++++++ src/proxy_lib/proxy_lib.c | 113 +++++++++++-- src/proxy_lib/proxy_lib.def | 6 + src/proxy_lib/proxy_lib.h | 22 --- src/proxy_lib/proxy_lib.map | 6 + src/proxy_lib/proxy_lib_linux.c | 5 +- src/proxy_lib/proxy_lib_windows.c | 5 +- test/CMakeLists.txt | 9 +- .../proxy_lib.cpp} | 0 test/proxy_lib/proxy_lib_handlers.cpp | 148 ++++++++++++++++++ .../proxy_lib_size_threshold.cpp} | 0 11 files changed, 331 insertions(+), 37 deletions(-) create mode 100644 include/umf/proxy_lib_handlers.h delete mode 100644 src/proxy_lib/proxy_lib.h rename test/{test_proxy_lib.cpp => proxy_lib/proxy_lib.cpp} (100%) create mode 100644 test/proxy_lib/proxy_lib_handlers.cpp rename test/{test_proxy_lib_size_threshold.cpp => proxy_lib/proxy_lib_size_threshold.cpp} (100%) diff --git a/include/umf/proxy_lib_handlers.h b/include/umf/proxy_lib_handlers.h new file mode 100644 index 000000000..6a5ca6841 --- /dev/null +++ b/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 + +#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 */ diff --git a/src/proxy_lib/proxy_lib.c b/src/proxy_lib/proxy_lib.c index f8bae304d..6e061510b 100644 --- a/src/proxy_lib/proxy_lib.c +++ b/src/proxy_lib/proxy_lib.c @@ -50,9 +50,9 @@ #include #include #include +#include #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; +} diff --git a/src/proxy_lib/proxy_lib.def b/src/proxy_lib/proxy_lib.def index f30b40556..6a70cec1d 100644 --- a/src/proxy_lib/proxy_lib.def +++ b/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 diff --git a/src/proxy_lib/proxy_lib.h b/src/proxy_lib/proxy_lib.h deleted file mode 100644 index aca5e8b58..000000000 --- a/src/proxy_lib/proxy_lib.h +++ /dev/null @@ -1,22 +0,0 @@ -/* - * 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_H -#define UMF_PROXY_LIB_H 1 - -#ifdef __cplusplus -extern "C" { -#endif - -void proxy_lib_create_common(void); -void proxy_lib_destroy_common(void); - -#ifdef __cplusplus -} -#endif - -#endif /* UMF_PROXY_LIB_H */ diff --git a/src/proxy_lib/proxy_lib.map b/src/proxy_lib/proxy_lib.map index 5d93d03ba..8d4bf2698 100644 --- a/src/proxy_lib/proxy_lib.map +++ b/src/proxy_lib/proxy_lib.map @@ -12,6 +12,12 @@ malloc; malloc_usable_size; realloc; + # handlers + umfSetProxyLibHandlerMallocPre; + umfSetProxyLibHandlerAlignedMallocPre; + umfSetProxyLibHandlerFreePre; + umfSetProxyLibHandlerMallocPost; + umfSetProxyLibHandlerAlignedMallocPost; local: *; }; diff --git a/src/proxy_lib/proxy_lib_linux.c b/src/proxy_lib/proxy_lib_linux.c index 50332b49f..d1960d607 100644 --- a/src/proxy_lib/proxy_lib_linux.c +++ b/src/proxy_lib/proxy_lib_linux.c @@ -7,7 +7,10 @@ * */ -#include "proxy_lib.h" +#include + +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) diff --git a/src/proxy_lib/proxy_lib_windows.c b/src/proxy_lib/proxy_lib_windows.c index 8bbff609c..c396ed484 100644 --- a/src/proxy_lib/proxy_lib_windows.c +++ b/src/proxy_lib/proxy_lib_windows.c @@ -9,7 +9,10 @@ #include -#include "proxy_lib.h" +#include + +void proxy_lib_create_common(void); +void proxy_lib_destroy_common(void); static void proxy_lib_create(void) { proxy_lib_create_common(); } diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt index c8b854ba5..7426b1fca 100644 --- a/test/CMakeLists.txt +++ b/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") diff --git a/test/test_proxy_lib.cpp b/test/proxy_lib/proxy_lib.cpp similarity index 100% rename from test/test_proxy_lib.cpp rename to test/proxy_lib/proxy_lib.cpp diff --git a/test/proxy_lib/proxy_lib_handlers.cpp b/test/proxy_lib/proxy_lib_handlers.cpp new file mode 100644 index 000000000..bff39cd15 --- /dev/null +++ b/test/proxy_lib/proxy_lib_handlers.cpp @@ -0,0 +1,148 @@ +/* + * 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 +#else +#include +#endif + +#include +#include + +#include "base.hpp" +#include "test_helpers.h" +#include "utils_assert.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 *user_data, void **ptr, + umf_memory_pool_handle_t pool) { + // NOTE: pool could be empty in case the alloc < threshold + (void)pool; + (void)ptr; + + // NOTE: we can't use google asserts here as they could call new/delete + // which leads to an infinite recursion loop + ASSERT(user_data != nullptr); + ASSERT(ptr != nullptr); + + free_cnt += 1; + + // increment user data by 1 + (*(size_t *)user_data)++; +} + +TEST_F(test, proxyLib_handlers_free) { + void *ptr = ::malloc(SIZE_64); + ASSERT_NE(ptr, nullptr); + + size_t user_data = 5; + umfSetProxyLibHandlerFreePre(handler_free_pre, (void *)&user_data); + ASSERT_EQ(free_cnt, 0); + + ::free(ptr); + ASSERT_EQ(free_cnt, 1); + ASSERT_EQ(user_data, 6); + + umfSetProxyLibHandlerFreePre(NULL, NULL); +} + +struct user_data_t { + size_t data; +}; + +static int malloc_pre_cnt = 0; +void handler_malloc_pre(void *user_data, size_t *size) { + ASSERT(user_data != nullptr); + ASSERT(size != nullptr); + + // do larger alloc - first few bytes would be reserved for user data + // (incremented here) + *size += sizeof(user_data_t); + ((user_data_t *)user_data)->data += 1; + + malloc_pre_cnt += 1; +} + +static int malloc_post_cnt = 0; +void handler_malloc_post(void *user_data, void **ptr, + umf_memory_pool_handle_t pool) { + (void)pool; + + ASSERT(user_data != nullptr); + ASSERT(ptr != nullptr); + + // fill first bytes would be filled with user_data + memcpy(*ptr, user_data, sizeof(user_data_t)); + + // shift ptr by 8 bytes so the app would not see user data + *(uint8_t *)ptr += sizeof(user_data_t); + + malloc_post_cnt += 1; +} + +void handler_free_pre2(void *user_data, void **ptr, + umf_memory_pool_handle_t pool) { + // NOTE: pool could be empty in case the alloc < threshold + (void)pool; + (void)user_data; + + ASSERT(user_data != nullptr); + ASSERT(ptr != nullptr); + + // in malloc we changed the size of ptr - we have to update the ptr here to + // free the whole allocation + *(uint8_t *)ptr -= sizeof(user_data_t); +} + +TEST_F(test, proxyLib_handlers_malloc) { + + user_data_t user_data = {5}; + umfSetProxyLibHandlerMallocPre(handler_malloc_pre, &user_data); + umfSetProxyLibHandlerMallocPost(handler_malloc_post, &user_data); + + void *ptr = ::malloc(SIZE_64); + ASSERT_NE(ptr, nullptr); + + if (ptr == NULL) { + // WA for windows checks + return; + } + + // set first and last byte to 1 + *(uint8_t *)ptr = 1; + *((uint8_t *)ptr + SIZE_64 - 1) = 1; + + // check if user data is present in the first few bytes of allocation + uint8_t *user_ptr = (uint8_t *)ptr - sizeof(user_data_t); + ASSERT_EQ(((user_data_t *)user_ptr)->data, 6); + ASSERT_EQ(user_data.data, 6); + + // check if the allaction data is correct (first and last byte) + ASSERT_EQ(*(uint8_t *)ptr, 1); + ASSERT_EQ(*((uint8_t *)ptr + SIZE_64 - 1), 1); + + ASSERT_EQ(malloc_pre_cnt, 1); + ASSERT_EQ(malloc_post_cnt, 1); + + umfSetProxyLibHandlerFreePre(handler_free_pre2, (void *)&user_data); + ::free(ptr); + + // IMPORTANT! + // We have to strictly control which allocation will be modified by our + // handlers. Cleanup them here so no further test/system allocations would + // be affected + umfSetProxyLibHandlerMallocPre(NULL, NULL); + umfSetProxyLibHandlerMallocPost(NULL, NULL); + umfSetProxyLibHandlerFreePre(NULL, NULL); +} diff --git a/test/test_proxy_lib_size_threshold.cpp b/test/proxy_lib/proxy_lib_size_threshold.cpp similarity index 100% rename from test/test_proxy_lib_size_threshold.cpp rename to test/proxy_lib/proxy_lib_size_threshold.cpp