Merge pull request #380 from igchor/fix_rec_tracking2

Add runtime option to disable pool tracking per-pool and disable pool tracking for Proxy_pool

Author: ldorau

include/umf/memory_pool.h

@@ -35,6 +35,8 @@ typedef enum umf_pool_create_flag_t {
     UMF_POOL_CREATE_FLAG_OWN_PROVIDER =
         (1
          << 0), ///< Pool will own the specified provider and destroy it in umfPoolDestroy
+    UMF_POOL_CREATE_FLAG_DISABLE_TRACKING =
+        (1 << 1), ///< Pool will not track memory allocations
     /// @cond
     UMF_POOL_CREATE_FLAG_FORCE_UINT32 = 0x7fffffff
     /// @endcond

src/CMakeLists.txt

@@ -118,16 +118,15 @@ else()
         LIBS ${UMF_LIBS})
 endif()
 
+if(UMF_ENABLE_POOL_TRACKING)
+    set(UMF_PUBLIC_COMPILE_DEFINITIONS ${UMF_PUBLIC_COMPILE_DEFINITIONS}
+                                       "UMF_ENABLE_POOL_TRACKING")
+endif()
+
 target_link_directories(umf PRIVATE ${UMF_PRIVATE_LIBRARY_DIRS})
 
 target_compile_definitions(umf PUBLIC ${UMF_PUBLIC_COMPILE_DEFINITIONS})
 
-if(UMF_ENABLE_POOL_TRACKING)
-    target_sources(umf PRIVATE memory_pool_tracking.c)
-else()
-    target_sources(umf PRIVATE memory_pool_default.c)
-endif()
-
 if(UMF_BUILD_LEVEL_ZERO_PROVIDER)
     target_sources(umf PRIVATE provider/provider_level_zero.c)
 

src/memory_pool.c

@@ -17,17 +17,125 @@
 #include <assert.h>
 #include <stdlib.h>
 
