Use global base_alloc in proxy lib

src/CMakeLists.txt

@@ -15,7 +15,6 @@ set(UMF_LIBS umf_utils)
 
 set(BA_SOURCES
     ${CMAKE_CURRENT_SOURCE_DIR}/base_alloc/base_alloc.c
-    ${CMAKE_CURRENT_SOURCE_DIR}/base_alloc/base_alloc_linear.c
     ${CMAKE_CURRENT_SOURCE_DIR}/base_alloc/base_alloc_global.c)
 
 if (LINUX)

src/base_alloc/base_alloc.c

@@ -238,13 +238,6 @@ void umf_ba_free(umf_ba_pool_t *pool, void *ptr) {
 }
 
 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
     ba_debug_checks(pool);
     if (pool->metadata.n_allocs) {

src/base_alloc/base_alloc_global.c

@@ -14,16 +14,17 @@
 #include "base_alloc.h"
 #include "base_alloc_global.h"
 #include "base_alloc_internal.h"
+#include "utils_common.h"
 #include "utils_concurrency.h"
 #include "utils_math.h"
 
 // global base allocator used by all providers and pools
 static UTIL_ONCE_FLAG ba_is_initialized = UTIL_ONCE_FLAG_INIT;
 
-// allocation classes need to be powers of 2
+// allocation classes need to be consecutive powers of 2
 #define ALLOCATION_CLASSES                                                     \
-    { 16, 32, 64, 128 }
-#define NUM_ALLOCATION_CLASSES 4
+    { 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192 }
+#define NUM_ALLOCATION_CLASSES 10
 
 struct base_alloc_t {
     size_t ac_sizes[NUM_ALLOCATION_CLASSES];
@@ -82,41 +83,115 @@ static int size_to_idx(size_t size) {
     return index;
 }
 
-void *umf_ba_global_alloc(size_t size) {
+static void *transform_ptr(void *ptr, size_t size, size_t alignment) {
+    assert(ptr);
+
+    void *user_ptr;
+    if (alignment <= sizeof(size_t)) {
+        user_ptr = (void *)((uintptr_t)ptr + sizeof(size_t));
+    } else {
+        user_ptr = (void *)ALIGN_UP((uintptr_t)ptr + sizeof(size_t), alignment);
+    }
+
+    size_t ptr_offset_from_original = (uintptr_t)user_ptr - (uintptr_t)ptr;
+
+    size_t *metadata_loc = (size_t *)((char *)user_ptr - sizeof(size_t));
+    *metadata_loc = size | (ptr_offset_from_original << 32);
+
+    return user_ptr;
+}
+
+static void *get_original_alloc(void *user_ptr, size_t *total_size,
+                                size_t *usable_size) {
+    assert(user_ptr);
+
+    size_t *metadata_loc = (size_t *)((char *)user_ptr - sizeof(size_t));
+
+    size_t stored_size = *metadata_loc & ((1ULL << 32) - 1);
+    size_t ptr_offset_from_original = *metadata_loc >> 32;
+
+    void *original_ptr =
+        (void *)((uintptr_t)user_ptr - ptr_offset_from_original);
+
+    if (total_size) {
+        *total_size = stored_size;
+    }
+
+    if (usable_size) {
+        *usable_size = stored_size - ptr_offset_from_original;
+    }
+
+    return original_ptr;
+}
+
+void *umf_ba_global_aligned_alloc(size_t size, size_t alignment) {
     util_init_once(&ba_is_initialized, umf_ba_create_global);
 
+    // for metadata
+    size += sizeof(size_t);
+
+    if (alignment > sizeof(size_t)) {
+        size += alignment;
+    }
+
     if (size > BASE_ALLOC.ac_sizes[NUM_ALLOCATION_CLASSES - 1]) {
 #ifndef NDEBUG
         fprintf(stderr,
-                "base_alloc: allocation size larger than the biggest "
-                "allocation class. Falling back to OS memory allocation.\n");
+                "base_alloc: allocation size (%zu) larger than the biggest "
+                "allocation class. Falling back to OS memory allocation.\n",
+                size);
 #endif
-        return ba_os_alloc(size);
+        return transform_ptr(ba_os_alloc(size), size, alignment);
     }
 
     int ac_index = size_to_idx(size);
     if (!BASE_ALLOC.ac[ac_index]) {
         // if creating ac failed, fall back to os allocation
         fprintf(stderr, "base_alloc: allocation class not created. Falling "
                         "back to OS memory allocation.\n");
-        return ba_os_alloc(size);
+        return transform_ptr(ba_os_alloc(size), size, alignment);
     }
 
-    return umf_ba_alloc(BASE_ALLOC.ac[ac_index]);
+    return transform_ptr(umf_ba_alloc(BASE_ALLOC.ac[ac_index]), size,
+                         alignment);
 }
 
-void umf_ba_global_free(void *ptr, size_t size) {
-    if (size > BASE_ALLOC.ac_sizes[NUM_ALLOCATION_CLASSES - 1]) {
-        ba_os_free(ptr, size);
+void *umf_ba_global_alloc(size_t size) {
+    return umf_ba_global_aligned_alloc(size, sizeof(size_t));
+}
+
+void umf_ba_global_free(void *ptr) {
+    if (!ptr) {
         return;
     }
 
-    int ac_index = size_to_idx(size);
+    size_t total_size;
+    ptr = get_original_alloc(ptr, &total_size, NULL);
+
+    if (total_size > BASE_ALLOC.ac_sizes[NUM_ALLOCATION_CLASSES - 1]) {
+        ba_os_free(ptr, total_size);
+        return;
+    }
+
+    int ac_index = size_to_idx(total_size);
     if (!BASE_ALLOC.ac[ac_index]) {
         // if creating ac failed, memory must have been allocated by os
-        ba_os_free(ptr, size);
+        ba_os_free(ptr, total_size);
         return;
     }
 
     umf_ba_free(BASE_ALLOC.ac[ac_index], ptr);
 }
+
+size_t umf_ba_global_malloc_usable_size(void *ptr) {
+    if (!ptr) {
+        return 0;
+    }
+
+    size_t usable_size;
+    get_original_alloc(ptr, NULL, &usable_size);
+
+    assert(usable_size >= sizeof(size_t));
+
+    return usable_size;
+}

src/base_alloc/base_alloc_global.h

@@ -15,8 +15,10 @@ extern "C" {
 #endif
 
 void *umf_ba_global_alloc(size_t size);
-void umf_ba_global_free(void *ptr, size_t size);
+void umf_ba_global_free(void *ptr);
 void umf_ba_destroy_global(void);
+size_t umf_ba_global_malloc_usable_size(void *ptr);
+void *umf_ba_global_aligned_alloc(size_t size, size_t alignment);
 
 #ifdef __cplusplus
 }