Skip to content

Commit

Permalink
Add size threshold to proxy lib to call system allocator
Browse files Browse the repository at this point in the history 
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 <[email protected]>
  • Loading branch information
@ldorau
ldorau committed Nov 13, 2024
1 parent 275af24 commit 35d1da6
Show file tree
Hide file tree
Showing 2 changed files with 186 additions and 17 deletions.
10 changes: 10 additions & 0 deletions .github/workflows/reusable_proxy_lib.yml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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}}
Expand Down
193 changes: 176 additions & 17 deletions src/proxy_lib/proxy_lib.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,12 @@
* - _aligned_offset_recalloc()
*/

#ifndef _WIN32
#define _GNU_SOURCE // for RTLD_NEXT
#include <dlfcn.h>
#undef _GNU_SOURCE
#endif /* _WIN32 */

#if (defined PROXY_LIB_USES_JEMALLOC_POOL)
#include <umf/pools/pool_jemalloc.h>
#define umfPoolManagerOps umfJemallocPoolOps
Expand All @@ -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 ***************************************/
Expand Down Expand Up @@ -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;
Expand All @@ -107,34 +132,123 @@ 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;

} else if (utils_env_var_has_str("UMF_PROXY",
"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);
Expand All @@ -149,16 +263,18 @@ 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) {
if (utils_is_running_in_proxy_lib()) {
// 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;
Expand All @@ -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");
}

/*****************************************************************************/
Expand Down Expand Up @@ -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);
Expand All @@ -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);
Expand All @@ -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;
}

Expand All @@ -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);
Expand All @@ -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
}

Expand Down

0 comments on commit 35d1da6

Please sign in to comment.