From 7c781ef542004352acff2f0e61bbff398e6f11ee Mon Sep 17 00:00:00 2001 From: Lukasz Dorau Date: Tue, 12 Nov 2024 12:34:42 +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 | 150 ++++++++++++++++++++--- 2 files changed, 144 insertions(+), 16 deletions(-) diff --git a/.github/workflows/reusable_proxy_lib.yml b/.github/workflows/reusable_proxy_lib.yml index 56211b97d..ea212633e 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 2730b9f17..7eb0b282e 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 +#include +#undef _GNU_SOURCE +#endif /* _WIN32 */ + #if (defined PROXY_LIB_USES_JEMALLOC_POOL) #include #define umfPoolManagerOps umfJemallocPoolOps @@ -103,10 +109,73 @@ static umf_memory_pool_handle_t Proxy_pool = NULL; // it protects us from recursion in umfPool*() static __TLS int was_called_from_umfPool = 0; +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); + +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 threshold_value = 0; + /*****************************************************************************/ /*** The constructor and destructor of the proxy library *********************/ /*****************************************************************************/ +#ifndef _WIN32 +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)atoi(str_threshold); + LOG_DEBUG("threshold_value = (char *) %s, (int) %zu", str_threshold, + int_threshold); + + return int_threshold; +} + +static int dlsym_system_allocator(void) { + *((void **)(&system_aligned_alloc)) = dlsym(RTLD_NEXT, "aligned_alloc"); + *((void **)(&system_calloc)) = dlsym(RTLD_NEXT, "calloc"); + *((void **)(&system_free)) = dlsym(RTLD_NEXT, "free"); + *((void **)(&system_malloc)) = dlsym(RTLD_NEXT, "malloc"); + *((void **)(&system_malloc_usable_size)) = + dlsym(RTLD_NEXT, "malloc_usable_size"); + *((void **)(&system_realloc)) = dlsym(RTLD_NEXT, "realloc"); + + 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; +} +#endif /* _WIN32 */ + void proxy_lib_create_common(void) { utils_log_init(); umf_os_memory_provider_params_t os_params = @@ -114,11 +183,21 @@ void proxy_lib_create_common(void) { umf_result_t umf_result; #ifndef _WIN32 - char shm_name[NAME_MAX]; + size_t _threshold = get_size_threshold(); + if (_threshold > 0) { + if (dlsym_system_allocator()) { + LOG_ERR("initialization of the system allocator failed!"); + exit(-1); + } + + threshold_value = _threshold; + LOG_INFO("system allocator initialized, size threshold value = %zu", + threshold_value); + } 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 +205,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 +229,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 +240,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 +250,10 @@ void proxy_lib_destroy_common(void) { umf_memory_provider_handle_t provider = OS_memory_provider; OS_memory_provider = NULL; umfMemoryProviderDestroy(provider); + LOG_DEBUG("proxy library destroyed"); + +fini_proxy_lib_destroy_common: + LOG_DEBUG("proxy library finalized"); } /*****************************************************************************/ @@ -246,6 +332,10 @@ static inline size_t ba_leak_pool_contains_pointer(void *ptr) { /*****************************************************************************/ void *malloc(size_t size) { + if (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 +347,10 @@ void *malloc(size_t size) { } void *calloc(size_t nmemb, size_t size) { + if ((nmemb * 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 +370,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 (threshold_value) { + system_free(ptr); + return; + } + + LOG_ERR("free() failed: %p", ptr); + return; } @@ -303,18 +402,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 (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 < 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 +438,29 @@ 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 (!was_called_from_umfPool && 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 (threshold_value) { + return system_malloc_usable_size(ptr); + } + + LOG_ERR("malloc_usable_size() failed: %p", ptr); + return 0; // unsupported in this case }