Add proxy library for intercepting user allocation requests

UMF proxy library is a library for intercepting user allocation requests.
It intercepts following APIs:
- aligned_alloc()
- calloc()
- free()
- malloc()
- malloc_usable_size()
- realloc()

Signed-off-by: Lukasz Dorau <[email protected]>

Author: ldorau

.github/workflows/pr_push.yml

@@ -91,3 +91,6 @@ jobs:
     Benchmark:
         needs: [Build]
         uses: ./.github/workflows/benchmarks.yml
+    Proxy_lib:
+        needs: [Build]
+        uses: ./.github/workflows/proxy_lib.yml

.github/workflows/proxy_lib.yml

@@ -0,0 +1,62 @@
+name: Proxy library
+
+on: workflow_call
+
+permissions:
+  contents: read
+
+jobs:
+  ubuntu-build:
+    name: Build - Ubuntu
+
+    strategy:
+      matrix:
+        os: ['ubuntu-22.04']
+        build_type: [Release, Debug]
+        compiler: [{c: gcc, cxx: g++}]
+    runs-on: ${{matrix.os}}
+
+    steps:
+      - name: Checkout
+        uses: actions/checkout@b4ffde65f46336ab88eb53be808477a3936bae11 # v4.1.1
+
+      - name: Install apt packages
+        run: |
+          sudo apt-get update
+          sudo apt-get install -y cmake libhwloc-dev libjemalloc-dev libtbb-dev
+
+      - name: Configure build
+        run: >
+          cmake
+          -B ${{github.workspace}}/build
+          -DCMAKE_BUILD_TYPE=${{matrix.build_type}}
+          -DCMAKE_C_COMPILER=${{matrix.compiler.c}}
+          -DCMAKE_CXX_COMPILER=${{matrix.compiler.cxx}}
+          -DUMF_BUILD_SHARED_LIBRARY=ON
+          -DUMF_BUILD_BENCHMARKS=ON
+          -DUMF_BUILD_TESTS=ON
+          -DUMF_FORMAT_CODE_STYLE=OFF
+          -DUMF_DEVELOPER_MODE=OFF
+          -DUMF_BUILD_LIBUMF_POOL_JEMALLOC=ON
+          -DUMF_BUILD_LIBUMF_POOL_DISJOINT=ON
+          -DUMF_BUILD_LIBUMF_POOL_SCALABLE=ON
+          -DUMF_ENABLE_POOL_TRACKING=OFF
+
+      - name: Build UMF
+        run: cmake --build ${{github.workspace}}/build -j $(nproc)
+
+      - name: Run "ctest --output-on-failure" with proxy library
+        working-directory: ${{github.workspace}}/build
+        run: LD_PRELOAD=./lib/libumf_proxy.so ctest --output-on-failure
+
+      - name: Run "./test/umf_test-memoryPool" with proxy library
+        working-directory: ${{github.workspace}}/build
+        run: LD_PRELOAD=./lib/libumf_proxy.so ./test/umf_test-memoryPool
+
+      - name: Run "/usr/bin/ls" with proxy library
+        working-directory: ${{github.workspace}}/build
+        run: LD_PRELOAD=./lib/libumf_proxy.so /usr/bin/ls
+
+      - name: Run "/usr/bin/date" with proxy library
+        working-directory: ${{github.workspace}}/build
+        run: LD_PRELOAD=./lib/libumf_proxy.so /usr/bin/date

src/CMakeLists.txt

@@ -4,6 +4,11 @@
 
 include(${UMF_CMAKE_SOURCE_DIR}/cmake/helpers.cmake)
 
+# set UMF_PROXY_LIB_BASED_ON_POOL to one of:
+# - SCALABLE
+# - JEMALLOC
+set(UMF_PROXY_LIB_BASED_ON_POOL SCALABLE)
+
 add_subdirectory(utils)
 
 set(UMF_LIBS umf_utils)
@@ -103,3 +108,26 @@ install(TARGETS umf
 )
 
 add_subdirectory(pool)
