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 12, 2024
1 parent f24b0f7 commit 0d3f0c7
Show file tree
Hide file tree
Showing 2 changed files with 144 additions and 16 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
150 changes: 134 additions & 16 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
#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 Down Expand Up @@ -103,38 +109,112 @@ 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 =
umfOsMemoryProviderParamsDefault();
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;

} 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 +229,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 +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");
}

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

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

Expand Down

0 comments on commit 0d3f0c7

Please sign in to comment.