[CTL] Add allocation counter for disjoint pool

Author: KFilipek

src/pool/pool_disjoint.c

@@ -9,6 +9,7 @@
 #include <ctype.h>
 #include <errno.h>
 #include <inttypes.h>
+#include <stdatomic.h>
 #include <stdlib.h>
 #include <string.h>
 
@@ -33,6 +34,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,10 +68,25 @@ 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
+atomic_int counter_alloc_balance;
+// atomic_fetch_add(&counter_alloc_balance, 1);
+
+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;
+    *(int *)arg = atomic_load(&counter_alloc_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) {
+    atomic_init(&counter_alloc_balance, 0);
     CTL_REGISTER_MODULE(&disjoint_ctl_root, disjoint);
 }
 
@@ -745,7 +762,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);
-
+    atomic_fetch_add(&counter_alloc_balance, 1);
     return ptr;
 }
 
@@ -779,6 +796,7 @@ void *disjoint_pool_aligned_malloc(void *pool, size_t size, size_t alignment) {
     }
 
     if (alignment <= 1) {
+        atomic_fetch_add(&counter_alloc_balance, 1);
         return disjoint_pool_allocate(pool, size);
     }
 
@@ -809,6 +827,7 @@ void *disjoint_pool_aligned_malloc(void *pool, size_t size, size_t alignment) {
 
         assert(ptr);
         utils_annotate_memory_undefined(ptr, size);
+        atomic_fetch_add(&counter_alloc_balance, 1);
         return ptr;
     }
 
@@ -847,6 +866,7 @@ void *disjoint_pool_aligned_malloc(void *pool, size_t size, size_t alignment) {
                   (from_pool ? "pool" : "provider"), ptr);
     }
 
+    atomic_fetch_add(&counter_alloc_balance, 1);
     return aligned_ptr;
 }
 
@@ -914,6 +934,8 @@ 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 {
+            atomic_fetch_sub(&counter_alloc_balance, 1);
         }
 
         return ret;
@@ -958,6 +980,7 @@ umf_result_t disjoint_pool_free(void *pool, void *ptr) {
                   disjoint_pool->params.cur_pool_size);
     }
 
+    atomic_fetch_sub(&counter_alloc_balance, 1);
     return UMF_RESULT_SUCCESS;
 }
 

test/pools/disjoint_pool_ctl.cpp

@@ -85,6 +85,63 @@ 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));
+    PoolWrapper poolWrapper(providerWrapper.get(), umfDisjointPoolOps(),
+                            params);
+
+    // Check that the allocation balance is zero
+    int 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 = 1024; // 1 KB
+    const int max_allocations = 10;
+    void *ptr[max_allocations] = {nullptr};
+    int i = 0;
+    while (i < max_allocations) {
+        ptr[i] = umfPoolMalloc(poolWrapper.get(), allocation_size);
+        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, i);
+
+    // Check balance after freeing the allocations
+    for (int j = 0; j < max_allocations; ++j) {
+        if (ptr[j]) {
+            umfPoolFree(poolWrapper.get(), ptr[j]);
+        }
+    }
+    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 ==