Add size threshold to proxy lib to call system allocator (Linux only)

Add a size threshold to proxy lib to call system allocator when the size is less than the given threshold (Linux only yet). Signed-off-by: Lukasz Dorau <[email protected]>
ldorau · Nov 18, 2024 · 0d69471 · 0d69471
1 parent 2b5756a
commit 0d69471
Show file tree

Hide file tree

Showing 6 changed files with 197 additions and 16 deletions.
diff --git a/README.md b/README.md
@@ -317,6 +317,13 @@ The memory used by the proxy memory allocator is mmap'ed:
    - `page.disposition=shared-shm` - IPC uses the named shared memory. An SHM name is generated using the `umf_proxy_lib_shm_pid_$PID` pattern, where `$PID` is the PID of the process. It creates the `/dev/shm/umf_proxy_lib_shm_pid_$PID` file.
    - `page.disposition=shared-fd` - IPC uses the file descriptor duplication. It requires using `pidfd_getfd(2)` to obtain a duplicate of another process's file descriptor. Permission to duplicate another process's file descriptor is governed by a ptrace access mode `PTRACE_MODE_ATTACH_REALCREDS` check (see `ptrace(2)`) that can be changed using the `/proc/sys/kernel/yama/ptrace_scope` interface. `pidfd_getfd(2)` is supported since Linux 5.6.
 
+**Size threshold**
+
+The **size threshold** feature (Linux only) causes that all allocations of size less than the given threshold value go to the default system allocator instead of the proxy library.
+It can be enabled by adding the `size.threshold=<value>` string to the `UMF_PROXY` environment variable (with `';'` as a separator), for example: `UMF_PROXY="page.disposition=shared-shm;size.threshold=64"`.
+
+**Remark:** changing a size of allocation (using `realloc()` ) does not change the allocator (`realloc(malloc(threshold - 1), threshold + 1)` still belongs to the default system allocator and `realloc(malloc(threshold + 1), threshold - 1)` still belongs to the proxy library pool allocator).
+
 #### Windows
 
 In case of Windows it requires:

diff --git 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 <dlfcn.h>
+#undef _GNU_SOURCE
+#endif /* _WIN32 */
+
 #if (defined PROXY_LIB_USES_JEMALLOC_POOL)
 #include <umf/pools/pool_jemalloc.h>
 #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 ***************************************/
@@ -94,6 +101,24 @@ void utils_init_once(UTIL_ONCE_FLAG *flag, void (*onceCb)(void));
  * of a UMF pool to allocate memory needed by an application. It should be freed
  * by an application.
  */
+#ifndef _WIN32
+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 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;
+#endif /* _WIN32 */
 
 static UTIL_ONCE_FLAG Base_alloc_leak_initialized = UTIL_ONCE_FLAG_INIT;
 static umf_ba_linear_pool_t *Base_alloc_leak = NULL;
@@ -107,34 +132,86 @@ static __TLS int was_called_from_umfPool = 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=");
+    long threshold = utils_get_size_threshold(str_threshold);
+    if (threshold < 0) {
+        LOG_ERR("incorrect size threshold: %s", str_threshold);
+        exit(-1);
+    }
+
+    return (size_t)threshold;
+}
+
+static int get_system_allocator_symbols(void) {
+    *((void **)(&System_aligned_alloc)) =
+        utils_get_symbol_addr(RTLD_NEXT, "aligned_alloc", NULL);
+    *((void **)(&System_calloc)) =
+        utils_get_symbol_addr(RTLD_NEXT, "calloc", NULL);
+    *((void **)(&System_free)) = utils_get_symbol_addr(RTLD_NEXT, "free", NULL);
+    *((void **)(&System_malloc)) =
+        utils_get_symbol_addr(RTLD_NEXT, "malloc", NULL);
+    *((void **)(&System_malloc_usable_size)) =
+        utils_get_symbol_addr(RTLD_NEXT, "malloc_usable_size", NULL);
+    *((void **)(&System_realloc)) =
+        utils_get_symbol_addr(RTLD_NEXT, "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;
+}
+#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 (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);
+    }
 
     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);
