Implement umfPool[Set/Get]Tag

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

src/libumf.def

@@ -104,10 +104,12 @@ EXPORTS
     umfPoolFree
     umfPoolGetIPCHandleSize
     umfPoolGetLastAllocationError
+    umfPoolGetTag
     umfPoolGetMemoryProvider
     umfPoolMalloc
     umfPoolMallocUsableSize
     umfPoolRealloc
+    umfPoolSetTag
     umfProxyPoolOps
     umfPutIPCHandle
     umfScalablePoolOps

src/libumf.map

@@ -99,9 +99,11 @@ UMF_1.0 {
         umfPoolGetIPCHandleSize;
         umfPoolGetLastAllocationError;
         umfPoolGetMemoryProvider;
+        umfPoolGetTag;
         umfPoolMalloc;
         umfPoolMallocUsableSize;
         umfPoolRealloc;
+        umfPoolSetTag;
         umfProxyPoolOps;
         umfPutIPCHandle;
         umfScalablePoolOps;

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;
+}

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

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<std::thread> threads;
+
+    auto encodeTag = [](size_t thread, size_t op) -> void * {
+        return reinterpret_cast<void *>(thread * NUM_OPS_PER_THREAD + op);
+    };
+
+    auto decodeTag = [](void *tag) -> std::pair<size_t, size_t> {
+        auto op = reinterpret_cast<size_t>(tag) & (NUM_OPS_PER_THREAD - 1);
+        auto thread = reinterpret_cast<size_t>(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<void *>(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,