Merge pull request #457 from vinser52/svinogra_ipc_fix

Fix umfMemoryTrackerGetPool function and minor refactoring

Author: lplewa

src/ipc.c

@@ -47,16 +47,16 @@ umf_result_t umfGetIPCHandle(const void *ptr, umf_ipc_handle_t *umfIPCHandle,
         return UMF_RESULT_ERROR_OUT_OF_HOST_MEMORY;
     }
 
-    ret =
-        umfMemoryProviderGetIPCHandle(provider, allocInfo.base, allocInfo.size,
-                                      (void *)ipcData->providerIpcData);
+    ret = umfMemoryProviderGetIPCHandle(provider, allocInfo.base,
+                                        allocInfo.baseSize,
+                                        (void *)ipcData->providerIpcData);
     if (ret != UMF_RESULT_SUCCESS) {
         LOG_ERR("umfGetIPCHandle: failed to get IPC handle.");
         umf_ba_global_free(ipcData);
         return ret;
     }
 
-    ipcData->size = allocInfo.size;
+    ipcData->baseSize = allocInfo.baseSize;
     ipcData->offset = (uintptr_t)ptr - (uintptr_t)allocInfo.base;
 
     *umfIPCHandle = ipcData;
@@ -114,5 +114,5 @@ umf_result_t umfCloseIPCHandle(void *ptr) {
     umf_memory_provider_handle_t hProvider = allocInfo.pool->provider;
 
     return umfMemoryProviderCloseIPCHandle(hProvider, allocInfo.base,
-                                           allocInfo.size);
+                                           allocInfo.baseSize);
 }

src/ipc_internal.h

@@ -21,7 +21,7 @@ extern "C" {
 // providerIpcData is a Flexible Array Member because its size varies
 // depending on the provider.
 typedef struct umf_ipc_data_t {
-    size_t size; // size of base allocation
+    size_t baseSize; // size of base (coarse-grain) allocation
     uint64_t offset;
     char providerIpcData[];
 } umf_ipc_data_t;

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,
@@ -117,7 +113,7 @@ umf_result_t umfMemoryTrackerGetAllocInfo(const void *ptr,
     }
 
     pAllocInfo->base = (void *)rkey;
-    pAllocInfo->size = rvalue->size;
+    pAllocInfo->baseSize = rvalue->size;
     pAllocInfo->pool = rvalue->pool;
 
     return UMF_RESULT_SUCCESS;
@@ -577,7 +573,7 @@ static size_t getDataSizeFromIpcHandle(const void *providerIpcData) {
     // the Flexible Array Member of umf_ipc_data_t.
     umf_ipc_data_t *ipcUmfData =
         (umf_ipc_data_t *)((uint8_t *)providerIpcData - sizeof(umf_ipc_data_t));
-    return ipcUmfData->size;
+    return ipcUmfData->baseSize;
 }
 
 static umf_result_t trackingOpenIpcHandle(void *provider, void *providerIpcData,

src/provider/provider_tracking.h

@@ -42,7 +42,7 @@ umf_memory_pool_handle_t umfMemoryTrackerGetPool(const void *ptr);
 
 typedef struct umf_alloc_info_t {
     void *base;
-    size_t size;
+    size_t baseSize;
     umf_memory_pool_handle_t pool;
 } umf_alloc_info_t;
 

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,