+
+# TODO: enable proxy_lib on Windows
+if(LINUX)
+    if(UMF_PROXY_LIB_BASED_ON_POOL STREQUAL SCALABLE)
+        set(PROXY_LIB_USES_SCALABLE_POOL ON)
+        set(PROXY_LIBS umf scalable_pool)
+        if(UMF_BUILD_LIBUMF_POOL_SCALABLE)
+            add_subdirectory(proxy_lib)
+        else()
+            message(WARNING "Disabling the proxy library, because UMF_PROXY_LIB_BASED_ON_POOL==SCALABLE but UMF_BUILD_LIBUMF_POOL_SCALABLE is OFF")
+        endif()
+    elseif(UMF_PROXY_LIB_BASED_ON_POOL STREQUAL JEMALLOC)
+        set(PROXY_LIB_USES_JEMALLOC_POOL ON)
+        set(PROXY_LIBS umf jemalloc_pool)
+        if(UMF_BUILD_LIBUMF_POOL_JEMALLOC)
+            add_subdirectory(proxy_lib)
+        else()
+            message(WARNING "Disabling the proxy library, because UMF_PROXY_LIB_BASED_ON_POOL==JEMALLOC but UMF_BUILD_LIBUMF_POOL_JEMALLOC is OFF")
+        endif()
+    else()
+        message(FATAL_ERROR "Proxy library: pool manager not chosen")
+    endif()
+endif()

src/base_alloc/base_alloc.c

@@ -36,9 +36,9 @@ struct umf_ba_main_pool_meta_t {
     size_t chunk_size;         // size of all memory chunks in this pool
     os_mutex_t free_lock;      // lock of free_list
     umf_ba_chunk_t *free_list; // list of free chunks
+    size_t n_allocs;           // number of allocated chunks
 #ifndef NDEBUG
     size_t n_pools;
-    size_t n_allocs;
     size_t n_chunks;
 #endif /* NDEBUG */
 };
