diff --git a/src/libumf.c b/src/libumf.c index f5e2004ede..0cfb62062f 100644 --- a/src/libumf.c +++ b/src/libumf.c @@ -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) diff --git a/src/pool/pool_disjoint.c b/src/pool/pool_disjoint.c index 6629c9a5b0..8aa38c86fe 100644 --- a/src/pool/pool_disjoint.c +++ b/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,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)( + 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); @@ -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); @@ -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; } @@ -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; } @@ -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++; } @@ -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; } diff --git a/test/pools/disjoint_pool_ctl.cpp b/test/pools/disjoint_pool_ctl.cpp index e683b9bc0f..ad501241ab 100644 --- a/test/pools/disjoint_pool_ctl.cpp +++ b/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 #include #include #include @@ -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(¶ms)); + + // 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 ==