implement malloc_usable_size and enable pool tests for disjoint

Author: pbalcer

src/pool/pool_disjoint.c

@@ -752,8 +752,10 @@ void *disjoint_pool_aligned_malloc(void *pool, size_t size, size_t alignment) {
     }
 
     void *aligned_ptr = (void *)ALIGN_UP_SAFE((size_t)ptr, alignment);
-    VALGRIND_DO_MEMPOOL_ALLOC(disjoint_pool, aligned_ptr, size);
-    utils_annotate_memory_undefined(aligned_ptr, size);
+    size_t diff = (ptrdiff_t)aligned_ptr - (ptrdiff_t)ptr;
+    size_t real_size = bucket->size - diff;
+    VALGRIND_DO_MEMPOOL_ALLOC(disjoint_pool, aligned_ptr, real_size);
+    utils_annotate_memory_undefined(aligned_ptr, real_size);
 
     utils_mutex_unlock(&bucket->bucket_lock);
 
@@ -767,11 +769,34 @@ void *disjoint_pool_aligned_malloc(void *pool, size_t size, size_t alignment) {
 }
 
 size_t disjoint_pool_malloc_usable_size(void *pool, void *ptr) {
-    (void)pool;
-    (void)ptr;
+    disjoint_pool_t *disjoint_pool = (disjoint_pool_t *)pool;
+    if (ptr == NULL) {
+        return 0;
+    }
 
-    // Not supported
-    return 0;
+    // check if given pointer is allocated inside any Disjoint Pool slab
+    slab_t *slab =
+        (slab_t *)critnib_find_le(disjoint_pool->known_slabs, (uintptr_t)ptr);
+    if (slab == NULL || ptr >= slab_get_end(slab)) {
+        // memory comes directly from the provider
+        umf_alloc_info_t allocInfo = {NULL, 0, NULL};
+        umf_result_t ret = umfMemoryTrackerGetAllocInfo(ptr, &allocInfo);
+        if (ret != UMF_RESULT_SUCCESS) {
+            return 0;
+        }
+
+        return allocInfo.baseSize;
+    }
+    // Get the unaligned pointer
+    // NOTE: the base pointer slab->mem_ptr needn't to be aligned to bucket size
+    size_t chunk_idx =
+        (((uintptr_t)ptr - (uintptr_t)slab->mem_ptr) / slab->bucket->size);
+    void *unaligned_ptr =
+        (void *)((uintptr_t)slab->mem_ptr + chunk_idx * slab->bucket->size);
+
+    ptrdiff_t diff = (ptrdiff_t)ptr - (ptrdiff_t)unaligned_ptr;
+
+    return slab->bucket->size - diff;
 }
 
 umf_result_t disjoint_pool_free(void *pool, void *ptr) {

test/common/pool.hpp

@@ -1,6 +1,6 @@
 /*
  *
- * Copyright (C) 2023-2024 Intel Corporation
+ * Copyright (C) 2023-2025 Intel Corporation
  *
  * Under the Apache License v2.0 with LLVM Exceptions. See LICENSE.TXT.
  * SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
@@ -19,6 +19,7 @@
 
 #include <umf/base.h>
 #include <umf/memory_provider.h>
+#include <umf/pools/pool_disjoint.h>
 
 #include "base.hpp"
 #include "cpp_helpers.hpp"
@@ -150,6 +151,49 @@ struct malloc_pool : public pool_base_t {
 umf_memory_pool_ops_t MALLOC_POOL_OPS =
     umf::poolMakeCOps<umf_test::malloc_pool, void>();
 
+static constexpr size_t DEFAULT_DISJOINT_SLAB_MIN_SIZE = 4096;
+static constexpr size_t DEFAULT_DISJOINT_MAX_POOLABLE_SIZE = 4096;
+static constexpr size_t DEFAULT_DISJOINT_CAPACITY = 4;
+static constexpr size_t DEFAULT_DISJOINT_MIN_BUCKET_SIZE = 64;
+
+inline void *defaultDisjointPoolConfig() {
+    umf_disjoint_pool_params_handle_t config = nullptr;
+    umf_result_t res = umfDisjointPoolParamsCreate(&config);
+    if (res != UMF_RESULT_SUCCESS) {
+        throw std::runtime_error("Failed to create pool params");
+    }
+    res = umfDisjointPoolParamsSetSlabMinSize(config,
+                                              DEFAULT_DISJOINT_SLAB_MIN_SIZE);
+    if (res != UMF_RESULT_SUCCESS) {
+        umfDisjointPoolParamsDestroy(config);
+        throw std::runtime_error("Failed to set slab min size");
+    }
+    res = umfDisjointPoolParamsSetMaxPoolableSize(
+        config, DEFAULT_DISJOINT_MAX_POOLABLE_SIZE);
+    if (res != UMF_RESULT_SUCCESS) {
+        umfDisjointPoolParamsDestroy(config);
+        throw std::runtime_error("Failed to set max poolable size");
+    }
+    res = umfDisjointPoolParamsSetCapacity(config, DEFAULT_DISJOINT_CAPACITY);
+    if (res != UMF_RESULT_SUCCESS) {
+        umfDisjointPoolParamsDestroy(config);
+        throw std::runtime_error("Failed to set capacity");
+    }
+    res = umfDisjointPoolParamsSetMinBucketSize(
+        config, DEFAULT_DISJOINT_MIN_BUCKET_SIZE);
+    if (res != UMF_RESULT_SUCCESS) {
+        umfDisjointPoolParamsDestroy(config);
+        throw std::runtime_error("Failed to set min bucket size");
+    }
+
+    return config;
+}
+
+inline umf_result_t defaultDisjointPoolConfigDestroy(void *config) {
+    return umfDisjointPoolParamsDestroy(
+        static_cast<umf_disjoint_pool_params_handle_t>(config));
+}
+
 } // namespace umf_test
 
 #endif /* UMF_TEST_POOL_HPP */

test/memoryPoolAPI.cpp

@@ -12,6 +12,7 @@
 #include "test_helpers.h"
 
 #include <umf/memory_provider.h>
+#include <umf/pools/pool_disjoint.h>
 #include <umf/pools/pool_proxy.h>
 
 #ifdef UMF_PROXY_LIB_ENABLED
@@ -295,12 +296,14 @@ TEST_F(tagTest, SetAndGetInvalidPool) {
 
 INSTANTIATE_TEST_SUITE_P(
     mallocPoolTest, umfPoolTest,
-    ::testing::Values(poolCreateExtParams{&MALLOC_POOL_OPS, nullptr, nullptr,
-                                          &UMF_NULL_PROVIDER_OPS, nullptr,
-                                          nullptr},
-                      poolCreateExtParams{umfProxyPoolOps(), nullptr, nullptr,
-                                          &BA_GLOBAL_PROVIDER_OPS, nullptr,
-                                          nullptr}));
+    ::testing::Values(
+        poolCreateExtParams{&MALLOC_POOL_OPS, nullptr, nullptr,
+                            &UMF_NULL_PROVIDER_OPS, nullptr, nullptr},
+        poolCreateExtParams{umfProxyPoolOps(), nullptr, nullptr,
+                            &BA_GLOBAL_PROVIDER_OPS, nullptr, nullptr},
+        poolCreateExtParams{umfDisjointPoolOps(), defaultDisjointPoolConfig,
+                            defaultDisjointPoolConfigDestroy,
+                            &BA_GLOBAL_PROVIDER_OPS, nullptr, nullptr}));
 
 INSTANTIATE_TEST_SUITE_P(mallocMultiPoolTest, umfMultiPoolTest,
                          ::testing::Values(poolCreateExtParams{

test/poolFixtures.hpp

@@ -452,26 +452,45 @@ TEST_P(umfPoolTest, allocMaxSize) {
 }
 
 TEST_P(umfPoolTest, mallocUsableSize) {
+    [[maybe_unused]] auto pool_ops = std::get<0>(this->GetParam());
+#ifdef _WIN32
+    if (pool_ops == &umf_test::MALLOC_POOL_OPS) {
+        GTEST_SKIP()
+            << "Windows Malloc Pool does not support umfPoolAlignedMalloc";
+    }
+#endif
+    if (!umf_test::isAlignedAllocSupported(pool.get())) {
+        GTEST_SKIP();
+    }
 #ifdef __SANITIZE_ADDRESS__
-    // Sanitizer replaces malloc_usable_size implementation with its own
-    GTEST_SKIP()
-        << "This test is invalid with AddressSanitizer instrumentation";
-#else
+    if (pool_ops == &umf_test::MALLOC_POOL_OPS) {
+        // Sanitizer replaces malloc_usable_size implementation with its own
+        GTEST_SKIP()
+            << "This test is invalid with AddressSanitizer instrumentation";
+    }
+#endif
+    for (size_t allocSize :
+         {32, 64, 1 << 6, 1 << 10, 1 << 13, 1 << 16, 1 << 19}) {
+        for (size_t alignment : {0, 1 << 6, 1 << 8, 1 << 12}) {
+            if (alignment >= allocSize) {
+                continue;
+            }
+            void *ptr = nullptr;
+            if (alignment == 0) {
+                ptr = umfPoolMalloc(pool.get(), allocSize);
+            } else {
+                ptr = umfPoolAlignedMalloc(pool.get(), allocSize, alignment);
+            }
+            ASSERT_NE(ptr, nullptr);
+            size_t result = umfPoolMallocUsableSize(pool.get(), ptr);
+            ASSERT_TRUE(result == 0 || result >= allocSize);
 
-    for (size_t allocSize : {32, 48, 1024, 8192}) {
-        char *ptr = static_cast<char *>(umfPoolMalloc(pool.get(), allocSize));
-        ASSERT_NE(ptr, nullptr);
-        size_t result = umfPoolMallocUsableSize(pool.get(), ptr);
-        ASSERT_TRUE(result == 0 || result >= allocSize);
+            // Make sure we can write to this memory
+            memset(ptr, 123, result);
 
-        // Make sure we can write to this memory
-        for (size_t i = 0; i < result; i++) {
-            ptr[i] = 123;
+            umfPoolFree(pool.get(), ptr);
         }
-
-        umfPoolFree(pool.get(), ptr);
     }
-#endif
 }
 
 #endif /* UMF_TEST_POOL_FIXTURES_HPP */

test/pools/disjoint_pool.cpp

@@ -13,49 +13,6 @@
 #include "provider_null.h"
 #include "provider_trace.h"
 
-static constexpr size_t DEFAULT_DISJOINT_SLAB_MIN_SIZE = 4096;
-static constexpr size_t DEFAULT_DISJOINT_MAX_POOLABLE_SIZE = 4096;
-static constexpr size_t DEFAULT_DISJOINT_CAPACITY = 4;
-static constexpr size_t DEFAULT_DISJOINT_MIN_BUCKET_SIZE = 64;
-
-void *defaultPoolConfig() {
-    umf_disjoint_pool_params_handle_t config = nullptr;
-    umf_result_t res = umfDisjointPoolParamsCreate(&config);
-    if (res != UMF_RESULT_SUCCESS) {
-        throw std::runtime_error("Failed to create pool params");
-    }
-    res = umfDisjointPoolParamsSetSlabMinSize(config,
-                                              DEFAULT_DISJOINT_SLAB_MIN_SIZE);
-    if (res != UMF_RESULT_SUCCESS) {
-        umfDisjointPoolParamsDestroy(config);
-        throw std::runtime_error("Failed to set slab min size");
-    }
-    res = umfDisjointPoolParamsSetMaxPoolableSize(
-        config, DEFAULT_DISJOINT_MAX_POOLABLE_SIZE);
-    if (res != UMF_RESULT_SUCCESS) {
-        umfDisjointPoolParamsDestroy(config);
-        throw std::runtime_error("Failed to set max poolable size");
-    }
-    res = umfDisjointPoolParamsSetCapacity(config, DEFAULT_DISJOINT_CAPACITY);
-    if (res != UMF_RESULT_SUCCESS) {
-        umfDisjointPoolParamsDestroy(config);
-        throw std::runtime_error("Failed to set capacity");
-    }
-    res = umfDisjointPoolParamsSetMinBucketSize(
-        config, DEFAULT_DISJOINT_MIN_BUCKET_SIZE);
-    if (res != UMF_RESULT_SUCCESS) {
-        umfDisjointPoolParamsDestroy(config);
-        throw std::runtime_error("Failed to set min bucket size");
-    }
-
-    return config;
-}
-
-umf_result_t poolConfigDestroy(void *config) {
-    return umfDisjointPoolParamsDestroy(
-        static_cast<umf_disjoint_pool_params_handle_t>(config));
-}
-
 using umf_test::test;
 using namespace umf_test;
 
@@ -92,7 +49,7 @@ TEST_F(test, internals) {
     provider_handle = providerUnique.get();
 
     umf_disjoint_pool_params_handle_t params =
-        (umf_disjoint_pool_params_handle_t)defaultPoolConfig();
+        (umf_disjoint_pool_params_handle_t)defaultDisjointPoolConfig();
     // set to maximum tracing
     params->pool_trace = 3;
     params->max_poolable_size = 1024 * 1024;
@@ -256,7 +213,7 @@ TEST_F(test, sharedLimits) {
     static constexpr size_t MaxSize = 4 * SlabMinSize;
 
     umf_disjoint_pool_params_handle_t params =
-        (umf_disjoint_pool_params_handle_t)defaultPoolConfig();
+        (umf_disjoint_pool_params_handle_t)defaultDisjointPoolConfig();
     umf_result_t ret = umfDisjointPoolParamsSetSlabMinSize(params, SlabMinSize);
     EXPECT_EQ(ret, UMF_RESULT_SUCCESS);
 
@@ -373,22 +330,23 @@ TEST_F(test, disjointPoolInvalidBucketSize) {
 
 INSTANTIATE_TEST_SUITE_P(disjointPoolTests, umfPoolTest,
                          ::testing::Values(poolCreateExtParams{
-                             umfDisjointPoolOps(), defaultPoolConfig,
-                             poolConfigDestroy, &BA_GLOBAL_PROVIDER_OPS,
-                             nullptr, nullptr}));
+                             umfDisjointPoolOps(), defaultDisjointPoolConfig,
+                             defaultDisjointPoolConfigDestroy,
+                             &BA_GLOBAL_PROVIDER_OPS, nullptr, nullptr}));
 
 void *memProviderParams() { return (void *)&DEFAULT_DISJOINT_CAPACITY; }
 
 INSTANTIATE_TEST_SUITE_P(
     disjointPoolTests, umfMemTest,
     ::testing::Values(std::make_tuple(
-        poolCreateExtParams{umfDisjointPoolOps(), defaultPoolConfig,
-                            poolConfigDestroy, &MOCK_OUT_OF_MEM_PROVIDER_OPS,
-                            memProviderParams, nullptr},
+        poolCreateExtParams{umfDisjointPoolOps(), defaultDisjointPoolConfig,
+                            defaultDisjointPoolConfigDestroy,
+                            &MOCK_OUT_OF_MEM_PROVIDER_OPS, memProviderParams,
+                            nullptr},
         static_cast<int>(DEFAULT_DISJOINT_CAPACITY) / 2)));
 
 INSTANTIATE_TEST_SUITE_P(disjointMultiPoolTests, umfMultiPoolTest,
                          ::testing::Values(poolCreateExtParams{
-                             umfDisjointPoolOps(), defaultPoolConfig,
-                             poolConfigDestroy, &BA_GLOBAL_PROVIDER_OPS,
-                             nullptr, nullptr}));
+                             umfDisjointPoolOps(), defaultDisjointPoolConfig,
+                             defaultDisjointPoolConfigDestroy,
+                             &BA_GLOBAL_PROVIDER_OPS, nullptr, nullptr}));