Skip to content

Commit

Permalink
Add threshold to proxy lib to call system allocator
Browse files Browse the repository at this point in the history 
Add threshold to proxy lib to call system allocator
when a size is less than the given threshold.

Signed-off-by: Lukasz Dorau <[email protected]>
  • Loading branch information
@ldorau
ldorau committed Nov 10, 2024
1 parent afbab27 commit 8bfade2
Show file tree
Hide file tree
Showing 3 changed files with 108 additions and 15 deletions.
11 changes: 11 additions & 0 deletions src/libumf.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,9 @@
#include "base_alloc_global.h"
#include "memspace_internal.h"
#include "provider_tracking.h"
#include "utils_common.h"
#include "utils_log.h"

#if !defined(UMF_NO_HWLOC)
#include "topology.h"
#endif
Expand All @@ -39,6 +41,15 @@ void umfTearDown(void) {
umfMemspaceLowestLatencyDestroy();
umfDestroyTopology();
#endif

if (utils_is_running_in_proxy_lib()) {
// We cannot destroy the TRACKER nor the base allocator
// when we are running in the proxy library,
// because it could lead to use-after-free in the program's unloader
// (for example _dl_fini() on Linux).
return;
}

// make sure TRACKER is not used after being destroyed
umf_memory_tracker_handle_t t = TRACKER;
TRACKER = NULL;
Expand Down
14 changes: 8 additions & 6 deletions src/provider/provider_tracking.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -103,16 +103,19 @@ umf_memory_pool_handle_t umfMemoryTrackerGetPool(const void *ptr) {

umf_result_t umfMemoryTrackerGetAllocInfo(const void *ptr,
umf_alloc_info_t *pAllocInfo) {
assert(ptr);
assert(pAllocInfo);

if (ptr == NULL) {
return UMF_RESULT_ERROR_INVALID_ARGUMENT;
}

if (TRACKER == NULL) {
LOG_ERR("tracker is not created");
LOG_ERR("tracker does not exist");
return UMF_RESULT_ERROR_NOT_SUPPORTED;
}

if (TRACKER->map == NULL) {
LOG_ERR("tracker's map is not created");
LOG_ERR("tracker's map does not exist");
return UMF_RESULT_ERROR_NOT_SUPPORTED;
}

Expand All @@ -121,9 +124,8 @@ umf_result_t umfMemoryTrackerGetAllocInfo(const void *ptr,
int found = critnib_find(TRACKER->map, (uintptr_t)ptr, FIND_LE,
(void *)&rkey, (void **)&rvalue);
if (!found || (uintptr_t)ptr >= rkey + rvalue->size) {
LOG_WARN("pointer %p not found in the "
"tracker, TRACKER=%p",
ptr, (void *)TRACKER);
LOG_DEBUG("pointer %p not found in the tracker, TRACKER=%p", ptr,
(void *)TRACKER);
return UMF_RESULT_ERROR_INVALID_ARGUMENT;
}

Expand Down
98 changes: 89 additions & 9 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,6 +109,21 @@ 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 *********************/
/*****************************************************************************/
Expand Down Expand Up @@ -149,16 +170,40 @@ void proxy_lib_create_common(void) {
LOG_ERR("creating UMF pool manager failed");
exit(-1);
}

LOG_DEBUG("proxy pool created");

#ifndef _WIN32
*((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) {
// TODO threshold_value will be read from the environment variable
threshold_value = 128;
LOG_DEBUG("all system hooks found, threshold_value = %zu",
threshold_value);
} else {
LOG_WARN("system hooks NOT found");
}
#endif

// 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 +213,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 +295,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 +310,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 +333,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 +365,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 +401,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 (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 8bfade2

Please sign in to comment.