@@ -149,16 +226,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;
@@ -168,6 +247,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 +329,12 @@ static inline size_t ba_leak_pool_contains_pointer(void *ptr) {
 /*****************************************************************************/
 
 void *malloc(size_t size) {
+#ifndef _WIN32
+    if (size < Size_threshold_value) {
+        return System_malloc(size);
+    }
+#endif /* _WIN32 */
+
     if (!was_called_from_umfPool && Proxy_pool) {
         was_called_from_umfPool = 1;
         void *ptr = umfPoolMalloc(Proxy_pool, size);
@@ -257,6 +346,12 @@ void *malloc(size_t size) {
 }
 
 void *calloc(size_t nmemb, size_t size) {
+#ifndef _WIN32
+    if ((nmemb * size) < Size_threshold_value) {
+        return System_calloc(nmemb, size);
+    }
+#endif /* _WIN32 */
+
     if (!was_called_from_umfPool && Proxy_pool) {
         was_called_from_umfPool = 1;
         void *ptr = umfPoolCalloc(Proxy_pool, nmemb, size);
@@ -276,15 +371,22 @@ 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);
+#ifndef _WIN32
+    if (Size_threshold_value) {
+        System_free(ptr);
+        return;
+    }
+#endif /* _WIN32 */
+
+    LOG_ERR("free() failed: %p", ptr);
+
     return;
 }
 
@@ -303,18 +405,31 @@ 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);
+#ifndef _WIN32
+    if (Size_threshold_value) {
+        return System_realloc(ptr, size);
+    }
+#endif /* _WIN32 */
+
+    LOG_ERR("realloc() failed: %p", ptr);
+
     return NULL;
 }
 
 void *aligned_alloc(size_t alignment, size_t size) {
+#ifndef _WIN32
+    if (size < Size_threshold_value) {
+        return System_aligned_alloc(alignment, size);
+    }
+#endif /* _WIN32 */
+
     if (!was_called_from_umfPool && Proxy_pool) {
         was_called_from_umfPool = 1;
         void *ptr = umfPoolAlignedMalloc(Proxy_pool, size, alignment);
@@ -330,19 +445,30 @@ 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;
     }
 
+#ifndef _WIN32
+    if (Size_threshold_value) {
+        return System_malloc_usable_size(ptr);
+    }
+#endif /* _WIN32 */
+
+    LOG_ERR("malloc_usable_size() failed: %p", ptr);
+
     return 0; // unsupported in this case
 }
 

diff --git a/src/utils/utils_common.h b/src/utils/utils_common.h
@@ -168,6 +168,8 @@ int utils_file_open_or_create(const char *path);
 
 int utils_fallocate(int fd, long offset, long len);
 
+long utils_get_size_threshold(char *str_threshold);
+
 #ifdef __cplusplus
 }
 #endif

diff --git a/src/utils/utils_posix_common.c b/src/utils/utils_posix_common.c
@@ -7,6 +7,7 @@
  *
  */
 
+#include <ctype.h>
 #include <errno.h>
 #include <fcntl.h>
 #include <limits.h>
@@ -266,3 +267,27 @@ int utils_file_open_or_create(const char *path) {
 
     return fd;
 }
+
+// Expected input:
+// char *str_threshold = utils_env_var_get_str("UMF_PROXY", "size.threshold=");
+long utils_get_size_threshold(char *str_threshold) {
+    if (!str_threshold) {
+        return 0;
+    }
+
+    // move to the beginning of the number
+    str_threshold += strlen("size.threshold=");
+
+    // check if the first character is a digit
+    if (!isdigit(str_threshold[0])) {
+        LOG_ERR("incorrect size threshold, expected numerical value >=0: %s",
+                str_threshold);
+        return -1;
+    }
+
+    size_t threshold = atol(str_threshold);
+    LOG_DEBUG("Size_threshold_value = (char *) %s, (size_t) %zu", str_threshold,
+              threshold);
+
+    return threshold;
+}
diff --git a/src/utils/utils_windows_common.c b/src/utils/utils_windows_common.c
@@ -215,3 +215,10 @@ int utils_fallocate(int fd, long offset, long len) {
 
     return -1;
 }
+
+// Expected input:
+// char *str_threshold = utils_env_var_get_str("UMF_PROXY", "size.threshold=");
+long utils_get_size_threshold(char *str_threshold) {
+    (void)str_threshold; // unused
+    return 0;
+}
diff --git a/test/utils/utils_linux.cpp b/test/utils/utils_linux.cpp
@@ -56,3 +56,17 @@ TEST_F(test, utils_shm_create_invalid_args) {
     ret = utils_shm_create("/abc", -1);
     EXPECT_EQ(ret, -1);
 }
+
+TEST_F(test, utils_get_size_threshold) {
+    // Expected input to utils_get_size_threshold():
+    // char *str_threshold = utils_env_var_get_str("UMF_PROXY", "size.threshold=");
+    // positive tests
+    EXPECT_EQ(utils_get_size_threshold((char *)"size.threshold=111"), 111);
+    EXPECT_EQ(utils_get_size_threshold((char *)"size.threshold=222;abcd"), 222);
+    EXPECT_EQ(utils_get_size_threshold((char *)"size.threshold=333;var=value"),
+              333);
+    // negative tests
+    EXPECT_EQ(utils_get_size_threshold(NULL), 0);
+    EXPECT_EQ(utils_get_size_threshold((char *)"size.threshold=abc"), -1);
+    EXPECT_EQ(utils_get_size_threshold((char *)"size.threshold=-111"), -1);
+}