@@ -135,9 +135,9 @@ umf_ba_pool_t *umf_ba_create(size_t size) {
     pool->metadata.pool_size = pool_size;
     pool->metadata.chunk_size = chunk_size;
     pool->next_pool = NULL; // this is the only pool now
+    pool->metadata.n_allocs = 0;
 #ifndef NDEBUG
     pool->metadata.n_pools = 1;
-    pool->metadata.n_allocs = 0;
     pool->metadata.n_chunks = 0;
 #endif /* NDEBUG */
 
@@ -187,8 +187,8 @@ void *umf_ba_alloc(umf_ba_pool_t *pool) {
 
     umf_ba_chunk_t *chunk = pool->metadata.free_list;
     pool->metadata.free_list = pool->metadata.free_list->next;
-#ifndef NDEBUG
     pool->metadata.n_allocs++;
+#ifndef NDEBUG
     ba_debug_checks(pool);
 #endif /* NDEBUG */
     util_mutex_unlock(&pool->metadata.free_lock);
@@ -230,18 +230,30 @@ void umf_ba_free(umf_ba_pool_t *pool, void *ptr) {
     assert(pool_contains_pointer(pool, ptr));
     chunk->next = pool->metadata.free_list;
     pool->metadata.free_list = chunk;
-#ifndef NDEBUG
     pool->metadata.n_allocs--;
+#ifndef NDEBUG
     ba_debug_checks(pool);
 #endif /* NDEBUG */
     util_mutex_unlock(&pool->metadata.free_lock);
 }
 
 void umf_ba_destroy(umf_ba_pool_t *pool) {
+    // Do not destroy if we are running in the proxy library,
+    // because it may need those resources till
+    // the very end of exiting the application.
+    if (pool->metadata.n_allocs && is_running_in_proxy_lib()) {
+        return;
+    }
+
 #ifndef NDEBUG
-    assert(pool->metadata.n_allocs == 0);
     ba_debug_checks(pool);
+    if (pool->metadata.n_allocs) {
+        fprintf(stderr, "umf_ba_destroy(): pool->metadata.n_allocs = %zu\n",
+                pool->metadata.n_allocs);
+        assert(pool->metadata.n_allocs == 0);
+    }
 #endif /* NDEBUG */
+
     size_t size = pool->metadata.pool_size;
     umf_ba_next_pool_t *current_pool;
     umf_ba_next_pool_t *next_pool = pool->next_pool;

src/base_alloc/base_alloc_linear.c

@@ -161,9 +161,17 @@ void *umf_ba_linear_alloc(umf_ba_linear_pool_t *pool, size_t size) {
 }
 
 void umf_ba_linear_destroy(umf_ba_linear_pool_t *pool) {
+    // Do not destroy if we are running in the proxy library,
+    // because it may need those resources till
+    // the very end of exiting the application.
+    if (is_running_in_proxy_lib()) {
+        return;
+    }
+
 #ifndef NDEBUG
     ba_debug_checks(pool);
 #endif /* NDEBUG */
+
     umf_ba_next_linear_pool_t *current_pool;
     umf_ba_next_linear_pool_t *next_pool = pool->next_pool;
     while (next_pool) {

src/provider/provider_tracking.c

@@ -10,6 +10,7 @@
 #include "provider_tracking.h"
 #include "base_alloc_global.h"
 #include "critnib.h"
+#include "utils_common.h"
 #include "utils_concurrency.h"
 
 #include <umf/memory_pool.h>
@@ -357,18 +358,23 @@ static void check_if_tracker_is_empty(umf_memory_tracker_handle_t hTracker,
     }
 
     if (n_items) {
-        if (pool) {
-            fprintf(stderr,
-                    "ASSERT: tracking provider of pool %p is not empty! (%zu "
-                    "items left)\n",
-                    (void *)pool, n_items);
-        } else {
-            fprintf(
-                stderr,
-                "ASSERT: tracking provider is not empty! (%zu items left)\n",
-                n_items);
+        // Do not assert if we are running in the proxy library,
+        // because it may need those resources till
+        // the very end of exiting the application.
+        if (!is_running_in_proxy_lib()) {
+            if (pool) {
+                fprintf(stderr,
+                        "ASSERT: tracking provider of pool %p is not empty! "
+                        "(%zu items left)\n",
+                        (void *)pool, n_items);
+            } else {
+                fprintf(stderr,
+                        "ASSERT: tracking provider is not empty! (%zu items "
+                        "left)\n",
+                        n_items);
+            }
+            assert(n_items == 0);
         }
-        assert(n_items == 0);
     }
 }
 #endif /* NDEBUG */
@@ -505,12 +511,24 @@ void umfMemoryTrackerDestroy(umf_memory_tracker_handle_t handle) {
         return;
     }
 
+    // Do not destroy if we are running in the proxy library,
+    // because it may need those resources till
+    // the very end of exiting the application.
+    if (is_running_in_proxy_lib()) {
+        return;
+    }
+
 #ifndef NDEBUG
     check_if_tracker_is_empty(handle, NULL);
 #endif /* NDEBUG */
 
+    // We have to zero all inner pointers,
+    // because the tracker handle can be copied
+    // and used in many places.
     critnib_delete(handle->map);
+    handle->map = NULL;
     util_mutex_destroy_not_free(&handle->splitMergeMutex);
     umf_ba_destroy(handle->tracker_allocator);
+    handle->tracker_allocator = NULL;
     umf_ba_global_free(handle, sizeof(struct umf_memory_tracker_t));
 }

src/proxy_lib/CMakeLists.txt

@@ -0,0 +1,54 @@
+# 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_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)
+
+if(PROXY_LIB_USES_SCALABLE_POOL)
+    target_compile_definitions(umf_proxy PRIVATE PROXY_LIB_USES_SCALABLE_POOL=1)
+elseif(PROXY_LIB_USES_JEMALLOC_POOL)
+    target_compile_definitions(umf_proxy PRIVATE PROXY_LIB_USES_JEMALLOC_POOL=1)
+endif()
+
+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)