Fix umfMemoryTrackerGetPool implementation

Author: vinser52

src/provider/provider_tracking.c

@@ -80,34 +80,30 @@ static umf_result_t umfMemoryTrackerRemove(umf_memory_tracker_handle_t hTracker,
 }
 
 umf_memory_pool_handle_t umfMemoryTrackerGetPool(const void *ptr) {
-    assert(ptr);
-
-    if (TRACKER == NULL) {
-        LOG_ERR("tracker is not created");
-        return NULL;
-    }
-
-    if (TRACKER->map == NULL) {
-        LOG_ERR("tracker's map is not created");
-        return NULL;
-    }
-
-    uintptr_t rkey;
-    tracker_value_t *rvalue;
-    int found = critnib_find(TRACKER->map, (uintptr_t)ptr, FIND_LE,
-                             (void *)&rkey, (void **)&rvalue);
-    if (!found) {
+    umf_alloc_info_t allocInfo = {0};
+    umf_result_t ret = umfMemoryTrackerGetAllocInfo(ptr, &allocInfo);
+    if (ret != UMF_RESULT_SUCCESS) {
         return NULL;
     }
 
-    return (rkey + rvalue->size >= (uintptr_t)ptr) ? rvalue->pool : NULL;
+    return allocInfo.pool;
 }
 
 umf_result_t umfMemoryTrackerGetAllocInfo(const void *ptr,
                                           umf_alloc_info_t *pAllocInfo) {
     assert(ptr);
     assert(pAllocInfo);
 
+    if (TRACKER == NULL) {
+        LOG_ERR("tracker is not created");
+        return UMF_RESULT_ERROR_NOT_SUPPORTED;
+    }
+
+    if (TRACKER->map == NULL) {
+        LOG_ERR("tracker's map is not created");
+        return UMF_RESULT_ERROR_NOT_SUPPORTED;
+    }
+
     uintptr_t rkey;
     tracker_value_t *rvalue;
     int found = critnib_find(TRACKER->map, (uintptr_t)ptr, FIND_LE,

test/memoryPoolAPI.cpp

@@ -150,6 +150,35 @@ TEST_F(test, retrieveMemoryProvider) {
     ASSERT_EQ(retProvider, provider);
 }
 
+#ifdef UMF_ENABLE_POOL_TRACKING_TESTS
+TEST_F(test, BasicPoolByPtrTest) {
+    constexpr size_t SIZE = 4096 * 1024;
+
+    umf_memory_provider_handle_t provider;
+    umf_result_t ret =
+        umfMemoryProviderCreate(&MALLOC_PROVIDER_OPS, NULL, &provider);
+    ASSERT_EQ(ret, UMF_RESULT_SUCCESS);
+    auto pool =
+        wrapPoolUnique(createPoolChecked(umfProxyPoolOps(), provider, nullptr,
+                                         UMF_POOL_CREATE_FLAG_OWN_PROVIDER));
+    auto expected_pool = pool.get();
+    char *ptr = (char *)umfPoolMalloc(expected_pool, SIZE);
+    ASSERT_NE(ptr, nullptr);
+
+    auto ret_pool = umfPoolByPtr(ptr);
+    EXPECT_EQ(ret_pool, expected_pool);
+
+    ret_pool = umfPoolByPtr(ptr + SIZE);
+    EXPECT_EQ(ret_pool, nullptr);
+
+    ret_pool = umfPoolByPtr(ptr + SIZE - 1);
+    EXPECT_EQ(ret_pool, expected_pool);
+
+    ret = umfFree(ptr);
+    ASSERT_EQ(ret, UMF_RESULT_SUCCESS);
+}
+#endif /* UMF_ENABLE_POOL_TRACKING_TESTS */
+
 INSTANTIATE_TEST_SUITE_P(
     mallocPoolTest, umfPoolTest,
     ::testing::Values(poolCreateExtParams{&MALLOC_POOL_OPS, nullptr,