From 35d1da633c4e0a9439b06b01d582875f9f23ac6b Mon Sep 17 00:00:00 2001 From: Lukasz Dorau Date: Wed, 13 Nov 2024 16:28:08 +0100 Subject: [PATCH] Add size threshold to proxy lib to call system allocator Add a size threshold to proxy lib to call system allocator when the size is less than the given threshold. Signed-off-by: Lukasz Dorau --- .github/workflows/reusable_proxy_lib.yml | 10 ++ src/proxy_lib/proxy_lib.c | 193 +++++++++++++++++++++-- 2 files changed, 186 insertions(+), 17 deletions(-) diff --git a/.github/workflows/reusable_proxy_lib.yml b/.github/workflows/reusable_proxy_lib.yml index 56211b97d2..ea212633e0 100644 --- a/.github/workflows/reusable_proxy_lib.yml +++ b/.github/workflows/reusable_proxy_lib.yml @@ -77,6 +77,16 @@ jobs: working-directory: ${{env.BUILD_DIR}} run: UMF_PROXY="page.disposition=shared-shm" LD_PRELOAD=./lib/libumf_proxy.so /usr/bin/date + # TODO enable the provider_file_memory_ipc test when the IPC tests with the proxy library are fixed + # see the issue: https://github.com/oneapi-src/unified-memory-framework/issues/864 + - name: Run "ctest --output-on-failure" with proxy library and size.threshold=128 + working-directory: ${{env.BUILD_DIR}} + run: > + UMF_PROXY="size.threshold=128" + UMF_LOG="level:debug;flush:debug;output:stderr;pid:yes" + LD_PRELOAD=./lib/libumf_proxy.so + ctest --output-on-failure -E provider_file_memory_ipc + - name: Check coverage if: ${{ matrix.build_type == 'Debug' }} working-directory: ${{env.BUILD_DIR}} diff --git a/src/proxy_lib/proxy_lib.c b/src/proxy_lib/proxy_lib.c index 5a5912b134..d272381f0c 100644 --- a/src/proxy_lib/proxy_lib.c +++ b/src/proxy_lib/proxy_lib.c @@ -27,6 +27,12 @@ * - _aligned_offset_recalloc() */ +#ifndef _WIN32 +#define _GNU_SOURCE // for RTLD_NEXT +#include +#undef _GNU_SOURCE +#endif /* _WIN32 */ + #if (defined PROXY_LIB_USES_JEMALLOC_POOL) #include #define umfPoolManagerOps umfJemallocPoolOps @@ -48,6 +54,7 @@ #include "base_alloc_linear.h" #include "proxy_lib.h" #include "utils_common.h" +#include "utils_load_library.h" #include "utils_log.h" #ifdef _WIN32 /* Windows ***************************************/ @@ -95,6 +102,24 @@ void utils_init_once(UTIL_ONCE_FLAG *flag, void (*onceCb)(void)); * by an application. */ +typedef void *(*system_aligned_alloc_t)(size_t alignment, size_t size); +typedef void *(*system_calloc_t)(size_t nmemb, size_t size); +typedef void (*system_free_t)(void *ptr); +typedef void *(*system_malloc_t)(size_t size); +typedef size_t (*system_malloc_usable_size_t)(void *ptr); +typedef void *(*system_realloc_t)(void *ptr, size_t size); + +// pointers to the default system allocator's API +static void *System_library_handle; +static system_aligned_alloc_t System_aligned_alloc; +static system_calloc_t System_calloc; +static system_free_t System_free; +static system_malloc_t System_malloc; +static system_malloc_usable_size_t System_malloc_usable_size; +static system_realloc_t System_realloc; + +static size_t Size_threshold_value = 0; + static UTIL_ONCE_FLAG Base_alloc_leak_initialized = UTIL_ONCE_FLAG_INIT; static umf_ba_linear_pool_t *Base_alloc_leak = NULL; static umf_memory_provider_handle_t OS_memory_provider = NULL; @@ -107,18 +132,106 @@ static __TLS int was_called_from_umfPool = 0; /*** The constructor and destructor of the proxy library *********************/ /*****************************************************************************/ +// atoi() is defined in stdlib.h, but we cannot include it on Windows +static size_t custom_atoi(const char *str) { + int len = strlen(str); + int i; + + size_t result = 0; + + for (i = len - 1; i >= 0; i--) { + if (str[i] < '0' || str[i] > '9') { + LOG_ERR("proxy_lib_create_common: size threshold is not a valid " + "number: %s", + str); + return 0; + } + + result = 10 * result + (size_t)(str[i] - '0'); + } + + return result; +} + +static size_t get_size_threshold(void) { + char *str_threshold = utils_env_var_get_str("UMF_PROXY", "size.threshold="); + if (!str_threshold) { + return 0; + } + + // move to the beginning of the number + str_threshold += strlen("size.threshold="); + // find ';' at the end + char *end = strstr(str_threshold, ";"); + if (end) { + // replace ';' with '\0' to mark end of the string + *end = '\0'; + } + + size_t int_threshold = (size_t)custom_atoi(str_threshold); + LOG_DEBUG("Size_threshold_value = (char *) %s, (int) %zu", str_threshold, + int_threshold); + + return int_threshold; +} + +static int get_system_allocator_symbols(void) { +#ifdef _WIN32 + System_library_handle = utils_open_library("msvcrt.dll", 0); +#else + System_library_handle = RTLD_NEXT; +#endif /* _WIN32 */ + + *((void **)(&System_aligned_alloc)) = + utils_get_symbol_addr(System_library_handle, "aligned_alloc", NULL); + *((void **)(&System_calloc)) = + utils_get_symbol_addr(System_library_handle, "calloc", NULL); + *((void **)(&System_free)) = + utils_get_symbol_addr(System_library_handle, "free", NULL); + *((void **)(&System_malloc)) = + utils_get_symbol_addr(System_library_handle, "malloc", NULL); + *((void **)(&System_malloc_usable_size)) = utils_get_symbol_addr( + System_library_handle, "malloc_usable_size", NULL); + *((void **)(&System_realloc)) = + utils_get_symbol_addr(System_library_handle, "realloc", NULL); + + if (System_aligned_alloc && System_calloc && System_free && System_malloc && + System_malloc_usable_size && System_realloc) { + return 0; + } + + *((void **)(&System_aligned_alloc)) = NULL; + *((void **)(&System_calloc)) = NULL; + *((void **)(&System_free)) = NULL; + *((void **)(&System_malloc)) = NULL; + *((void **)(&System_malloc_usable_size)) = NULL; + *((void **)(&System_realloc)) = NULL; + + return -1; +} + void proxy_lib_create_common(void) { utils_log_init(); umf_os_memory_provider_params_t os_params = umfOsMemoryProviderParamsDefault(); umf_result_t umf_result; -#ifndef _WIN32 - char shm_name[NAME_MAX]; + size_t _threshold = get_size_threshold(); + if (_threshold > 0) { + if (get_system_allocator_symbols()) { + LOG_ERR("initialization of the system allocator failed!"); + exit(-1); + } + + Size_threshold_value = _threshold; + LOG_INFO("system allocator initialized, size threshold value = %zu", + Size_threshold_value); + } +#ifndef _WIN32 if (utils_env_var_has_str("UMF_PROXY", "page.disposition=shared-fd")) { - LOG_DEBUG("proxy_lib: using the MAP_SHARED visibility mode with the " - "file descriptor duplication"); + LOG_INFO("proxy_lib: using the MAP_SHARED visibility mode with the " + "file descriptor duplication"); os_params.visibility = UMF_MEM_MAP_SHARED; os_params.shm_name = NULL; @@ -126,15 +239,16 @@ void proxy_lib_create_common(void) { "page.disposition=shared-shm")) { os_params.visibility = UMF_MEM_MAP_SHARED; + char shm_name[NAME_MAX]; memset(shm_name, 0, NAME_MAX); sprintf(shm_name, "umf_proxy_lib_shm_pid_%i", utils_getpid()); os_params.shm_name = shm_name; - LOG_DEBUG("proxy_lib: using the MAP_SHARED visibility mode with the " - "named shared memory: %s", - os_params.shm_name); + LOG_INFO("proxy_lib: using the MAP_SHARED visibility mode with the " + "named shared memory: %s", + os_params.shm_name); } -#endif +#endif /* _WIN32 */ umf_result = umfMemoryProviderCreate(umfOsMemoryProviderOps(), &os_params, &OS_memory_provider); @@ -149,8 +263,10 @@ void proxy_lib_create_common(void) { LOG_ERR("creating UMF pool manager failed"); exit(-1); } + // The UMF pool has just been created (Proxy_pool != NULL). Stop using // the linear allocator and start using the UMF pool allocator from now on. + LOG_DEBUG("proxy library initialized"); } void proxy_lib_destroy_common(void) { @@ -158,7 +274,7 @@ void proxy_lib_destroy_common(void) { // We cannot destroy 'Base_alloc_leak' nor 'Proxy_pool' nor 'OS_memory_provider', // because it could lead to use-after-free in the program's unloader // (for example _dl_fini() on Linux). - return; + goto fini_proxy_lib_destroy_common; } umf_memory_pool_handle_t pool = Proxy_pool; @@ -168,6 +284,18 @@ void proxy_lib_destroy_common(void) { umf_memory_provider_handle_t provider = OS_memory_provider; OS_memory_provider = NULL; umfMemoryProviderDestroy(provider); + +#ifdef _WIN32 + if (System_library_handle) { + utils_close_library(System_library_handle); + System_library_handle = NULL; + } +#endif /* _WIN32 */ + + LOG_DEBUG("proxy library destroyed"); + +fini_proxy_lib_destroy_common: + LOG_DEBUG("proxy library finalized"); } /*****************************************************************************/ @@ -246,6 +374,10 @@ static inline size_t ba_leak_pool_contains_pointer(void *ptr) { /*****************************************************************************/ void *malloc(size_t size) { + if (size < Size_threshold_value) { + return System_malloc(size); + } + if (!was_called_from_umfPool && Proxy_pool) { was_called_from_umfPool = 1; void *ptr = umfPoolMalloc(Proxy_pool, size); @@ -257,6 +389,10 @@ void *malloc(size_t size) { } void *calloc(size_t nmemb, size_t size) { + if ((nmemb * size) < Size_threshold_value) { + return System_calloc(nmemb, size); + } + if (!was_called_from_umfPool && Proxy_pool) { was_called_from_umfPool = 1; void *ptr = umfPoolCalloc(Proxy_pool, nmemb, size); @@ -276,15 +412,20 @@ void free(void *ptr) { return; } - if (Proxy_pool) { + if (Proxy_pool && (umfPoolByPtr(ptr) == Proxy_pool)) { if (umfPoolFree(Proxy_pool, ptr) != UMF_RESULT_SUCCESS) { LOG_ERR("umfPoolFree() failed"); - assert(0); } return; } - assert(0); + if (Size_threshold_value) { + System_free(ptr); + return; + } + + LOG_ERR("free() failed: %p", ptr); + return; } @@ -303,18 +444,27 @@ void *realloc(void *ptr, size_t size) { return ba_leak_realloc(ptr, size, leak_pool_contains_pointer); } - if (Proxy_pool) { + if (Proxy_pool && (umfPoolByPtr(ptr) == Proxy_pool)) { was_called_from_umfPool = 1; void *new_ptr = umfPoolRealloc(Proxy_pool, ptr, size); was_called_from_umfPool = 0; return new_ptr; } - assert(0); + if (Size_threshold_value) { + return System_realloc(ptr, size); + } + + LOG_ERR("realloc() failed: %p", ptr); + return NULL; } void *aligned_alloc(size_t alignment, size_t size) { + if (size < Size_threshold_value) { + return System_aligned_alloc(alignment, size); + } + if (!was_called_from_umfPool && Proxy_pool) { was_called_from_umfPool = 1; void *ptr = umfPoolAlignedMalloc(Proxy_pool, size, alignment); @@ -330,19 +480,28 @@ size_t _msize(void *ptr) { #else size_t malloc_usable_size(void *ptr) { #endif - - // a check to verify we are running the proxy library + // a check to verify if we are running the proxy library if (ptr == (void *)0x01) { return 0xDEADBEEF; } - if (!was_called_from_umfPool && Proxy_pool) { + if (ba_leak_pool_contains_pointer(ptr)) { + return 0; // unsupported in case of the ba_leak allocator + } + + if (Proxy_pool && (umfPoolByPtr(ptr) == Proxy_pool)) { was_called_from_umfPool = 1; size_t size = umfPoolMallocUsableSize(Proxy_pool, ptr); was_called_from_umfPool = 0; return size; } + if (Size_threshold_value) { + return System_malloc_usable_size(ptr); + } + + LOG_ERR("malloc_usable_size() failed: %p", ptr); + return 0; // unsupported in this case }