diff --git a/include/umf/memory_pool.h b/include/umf/memory_pool.h index a93d400f92..de045acf49 100644 --- a/include/umf/memory_pool.h +++ b/include/umf/memory_pool.h @@ -170,6 +170,22 @@ umf_memory_pool_handle_t umfPoolByPtr(const void *ptr); umf_result_t umfPoolGetMemoryProvider(umf_memory_pool_handle_t hPool, umf_memory_provider_handle_t *hProvider); +/// +/// @brief Set a custom tag on the memory pool that can be later retrieved using umfPoolGetTag. +/// @param hPool specified memory pool +/// @param tag tag to be set +/// @param oldTag [out][optional] previous tag set on the memory pool +/// @return UMF_RESULT_SUCCESS on success or appropriate error code on failure. +umf_result_t umfPoolSetTag(umf_memory_pool_handle_t hPool, void *tag, + void **oldTag); + +/// +/// @brief Retrieve the tag associated with the memory pool or NULL if no tag is set. +/// @param hPool specified memory pool +/// @param tag [out] tag associated with the memory pool +/// @return UMF_RESULT_SUCCESS on success. +umf_result_t umfPoolGetTag(umf_memory_pool_handle_t hPool, void **tag); + #ifdef __cplusplus } #endif diff --git a/src/libumf.def b/src/libumf.def index 0b4588bb81..f0f38ee15c 100644 --- a/src/libumf.def +++ b/src/libumf.def @@ -104,10 +104,12 @@ EXPORTS umfPoolFree umfPoolGetIPCHandleSize umfPoolGetLastAllocationError + umfPoolGetTag umfPoolGetMemoryProvider umfPoolMalloc umfPoolMallocUsableSize umfPoolRealloc + umfPoolSetTag umfProxyPoolOps umfPutIPCHandle umfScalablePoolOps diff --git a/src/libumf.map b/src/libumf.map index 41467bad59..fd1d48d348 100644 --- a/src/libumf.map +++ b/src/libumf.map @@ -99,9 +99,11 @@ UMF_1.0 { umfPoolGetIPCHandleSize; umfPoolGetLastAllocationError; umfPoolGetMemoryProvider; + umfPoolGetTag; umfPoolMalloc; umfPoolMallocUsableSize; umfPoolRealloc; + umfPoolSetTag; umfProxyPoolOps; umfPutIPCHandle; umfScalablePoolOps; diff --git a/src/memory_pool.c b/src/memory_pool.c index 4a85955efa..f4289c2151 100644 --- a/src/memory_pool.c +++ b/src/memory_pool.c @@ -55,6 +55,13 @@ static umf_result_t umfPoolCreateInternal(const umf_memory_pool_ops_t *ops, pool->flags = flags; pool->ops = *ops; + pool->tag = NULL; + + if (NULL == utils_mutex_init(&pool->lock)) { + LOG_ERR("Failed to initialize mutex for pool"); + ret = UMF_RESULT_ERROR_UNKNOWN; + goto err_lock_init; + } ret = ops->initialize(pool->provider, params, &pool->pool_priv); if (ret != UMF_RESULT_SUCCESS) { @@ -66,6 +73,8 @@ static umf_result_t umfPoolCreateInternal(const umf_memory_pool_ops_t *ops, return UMF_RESULT_SUCCESS; err_pool_init: + utils_mutex_destroy_not_free(&pool->lock); +err_lock_init: if (!(flags & UMF_POOL_CREATE_FLAG_DISABLE_TRACKING)) { umfMemoryProviderDestroy(pool->provider); } @@ -90,6 +99,8 @@ void umfPoolDestroy(umf_memory_pool_handle_t hPool) { umfMemoryProviderDestroy(hUpstreamProvider); } + utils_mutex_destroy_not_free(&hPool->lock); + LOG_INFO("Memory pool destroyed: %p", (void *)hPool); // TODO: this free keeps memory in base allocator, so it can lead to OOM in some scenarios (it should be optimized) @@ -175,3 +186,24 @@ umf_result_t umfPoolGetLastAllocationError(umf_memory_pool_handle_t hPool) { UMF_CHECK((hPool != NULL), UMF_RESULT_ERROR_INVALID_ARGUMENT); return hPool->ops.get_last_allocation_error(hPool->pool_priv); } + +umf_result_t umfPoolSetTag(umf_memory_pool_handle_t hPool, void *tag, + void **oldTag) { + UMF_CHECK((hPool != NULL), UMF_RESULT_ERROR_INVALID_ARGUMENT); + utils_mutex_lock(&hPool->lock); + if (oldTag) { + *oldTag = hPool->tag; + } + hPool->tag = tag; + utils_mutex_unlock(&hPool->lock); + return UMF_RESULT_SUCCESS; +} + +umf_result_t umfPoolGetTag(umf_memory_pool_handle_t hPool, void **tag) { + UMF_CHECK((hPool != NULL), UMF_RESULT_ERROR_INVALID_ARGUMENT); + UMF_CHECK((tag != NULL), UMF_RESULT_ERROR_INVALID_ARGUMENT); + utils_mutex_lock(&hPool->lock); + *tag = hPool->tag; + utils_mutex_unlock(&hPool->lock); + return UMF_RESULT_SUCCESS; +} diff --git a/src/memory_pool_internal.h b/src/memory_pool_internal.h index 90f2f16298..e556ace214 100644 --- a/src/memory_pool_internal.h +++ b/src/memory_pool_internal.h @@ -22,6 +22,7 @@ extern "C" { #endif #include "base_alloc.h" +#include "utils_concurrency.h" typedef struct umf_memory_pool_t { void *pool_priv; @@ -30,6 +31,9 @@ typedef struct umf_memory_pool_t { // Memory provider used by the pool. umf_memory_provider_handle_t provider; + + utils_mutex_t lock; + void *tag; } umf_memory_pool_t; #ifdef __cplusplus diff --git a/test/memoryPoolAPI.cpp b/test/memoryPoolAPI.cpp index 1c6d83f2af..ec137b5493 100644 --- a/test/memoryPoolAPI.cpp +++ b/test/memoryPoolAPI.cpp @@ -178,6 +178,121 @@ TEST_F(test, BasicPoolByPtrTest) { ASSERT_EQ(ret, UMF_RESULT_SUCCESS); } +struct tagTest : umf_test::test { + void SetUp() override { + test::SetUp(); + provider = umf_test::wrapProviderUnique(nullProviderCreate()); + pool = umf_test::wrapPoolUnique( + createPoolChecked(umfProxyPoolOps(), provider.get(), nullptr)); + } + + umf::provider_unique_handle_t provider; + umf::pool_unique_handle_t pool; +}; + +TEST_F(tagTest, SetAndGet) { + umf_result_t ret = umfPoolSetTag(pool.get(), (void *)0x99, nullptr); + ASSERT_EQ(ret, UMF_RESULT_SUCCESS); + + void *tag; + ret = umfPoolGetTag(pool.get(), &tag); + ASSERT_EQ(ret, UMF_RESULT_SUCCESS); + ASSERT_EQ(tag, (void *)0x99); + + void *oldTag; + ret = umfPoolSetTag(pool.get(), (void *)0x100, &oldTag); + ASSERT_EQ(ret, UMF_RESULT_SUCCESS); + ASSERT_EQ(oldTag, (void *)0x99); + + ret = umfPoolGetTag(pool.get(), &tag); + ASSERT_EQ(ret, UMF_RESULT_SUCCESS); + ASSERT_EQ(tag, (void *)0x100); +} + +TEST_F(tagTest, SetAndGetNull) { + umf_result_t ret = umfPoolSetTag(pool.get(), nullptr, nullptr); + ASSERT_EQ(ret, UMF_RESULT_SUCCESS); + + void *tag; + ret = umfPoolGetTag(pool.get(), &tag); + ASSERT_EQ(ret, UMF_RESULT_SUCCESS); + ASSERT_EQ(tag, nullptr); +} + +TEST_F(tagTest, NoSetAndGet) { + void *tag; + umf_result_t ret = umfPoolGetTag(pool.get(), &tag); + ASSERT_EQ(ret, UMF_RESULT_SUCCESS); + ASSERT_EQ(tag, nullptr); +} + +TEST_F(tagTest, SetAndGetMt) { + static constexpr size_t NUM_THREADS = 8; + static constexpr size_t NUM_OPS_PER_THREAD = 16; + + std::vector threads; + + auto encodeTag = [](size_t thread, size_t op) -> void * { + return reinterpret_cast(thread * NUM_OPS_PER_THREAD + op); + }; + + auto decodeTag = [](void *tag) -> std::pair { + auto op = reinterpret_cast(tag) & (NUM_OPS_PER_THREAD - 1); + auto thread = reinterpret_cast(tag) / NUM_OPS_PER_THREAD; + return {thread, op}; + }; + + for (size_t i = 0; i < NUM_THREADS; i++) { + threads.emplace_back([this, i, encodeTag, decodeTag] { + for (size_t j = 0; j < NUM_OPS_PER_THREAD; j++) { + void *oldTag; + umf_result_t ret = + umfPoolSetTag(pool.get(), encodeTag(i, j), &oldTag); + ASSERT_EQ(ret, UMF_RESULT_SUCCESS); + + void *queriedTag; + ret = umfPoolGetTag(pool.get(), &queriedTag); + ASSERT_EQ(ret, UMF_RESULT_SUCCESS); + + auto [t1, op1] = decodeTag(oldTag); + auto [t2, op2] = decodeTag(queriedTag); + // if the tag was set by the same thread, the op part should be the same or higher + ASSERT_TRUE(t1 != t2 || op2 >= op1); + } + }); + } + + for (auto &thread : threads) { + thread.join(); + } + + void *tag; + auto ret = umfPoolGetTag(pool.get(), &tag); + ASSERT_EQ(ret, UMF_RESULT_SUCCESS); + + auto [t, op] = decodeTag(tag); + ASSERT_TRUE(t < NUM_THREADS); + ASSERT_TRUE(op == NUM_OPS_PER_THREAD - 1); +} + +TEST_F(tagTest, SetAndGetInvalidPtr) { + umf_result_t ret = umfPoolSetTag(pool.get(), nullptr, nullptr); + ASSERT_EQ(ret, UMF_RESULT_SUCCESS); + + ret = umfPoolGetTag(pool.get(), nullptr); + ASSERT_EQ(ret, UMF_RESULT_ERROR_INVALID_ARGUMENT); +} + +TEST_F(tagTest, SetAndGetInvalidPool) { + umf_result_t ret = + umfPoolSetTag(nullptr, reinterpret_cast(0x1), nullptr); + ASSERT_EQ(ret, UMF_RESULT_ERROR_INVALID_ARGUMENT); + + void *tag; + ret = umfPoolGetTag(nullptr, &tag); + ASSERT_EQ(ret, UMF_RESULT_ERROR_INVALID_ARGUMENT); +} + INSTANTIATE_TEST_SUITE_P( mallocPoolTest, umfPoolTest, ::testing::Values(poolCreateExtParams{&MALLOC_POOL_OPS, nullptr,