Add proxy library

ldorau · ldorau · commit 253de2c6743e · 2024-02-12T21:09:24.000+01:00
Signed-off-by: Lukasz Dorau &lt;lukasz.dorau@intel.com&gt;
diff --git a/.github/workflows/codeql.yml b/.github/workflows/codeql.yml
@@ -53,6 +53,7 @@ jobs:
         -DUMF_FORMAT_CODE_STYLE=ON
         -DUMF_DEVELOPER_MODE=ON
         -DUMF_ENABLE_POOL_TRACKING=ON
+        -DUMF_BUILD_LIBUMF_POOL_JEMALLOC=ON
         -DUMF_BUILD_LIBUMF_POOL_SCALABLE=${{matrix.pool_scalable}}
 
     - name: Build
diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt
@@ -103,3 +103,4 @@ install(TARGETS umf
 )
 
 add_subdirectory(pool)
+add_subdirectory(proxy_lib)
diff --git a/src/proxy_lib/CMakeLists.txt b/src/proxy_lib/CMakeLists.txt
@@ -0,0 +1,50 @@
+# Copyright (C) 2023-2024 Intel Corporation
+# Under the Apache License v2.0 with LLVM Exceptions. See LICENSE.TXT.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+include(${UMF_CMAKE_SOURCE_DIR}/cmake/helpers.cmake)
+
+set(PROXY_LIBS umf jemalloc_pool)
+
+set(PROXY_SOURCES
+    proxy_lib.c
+)
+
+set(PROXY_SOURCES_LINUX
+    proxy_lib_linux.c
+)
+
+set(PROXY_SOURCES_WINDOWS
+    proxy_lib_windows.c
+)
+
+set(PROXY_SOURCES_MACOSX
+    proxy_lib_linux.c
+)
+
+if(LINUX)
+    set(PROXY_SOURCES ${PROXY_SOURCES} ${PROXY_SOURCES_LINUX})
+elseif(WINDOWS)
+    set(PROXY_SOURCES ${PROXY_SOURCES} ${PROXY_SOURCES_WINDOWS})
+elseif(MACOSX)
+    set(PROXY_SOURCES ${PROXY_SOURCES} ${PROXY_SOURCES_MACOSX})
+endif()
+
+add_umf_library(NAME umf_proxy
+                TYPE SHARED
+                SRCS ${PROXY_SOURCES}
+                LIBS ${PROXY_LIBS}
+                LINUX_MAP_FILE ${CMAKE_CURRENT_SOURCE_DIR}/proxy_lib.map
+                WINDOWS_DEF_FILE ${CMAKE_CURRENT_SOURCE_DIR}/proxy_lib.def)
+
+add_library(${PROJECT_NAME}::proxy ALIAS umf_proxy)
+
+target_include_directories(umf_proxy PUBLIC
+    $<BUILD_INTERFACE:${PROJECT_SOURCE_DIR}/include>
+    $<BUILD_INTERFACE:${PROJECT_SOURCE_DIR}/src/utils>
+    $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}>
+    $<INSTALL_INTERFACE:${CMAKE_INSTALL_INCLUDEDIR}>
+)
+
+install(TARGETS umf_proxy
+        EXPORT ${PROJECT_NAME}-targets)
diff --git a/src/proxy_lib/proxy_lib.c b/src/proxy_lib/proxy_lib.c
@@ -0,0 +1,168 @@
+/*
+ * Copyright (C) 2024 Intel Corporation
+ *
+ * Under the Apache License v2.0 with LLVM Exceptions. See LICENSE.TXT.
+ * SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+*/
+
+#include <assert.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+
+#include <umf/memory_pool.h>
+#include <umf/memory_provider.h>
+#include <umf/pools/pool_jemalloc.h>
+#include <umf/providers/provider_os_memory.h>
+
+#include "proxy_lib.h"
+#include "utils_common.h"
+#include "utils_concurrency.h"
+
+static UTIL_ONCE_FLAG alloc_hooks_are_initialized = UTIL_ONCE_FLAG_INIT;
+
+static umf_memory_provider_handle_t OS_memory_provider = NULL;
+static umf_memory_pool_handle_t Proxy_pool = NULL;
+
+// it protects us from recursion in umfPool*()
+static __TLS int was_called_from_umfPool = 0;
+
+malloc_func_t System_malloc_func;
+calloc_func_t System_calloc_func;
+realloc_func_t System_realloc_func;
+free_func_t System_free_func;
+aligned_alloc_func_t System_aligned_alloc_func;
+
+/******************************************************************************/
+
+void proxy_lib_create_common(void) {
+    umf_os_memory_provider_params_t os_params =
+        umfOsMemoryProviderParamsDefault();
+    enum umf_result_t umf_result;
+
+    umf_result = umfMemoryProviderCreate(umfOsMemoryProviderOps(), &os_params,
+                                         &OS_memory_provider);
+    if (umf_result != UMF_RESULT_SUCCESS) {
+        fprintf(stderr, "error: creating OS memory provider failed\n");
+        exit(-1);
+    }
+
+    umf_result = umfPoolCreate(umfJemallocPoolOps(), OS_memory_provider, NULL,
+                               0, &Proxy_pool);
+    if (umf_result != UMF_RESULT_SUCCESS) {
+        fprintf(stderr, "error: creating jemalloc pool manager failed\n");
+        exit(-1);
+    }
+}
+
+void proxy_lib_destroy_common(void) {
+    umf_memory_pool_handle_t pool = Proxy_pool;
+    Proxy_pool = NULL;
+    umfPoolDestroy(pool);
+    umfMemoryProviderDestroy(OS_memory_provider);
+}
+
+/******************************************************************************/
+
+static inline void *system_malloc(size_t size) {
+    util_init_once(&alloc_hooks_are_initialized, proxy_lib_system_alloc_init);
+    return System_malloc_func(size);
+}
+
+static inline void *system_calloc(size_t nmemb, size_t size) {
+    util_init_once(&alloc_hooks_are_initialized, proxy_lib_system_alloc_init);
+    return System_calloc_func(nmemb, size);
+}
+
+static inline void *system_realloc(void *ptr, size_t size) {
+    util_init_once(&alloc_hooks_are_initialized, proxy_lib_system_alloc_init);
+    return System_realloc_func(ptr, size);
+}
+
+static inline void system_free(void *ptr) {
+    util_init_once(&alloc_hooks_are_initialized, proxy_lib_system_alloc_init);
+    System_free_func(ptr);
+}
+
+static inline void *system_aligned_alloc(size_t alignment, size_t size) {
+    util_init_once(&alloc_hooks_are_initialized, proxy_lib_system_alloc_init);
+    if (!System_aligned_alloc_func) {
+        return NULL; // unsupported
+    }
+    return System_aligned_alloc_func(alignment, size);
+}
+
+/******************************************************************************/
+
+WIN32_dllexport void *malloc(size_t size) {
+    if (!was_called_from_umfPool && Proxy_pool) {
+        was_called_from_umfPool = 1;
+        void *ptr = umfPoolMalloc(Proxy_pool, size);
+        was_called_from_umfPool = 0;
+        return ptr;
+    }
+
+    return system_malloc(size);
+}
+
+WIN32_dllexport void *calloc(size_t nmemb, size_t size) {
+    if (!was_called_from_umfPool && Proxy_pool) {
+        was_called_from_umfPool = 1;
+        void *ptr = umfPoolCalloc(Proxy_pool, nmemb, size);
+        was_called_from_umfPool = 0;
+        return ptr;
+    }
+
+    return system_calloc(nmemb, size);
+}
+
+WIN32_dllexport void *realloc(void *ptr, size_t size) {
+    if (!was_called_from_umfPool && Proxy_pool) {
+        was_called_from_umfPool = 1;
+        void *new_ptr = umfPoolRealloc(Proxy_pool, ptr, size);
+        was_called_from_umfPool = 0;
+        return new_ptr;
+    }
+
+    return system_realloc(ptr, size);
+}
+
+WIN32_dllexport void free(void *ptr) {
+    if (ptr == NULL) {
+        return;
+    }
+
+    umf_memory_pool_handle_t pool = umfPoolByPtr(ptr);
+    if (pool == NULL) {
+        system_free(ptr);
+        return;
+    }
+
+    enum umf_result_t umf_result = umfPoolFree(pool, ptr);
+    if (umf_result != UMF_RESULT_SUCCESS) {
+        fprintf(stderr, "error: umfPoolFree() failed\n");
+        assert(0);
+    }
+}
+
+WIN32_dllexport void *aligned_alloc(size_t alignment, size_t size) {
+    if (!was_called_from_umfPool && Proxy_pool) {
+        was_called_from_umfPool = 1;
+        void *ptr = umfPoolAlignedMalloc(Proxy_pool, size, alignment);
+        was_called_from_umfPool = 0;
+        return ptr;
+    }
+
+    return system_aligned_alloc(alignment, size);
+}
+
+WIN32_dllexport size_t malloc_usable_size(void *ptr) {
+    if (!was_called_from_umfPool && Proxy_pool) {
+        was_called_from_umfPool = 1;
+        size_t size = umfPoolMallocUsableSize(Proxy_pool, ptr);
+        was_called_from_umfPool = 0;
+        return size;
+    }
+
+    return 0; // unsupported in this case
+}
diff --git a/src/proxy_lib/proxy_lib.def b/src/proxy_lib/proxy_lib.def
@@ -0,0 +1,14 @@
+;;;; Begin Copyright Notice
+; Copyright (C) 2024 Intel Corporation
+; Under the Apache License v2.0 with LLVM Exceptions. See LICENSE.TXT.
+; SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+;;;;  End Copyright Notice
+
+LIBRARY UMF_PROXY
+EXPORTS
+	aligned_alloc
+	calloc
+	free
+	malloc
+	malloc_usable_size
+	realloc
diff --git a/src/proxy_lib/proxy_lib.h b/src/proxy_lib/proxy_lib.h
@@ -0,0 +1,43 @@
+/*
+ * Copyright (C) 2024 Intel Corporation
+ *
+ * Under the Apache License v2.0 with LLVM Exceptions. See LICENSE.TXT.
+ * SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+*/
+
+#ifndef UMF_PROXY_LIB_H
+#define UMF_PROXY_LIB_H 1
+
+#include <stddef.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+#ifdef _WIN32
+#define WIN32_dllexport __declspec(dllexport)
+#else
+#define WIN32_dllexport
+#endif
+
+typedef void *(*malloc_func_t)(size_t size);
+typedef void *(*calloc_func_t)(size_t nmemb, size_t size);
+typedef void *(*realloc_func_t)(void *ptr, size_t size);
+typedef void (*free_func_t)(void *ptr);
+typedef void *(*aligned_alloc_func_t)(size_t alignment, size_t size);
+
+extern malloc_func_t System_malloc_func;
+extern calloc_func_t System_calloc_func;
+extern realloc_func_t System_realloc_func;
+extern free_func_t System_free_func;
+extern aligned_alloc_func_t System_aligned_alloc_func;
+
+void proxy_lib_create_common(void);
+void proxy_lib_destroy_common(void);
+void proxy_lib_system_alloc_init(void);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif /* UMF_PROXY_LIB_H */
diff --git a/src/proxy_lib/proxy_lib.map b/src/proxy_lib/proxy_lib.map
@@ -0,0 +1,17 @@
+# Copyright (C) 2024 Intel Corporation
+# Under the Apache License v2.0 with LLVM Exceptions. See LICENSE.TXT.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+# linker VERSION script
+
+{
+    global:
+        aligned_alloc;
+        calloc;
+        free;
+        malloc;
+        malloc_usable_size;
+        realloc;
+    local:
+        *;
+};
diff --git a/src/proxy_lib/proxy_lib_linux.c b/src/proxy_lib/proxy_lib_linux.c
@@ -0,0 +1,41 @@
+/*
+ *
+ * Copyright (C) 2024 Intel Corporation
+ *
+ * Under the Apache License v2.0 with LLVM Exceptions. See LICENSE.TXT.
+ * SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+ *
+ */
+
+#define _GNU_SOURCE
+#include <dlfcn.h>
+
+#include <stdio.h>
+#include <stdlib.h>
+
+#include "proxy_lib.h"
+
+void __attribute__((constructor(101))) proxy_lib_create(void) {
+    proxy_lib_create_common();
+}
+
+void __attribute__((destructor(101))) proxy_lib_destroy(void) {
+    proxy_lib_destroy_common();
+}
+
+void proxy_lib_system_alloc_init(void) {
+    // ISO C forbids conversion of object pointer to function pointer type [-Werror=pedantic]
+    *((void **)(&System_malloc_func)) = dlsym(RTLD_NEXT, "malloc");
+    *((void **)(&System_calloc_func)) = dlsym(RTLD_NEXT, "calloc");
+    *((void **)(&System_realloc_func)) = dlsym(RTLD_NEXT, "realloc");
+    *((void **)(&System_free_func)) = dlsym(RTLD_NEXT, "free");
+    *((void **)(&System_aligned_alloc_func)) =
+        dlsym(RTLD_NEXT, "aligned_alloc");
+
+    // System_aligned_alloc_func can be NULL
+    if (!System_malloc_func || !System_calloc_func || !System_realloc_func ||
+        !System_free_func) {
+        fprintf(stderr, "error: cannot find libc alloc functions\n");
+        exit(-1);
+    }
+}
diff --git a/src/proxy_lib/proxy_lib_windows.c b/src/proxy_lib/proxy_lib_windows.c
@@ -0,0 +1,47 @@
+/*
+ *
+ * Copyright (C) 2024 Intel Corporation
+ *
+ * Under the Apache License v2.0 with LLVM Exceptions. See LICENSE.TXT.
+ * SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+ *
+ */
+
+#include <Windows.h>
+
+#include <stdio.h>
+
+#include "proxy_lib.h"
+
+static void proxy_lib_create(void) { proxy_lib_create_common(); }
+
+static void proxy_lib_destroy(void) { proxy_lib_destroy_common(); }
+
+BOOL APIENTRY DllMain(HINSTANCE hinstDLL, DWORD fdwReason, LPVOID lpvReserved) {
+    if (fdwReason == DLL_PROCESS_DETACH) {
+        proxy_lib_destroy();
+    } else if (fdwReason == DLL_PROCESS_ATTACH) {
+        proxy_lib_create();
+    }
+    return TRUE;
+}
+
+void proxy_lib_system_alloc_init(void) {
+    /* there is a memleak here, we should unload this library in the destructor but there is no such thing */
+    HMODULE mod = LoadLibrary(TEXT("msvcrt.dll"));
+
+    // ISO C forbids conversion of object pointer to function pointer type [-Werror=pedantic]
+    *((void **)(&System_malloc_func)) = GetProcAddress(mod, "malloc");
+    *((void **)(&System_calloc_func)) = GetProcAddress(mod, "calloc");
+    *((void **)(&System_realloc_func)) = GetProcAddress(mod, "realloc");
+    *((void **)(&System_free_func)) = GetProcAddress(mod, "free");
+    *((void **)(&System_aligned_alloc_func)) =
+        GetProcAddress(mod, "aligned_alloc");
+
+    // System_aligned_alloc_func can be NULL
+    if (!System_malloc_func || !System_calloc_func || !System_realloc_func ||
+        !System_free_func) {
+        fprintf(stderr, "error: cannot find libc alloc functions\n");
+        exit(-1);
+    }
+}

Original file line number	Diff line number	Diff line change
`@@ -103,3 +103,4 @@ install(TARGETS umf`
`103`	`103`	`)`
`104`	`104`
`105`	`105`	`add_subdirectory(pool)`
	`106`	`+add_subdirectory(proxy_lib)`