[CTL] Add allocation counter for disjoint pool

Author: KFilipek

src/pool/pool_disjoint.c

@@ -33,6 +33,7 @@ static char *DEFAULT_NAME = "disjoint";
 struct ctl disjoint_ctl_root;
 static UTIL_ONCE_FLAG ctl_initialized = UTIL_ONCE_FLAG_INIT;
 
+// CTL: name attribute
 static int CTL_READ_HANDLER(name)(void *ctx, umf_ctl_query_source_t source,
                                   void *arg, size_t size,
                                   umf_ctl_index_utlist_t *indexes,
@@ -66,8 +67,30 @@ static int CTL_WRITE_HANDLER(name)(void *ctx, umf_ctl_query_source_t source,
     return 0;
 }
 
-static const umf_ctl_node_t CTL_NODE(disjoint)[] = {CTL_LEAF_RW(name),
-                                                    CTL_NODE_END};
+// CTL: allocation counters
+uint64_t allocation_balance = 0;
+
+static int CTL_READ_HANDLER(allocation_balance)(
+    void *ctx, umf_ctl_query_source_t source, void *arg, size_t size,
+    umf_ctl_index_utlist_t *indexes, const char *extra_name,
+    umf_ctl_query_type_t queryType) {
+    (void)ctx, (void)source, (void)size, (void)indexes, (void)extra_name,
+        (void)queryType;
+    // Calculate from pool buckets
+    disjoint_pool_t *pool = (disjoint_pool_t *)ctx;
+    printf("Calculating allocation balance for pool: %p\n", (void *)pool);
+    if (arg == NULL) {
+        return -1;
+    }
+    uint64_t *balance = (uint64_t *)arg;
+    *balance = 0;
+    utils_atomic_load_acquire_u64(&allocation_balance, balance);
+
+    return 0;
+}
+
+static const umf_ctl_node_t CTL_NODE(disjoint)[] = {
+    CTL_LEAF_RW(name), CTL_LEAF_RO(allocation_balance), CTL_NODE_END};
 
 static void initialize_disjoint_ctl(void) {
     CTL_REGISTER_MODULE(&disjoint_ctl_root, disjoint);
@@ -579,7 +602,6 @@ static void *disjoint_pool_allocate(disjoint_pool_t *pool, size_t size) {
     }
 
     void *ptr = NULL;
-
     if (size > pool->params.max_poolable_size) {
         umf_result_t ret =
             umfMemoryProviderAlloc(pool->provider, size, 0, &ptr);
@@ -755,7 +777,7 @@ umf_result_t disjoint_pool_initialize(umf_memory_provider_handle_t provider,
 void *disjoint_pool_malloc(void *pool, size_t size) {
     disjoint_pool_t *hPool = (disjoint_pool_t *)pool;
     void *ptr = disjoint_pool_allocate(hPool, size);
-
+    utils_atomic_increment_u64(&allocation_balance);
     return ptr;
 }
 
@@ -939,8 +961,9 @@ umf_result_t disjoint_pool_free(void *pool, void *ptr) {
         if (ret != UMF_RESULT_SUCCESS) {
             TLS_last_allocation_error = ret;
             LOG_ERR("deallocation from the memory provider failed");
+        } else {
+            utils_atomic_decrement_u64(&allocation_balance);
         }
-
         return ret;
     }
 
@@ -971,6 +994,9 @@ umf_result_t disjoint_pool_free(void *pool, void *ptr) {
     critnib_release(disjoint_pool->known_slabs, ref_slab);
 
     if (disjoint_pool->params.pool_trace > 1) {
+        printf("Freeing %8zu %s bytes from %s -> %p\n", bucket->size,
+               disjoint_pool->params.name, (to_pool ? "pool" : "provider"),
+               unaligned_ptr);
         bucket->free_count++;
     }
 
@@ -985,7 +1011,7 @@ umf_result_t disjoint_pool_free(void *pool, void *ptr) {
                   disjoint_pool_get_limits(disjoint_pool)->total_size, name,
                   disjoint_pool->params.cur_pool_size);
     }
-
+    utils_atomic_decrement_u64(&allocation_balance);
     return UMF_RESULT_SUCCESS;
 }
 

test/pools/disjoint_pool_ctl.cpp

@@ -2,6 +2,7 @@
 // Under the Apache License v2.0 with LLVM Exceptions. See LICENSE.TXT.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exceptiongi
 
+#include <cstdint>
 #include <gtest/gtest.h>
 #include <umf/memory_pool.h>
 #include <umf/memory_provider.h>
@@ -85,6 +86,79 @@ class ProviderWrapper {
     void *m_params;
 };
 
+TEST_F(test, disjointCtlAllocationBalance) {
+    umf_os_memory_provider_params_handle_t os_memory_provider_params = nullptr;
+    if (UMF_RESULT_ERROR_NOT_SUPPORTED ==
+        umfOsMemoryProviderParamsCreate(&os_memory_provider_params)) {
+        GTEST_SKIP() << "OS memory provider is not supported!";
+    }
+
+    ProviderWrapper providerWrapper(umfOsMemoryProviderOps(),
+                                    os_memory_provider_params);
+    if (providerWrapper.get() == NULL) {
+        GTEST_SKIP() << "OS memory provider is not supported!";
+    }
+
+    umf_disjoint_pool_params_handle_t params = nullptr;
+    ASSERT_SUCCESS(umfDisjointPoolParamsCreate(&params));
+
+    // Set max poolable size to a reasonable value
+    ASSERT_SUCCESS(
+        umfDisjointPoolParamsSetMaxPoolableSize(params, 1024 * 1024));
+
+    // Set the capacity of the pool to a reasonable value
+    ASSERT_SUCCESS(
+        umfDisjointPoolParamsSetSlabMinSize(params, 64 * 1024)); // 64 KiB
+
+    // Set the trace level to 3 to enable allocation balance tracking
+    ASSERT_SUCCESS(umfDisjointPoolParamsSetTrace(params, 3));
+
+    PoolWrapper poolWrapper(providerWrapper.get(), umfDisjointPoolOps(),
+                            params);
+
+    // Check that the allocation balance is zero
+    uint64_t allocation_balance = 0;
+    ASSERT_SUCCESS(umfCtlGet("umf.pool.by_handle.disjoint.allocation_balance",
+                             poolWrapper.get(), &allocation_balance,
+                             sizeof(allocation_balance)));
+    ASSERT_EQ(allocation_balance, 0);
+
+    // Allocate some memory from the pool
+    size_t allocation_size = 64; // 64 B
+    const int max_allocations = 2;
+    void *ptr[max_allocations] = {nullptr};
+    int i = 0;
+    while (i < max_allocations) {
+        ptr[i] = umfPoolMalloc(poolWrapper.get(), allocation_size);
+        printf("Allocated %d: %p\n", i, ptr[i]);
+        ASSERT_NE(ptr[i], nullptr);
+        ++i;
+    }
+
+    // Check the allocation balance after allocations
+    ASSERT_SUCCESS(umfCtlGet("umf.pool.by_handle.disjoint.allocation_balance",
+                             poolWrapper.get(), &allocation_balance,
+                             sizeof(allocation_balance)));
+    ASSERT_EQ(allocation_balance, max_allocations);
+
+    // Check balance after freeing the allocations
+    for (int j = 0; j < max_allocations; ++j) {
+        if (ptr[j]) {
+            printf("Freeing allocation %p with %d\n", ptr[j], j);
+            ASSERT_EQ(umfPoolFree(poolWrapper.get(), ptr[j]),
+                      UMF_RESULT_SUCCESS);
+        }
+    }
+    allocation_balance = 123;
+    ASSERT_SUCCESS(umfCtlGet("umf.pool.by_handle.disjoint.allocation_balance",
+                             poolWrapper.get(), &allocation_balance,
+                             sizeof(allocation_balance)));
+    ASSERT_EQ(allocation_balance, 0);
+
+    ASSERT_SUCCESS(umfDisjointPoolParamsDestroy(params));
+    ASSERT_SUCCESS(umfOsMemoryProviderParamsDestroy(os_memory_provider_params));
+}
+
 TEST_F(test, disjointCtlName) {
     umf_os_memory_provider_params_handle_t os_memory_provider_params = nullptr;
     if (UMF_RESULT_ERROR_NOT_SUPPORTED ==