Skip to content

Commit d30e9f9

Browse files
committed
[CTL] Add allocation counter for disjoint pool
1 parent 222e26d commit d30e9f9

File tree

2 files changed

+102
-6
lines changed

2 files changed

+102
-6
lines changed

src/pool/pool_disjoint.c

Lines changed: 30 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ static char *DEFAULT_NAME = "disjoint";
3333
struct ctl disjoint_ctl_root;
3434
static UTIL_ONCE_FLAG ctl_initialized = UTIL_ONCE_FLAG_INIT;
3535

36+
// CTL: name attribute
3637
static int CTL_READ_HANDLER(name)(void *ctx, umf_ctl_query_source_t source,
3738
void *arg, size_t size,
3839
umf_ctl_index_utlist_t *indexes,
@@ -66,8 +67,28 @@ static int CTL_WRITE_HANDLER(name)(void *ctx, umf_ctl_query_source_t source,
6667
return 0;
6768
}
6869

69-
static const umf_ctl_node_t CTL_NODE(disjoint)[] = {CTL_LEAF_RW(name),
70-
CTL_NODE_END};
70+
// CTL: allocation counters
71+
static uint64_t allocation_balance = 0;
72+
73+
static int CTL_READ_HANDLER(allocation_balance)(
74+
void *ctx, umf_ctl_query_source_t source, void *arg, size_t size,
75+
umf_ctl_index_utlist_t *indexes, const char *extra_name,
76+
umf_ctl_query_type_t queryType) {
77+
(void)ctx, (void)source, (void)size, (void)indexes, (void)extra_name,
78+
(void)queryType;
79+
disjoint_pool_t *pool = (disjoint_pool_t *)ctx;
80+
if (arg == NULL) {
81+
return -1;
82+
}
83+
uint64_t *balance = (uint64_t *)arg;
84+
*balance = 0;
85+
utils_atomic_load_acquire_u64(&allocation_balance, balance);
86+
87+
return 0;
88+
}
89+
90+
static const umf_ctl_node_t CTL_NODE(disjoint)[] = {
91+
CTL_LEAF_RW(name), CTL_LEAF_RO(allocation_balance), CTL_NODE_END};
7192

7293
static void initialize_disjoint_ctl(void) {
7394
CTL_REGISTER_MODULE(&disjoint_ctl_root, disjoint);
@@ -579,7 +600,6 @@ static void *disjoint_pool_allocate(disjoint_pool_t *pool, size_t size) {
579600
}
580601

581602
void *ptr = NULL;
582-
583603
if (size > pool->params.max_poolable_size) {
584604
umf_result_t ret =
585605
umfMemoryProviderAlloc(pool->provider, size, 0, &ptr);
@@ -755,7 +775,7 @@ umf_result_t disjoint_pool_initialize(umf_memory_provider_handle_t provider,
755775
void *disjoint_pool_malloc(void *pool, size_t size) {
756776
disjoint_pool_t *hPool = (disjoint_pool_t *)pool;
757777
void *ptr = disjoint_pool_allocate(hPool, size);
758-
778+
utils_atomic_increment_u64(&allocation_balance);
759779
return ptr;
760780
}
761781

@@ -939,8 +959,9 @@ umf_result_t disjoint_pool_free(void *pool, void *ptr) {
939959
if (ret != UMF_RESULT_SUCCESS) {
940960
TLS_last_allocation_error = ret;
941961
LOG_ERR("deallocation from the memory provider failed");
962+
} else {
963+
utils_atomic_decrement_u64(&allocation_balance);
942964
}
943-
944965
return ret;
945966
}
946967

@@ -971,6 +992,9 @@ umf_result_t disjoint_pool_free(void *pool, void *ptr) {
971992
critnib_release(disjoint_pool->known_slabs, ref_slab);
972993

973994
if (disjoint_pool->params.pool_trace > 1) {
995+
printf("Freeing %8zu %s bytes from %s -> %p\n", bucket->size,
996+
disjoint_pool->params.name, (to_pool ? "pool" : "provider"),
997+
unaligned_ptr);
974998
bucket->free_count++;
975999
}
9761000

@@ -985,7 +1009,7 @@ umf_result_t disjoint_pool_free(void *pool, void *ptr) {
9851009
disjoint_pool_get_limits(disjoint_pool)->total_size, name,
9861010
disjoint_pool->params.cur_pool_size);
9871011
}
988-
1012+
utils_atomic_decrement_u64(&allocation_balance);
9891013
return UMF_RESULT_SUCCESS;
9901014
}
9911015

test/pools/disjoint_pool_ctl.cpp

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
// Under the Apache License v2.0 with LLVM Exceptions. See LICENSE.TXT.
33
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exceptiongi
44

5+
#include <cstdint>
56
#include <gtest/gtest.h>
67
#include <umf/memory_pool.h>
78
#include <umf/memory_provider.h>
@@ -85,6 +86,77 @@ class ProviderWrapper {
8586
void *m_params;
8687
};
8788

89+
TEST_F(test, disjointCtlAllocationBalance) {
90+
umf_os_memory_provider_params_handle_t os_memory_provider_params = nullptr;
91+
if (UMF_RESULT_ERROR_NOT_SUPPORTED ==
92+
umfOsMemoryProviderParamsCreate(&os_memory_provider_params)) {
93+
GTEST_SKIP() << "OS memory provider is not supported!";
94+
}
95+
96+
ProviderWrapper providerWrapper(umfOsMemoryProviderOps(),
97+
os_memory_provider_params);
98+
if (providerWrapper.get() == NULL) {
99+
GTEST_SKIP() << "OS memory provider is not supported!";
100+
}
101+
102+
umf_disjoint_pool_params_handle_t params = nullptr;
103+
ASSERT_SUCCESS(umfDisjointPoolParamsCreate(&params));
104+
105+
// Set max poolable size to a reasonable value
106+
ASSERT_SUCCESS(
107+
umfDisjointPoolParamsSetMaxPoolableSize(params, 1024 * 1024));
108+
109+
// Set the capacity of the pool to a reasonable value
110+
ASSERT_SUCCESS(
111+
umfDisjointPoolParamsSetSlabMinSize(params, 64 * 1024)); // 64 KiB
112+
113+
// Set the trace level to 3 to enable allocation balance tracking
114+
ASSERT_SUCCESS(umfDisjointPoolParamsSetTrace(params, 3));
115+
116+
PoolWrapper poolWrapper(providerWrapper.get(), umfDisjointPoolOps(),
117+
params);
118+
119+
// Check that the allocation balance is zero
120+
uint64_t allocation_balance = 0;
121+
ASSERT_SUCCESS(umfCtlGet("umf.pool.by_handle.disjoint.allocation_balance",
122+
poolWrapper.get(), &allocation_balance,
123+
sizeof(allocation_balance)));
124+
ASSERT_EQ(allocation_balance, 0ull);
125+
126+
// Allocate some memory from the pool
127+
size_t allocation_size = 64; // 64 B
128+
const uint64_t max_allocations = 2;
129+
void *ptr[max_allocations] = {nullptr};
130+
int i = 0;
131+
while (i < max_allocations) {
132+
ptr[i] = umfPoolMalloc(poolWrapper.get(), allocation_size);
133+
ASSERT_NE(ptr[i], nullptr);
134+
++i;
135+
}
136+
137+
// Check the allocation balance after allocations
138+
ASSERT_SUCCESS(umfCtlGet("umf.pool.by_handle.disjoint.allocation_balance",
139+
poolWrapper.get(), &allocation_balance,
140+
sizeof(allocation_balance)));
141+
ASSERT_EQ(allocation_balance, max_allocations);
142+
143+
// Check balance after freeing the allocations
144+
for (int j = 0; j < max_allocations; ++j) {
145+
if (ptr[j]) {
146+
ASSERT_EQ(umfPoolFree(poolWrapper.get(), ptr[j]),
147+
UMF_RESULT_SUCCESS);
148+
}
149+
}
150+
allocation_balance = 123;
151+
ASSERT_SUCCESS(umfCtlGet("umf.pool.by_handle.disjoint.allocation_balance",
152+
poolWrapper.get(), &allocation_balance,
153+
sizeof(allocation_balance)));
154+
ASSERT_EQ(allocation_balance, 0ull);
155+
156+
ASSERT_SUCCESS(umfDisjointPoolParamsDestroy(params));
157+
ASSERT_SUCCESS(umfOsMemoryProviderParamsDestroy(os_memory_provider_params));
158+
}
159+
88160
TEST_F(test, disjointCtlName) {
89161
umf_os_memory_provider_params_handle_t os_memory_provider_params = nullptr;
90162
if (UMF_RESULT_ERROR_NOT_SUPPORTED ==

0 commit comments

Comments
 (0)