+#include "base_alloc_global.h"
+#include "memory_pool_internal.h"
+#include "memory_provider_internal.h"
+#include "provider_tracking.h"
+
+static umf_result_t umfPoolCreateInternal(const umf_memory_pool_ops_t *ops,
+                                          umf_memory_provider_handle_t provider,
+                                          void *params,
+                                          umf_pool_create_flags_t flags,
+                                          umf_memory_pool_handle_t *hPool) {
+    if (!ops || !provider || !hPool) {
+        return UMF_RESULT_ERROR_INVALID_ARGUMENT;
+    }
+
+    umf_result_t ret = UMF_RESULT_SUCCESS;
+    umf_memory_pool_handle_t pool =
+        umf_ba_global_alloc(sizeof(umf_memory_pool_t));
+    if (!pool) {
+        return UMF_RESULT_ERROR_OUT_OF_HOST_MEMORY;
+    }
+
+    assert(ops->version == UMF_VERSION_CURRENT);
+
+    if (!(flags & UMF_POOL_CREATE_FLAG_DISABLE_TRACKING)) {
+        // wrap provider with memory tracking provider
+        ret = umfTrackingMemoryProviderCreate(provider, pool, &pool->provider);
+        if (ret != UMF_RESULT_SUCCESS) {
+            goto err_provider_create;
+        }
+    } else {
+        pool->provider = provider;
+    }
+
+    pool->flags = flags;
+    pool->ops = *ops;
+
+    ret = ops->initialize(pool->provider, params, &pool->pool_priv);
+    if (ret != UMF_RESULT_SUCCESS) {
+        goto err_pool_init;
+    }
+
+    *hPool = pool;
+    return UMF_RESULT_SUCCESS;
+
+err_pool_init:
+    if (!(flags & UMF_POOL_CREATE_FLAG_DISABLE_TRACKING)) {
+        umfMemoryProviderDestroy(pool->provider);
+    }
+err_provider_create:
+    umf_ba_global_free(pool);
+    return ret;
+}
+
+void umfPoolDestroy(umf_memory_pool_handle_t hPool) {
+    hPool->ops.finalize(hPool->pool_priv);
+    if (hPool->flags & UMF_POOL_CREATE_FLAG_OWN_PROVIDER) {
+        // Destroy associated memory provider.
+        umf_memory_provider_handle_t hProvider = NULL;
+        umfPoolGetMemoryProvider(hPool, &hProvider);
+        umfMemoryProviderDestroy(hProvider);
+    }
+
+    if (!(hPool->flags & UMF_POOL_CREATE_FLAG_DISABLE_TRACKING)) {
+        // Destroy tracking provider.
+        umfMemoryProviderDestroy(hPool->provider);
+    }
+    // TODO: this free keeps memory in base allocator, so it can lead to OOM in some scenarios (it should be optimized)
+    umf_ba_global_free(hPool);
+}
+
+umf_result_t umfFree(void *ptr) {
+#ifndef UMF_ENABLE_POOL_TRACKING
+    return UMF_RESULT_ERROR_NOT_SUPPORTED;
+#endif
+
+    umf_memory_pool_handle_t hPool = umfPoolByPtr(ptr);
+    if (hPool) {
+        return umfPoolFree(hPool, ptr);
+    }
+    return UMF_RESULT_SUCCESS;
+}
+
+umf_memory_pool_handle_t umfPoolByPtr(const void *ptr) {
+    return umfMemoryTrackerGetPool(ptr);
+}
+
+umf_result_t umfPoolGetMemoryProvider(umf_memory_pool_handle_t hPool,
+                                      umf_memory_provider_handle_t *hProvider) {
+    if (!hProvider) {
+        return UMF_RESULT_ERROR_INVALID_ARGUMENT;
+    }
+
+    if (hPool->flags & UMF_POOL_CREATE_FLAG_DISABLE_TRACKING) {
+        *hProvider = hPool->provider;
+    } else {
+        umfTrackingMemoryProviderGetUpstreamProvider(
+            umfMemoryProviderGetPriv(hPool->provider), hProvider);
+    }
+
+    return UMF_RESULT_SUCCESS;
+}
+
 umf_result_t umfPoolCreate(const umf_memory_pool_ops_t *ops,
                            umf_memory_provider_handle_t provider, void *params,
                            umf_pool_create_flags_t flags,
                            umf_memory_pool_handle_t *hPool) {
     libumfInit();
-    umf_result_t ret = umfPoolCreateInternal(ops, provider, params, hPool);
+
+#ifndef UMF_ENABLE_POOL_TRACKING
+    // if tracking is not enabled during compilation, disable it for each pool here as well
+    flags |= UMF_POOL_CREATE_FLAG_DISABLE_TRACKING;
+#endif
+
+    umf_result_t ret =
+        umfPoolCreateInternal(ops, provider, params, flags, hPool);
     if (ret != UMF_RESULT_SUCCESS) {
         return ret;
     }
     assert(*hPool != NULL);
-    (*hPool)->own_provider = (flags & UMF_POOL_CREATE_FLAG_OWN_PROVIDER);
 
     return UMF_RESULT_SUCCESS;
 }

src/memory_pool_default.c

@@ -26,18 +26,12 @@ extern "C" {
 typedef struct umf_memory_pool_t {
     void *pool_priv;
     umf_memory_pool_ops_t ops;
+    umf_pool_create_flags_t flags;
 
     // Memory provider used by the pool.
     umf_memory_provider_handle_t provider;
-    // Tells whether memory provider is owned by the pool.
-    bool own_provider;
 } umf_memory_pool_t;
 
-umf_result_t umfPoolCreateInternal(const umf_memory_pool_ops_t *ops,
-                                   umf_memory_provider_handle_t provider,
-                                   void *params,
-                                   umf_memory_pool_handle_t *hPool);
-
 #ifdef __cplusplus
 }
 #endif

src/memory_pool_internal.h

@@ -107,8 +107,9 @@ void proxy_lib_create_common(void) {
         exit(-1);
     }
 
-    umf_result = umfPoolCreate(umfPoolManagerOps(), OS_memory_provider, NULL, 0,
-                               &Proxy_pool);
+    umf_result =
+        umfPoolCreate(umfPoolManagerOps(), OS_memory_provider, NULL,
+                      UMF_POOL_CREATE_FLAG_DISABLE_TRACKING, &Proxy_pool);
     if (umf_result != UMF_RESULT_SUCCESS) {
         fprintf(stderr, "error: creating UMF pool manager failed\n");
         exit(-1);