Skip to content

[CTL] Add disjoint pool allocation balance counter #1334

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open

[CTL] Add disjoint pool allocation balance counter #1334

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
222e26d
[CTL] Add size check for umfCtlExec
KFilipek May 28, 2025
2c7329a
[CTL] Add allocation counter for disjoint pool
KFilipek May 28, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions src/libumf.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -127,6 +127,9 @@ umf_result_t umfCtlExec(const char *name, void *ctx, void *arg, size_t size) {
if (name == NULL) {
return UMF_RESULT_ERROR_INVALID_ARGUMENT;
}
if ((arg == NULL && size != 0) || (arg != NULL && size == 0)) {
return UMF_RESULT_ERROR_INVALID_ARGUMENT;
}

return ctl_query(NULL, ctx, CTL_QUERY_PROGRAMMATIC, name,
CTL_QUERY_RUNNABLE, arg, size)
Expand Down
35 changes: 29 additions & 6 deletions src/pool/pool_disjoint.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

isn't this something we could implement generically for any pool implementation in the umfPoolAlloc/umfPoolFree interface?

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,
Expand Down Expand Up @@ -66,8 +67,27 @@ 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
static uint64_t allocation_balance = 0;

static int CTL_READ_HANDLER(allocation_balance)(
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

we should have docs that list all ctl interfaces.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I will prepare the docs after this PR, which we discussed previously (with @lplewa @bratpiorka).

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;
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);
Expand Down Expand Up @@ -579,7 +599,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);
Expand Down Expand Up @@ -755,7 +774,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;
}

Expand Down Expand Up @@ -939,8 +958,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;
}

Expand Down Expand Up @@ -971,6 +991,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++;
}

Expand All @@ -985,7 +1008,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;
}

Expand Down
72 changes: 72 additions & 0 deletions test/pools/disjoint_pool_ctl.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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>
Expand Down Expand Up @@ -85,6 +86,77 @@ 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, 0ull);

// Allocate some memory from the pool
size_t allocation_size = 64; // 64 B
const uint64_t max_allocations = 2;
void *ptr[max_allocations] = {nullptr};
uint64_t 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, max_allocations);

// Check balance after freeing the allocations
for (uint64_t j = 0; j < max_allocations; ++j) {
if (ptr[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, 0ull);

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 ==
Expand Down
Loading