Clear tracker for the pool if provider does not support free() operation

Clear the tracker from the current pool pointers,
if the provider does not support the free() operation

Signed-off-by: Lukasz Dorau <[email protected]>

Author: ldorau

src/memory_pool.c

@@ -41,8 +41,12 @@ static umf_result_t umfPoolCreateInternal(const umf_memory_pool_ops_t *ops,
     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);
+        // Wrap provider with memory tracking provider.
+        // Check if the provider does not support the free() operation.
+        bool upstreamDoesNotFree = (umfMemoryProviderFree(provider, NULL, 0) ==
+                                    UMF_RESULT_ERROR_NOT_SUPPORTED);
+        ret = umfTrackingMemoryProviderCreate(provider, pool, &pool->provider,
+                                              upstreamDoesNotFree);
         if (ret != UMF_RESULT_SUCCESS) {
             goto err_provider_create;
         }

src/provider/provider_tracking.c

@@ -141,6 +141,9 @@ typedef struct umf_tracking_memory_provider_t {
     umf_memory_tracker_handle_t hTracker;
     umf_memory_pool_handle_t pool;
     critnib *ipcCache;
+
+    // the upstream provider does not support the free() operation
+    bool upstreamDoesNotFree;
 } umf_tracking_memory_provider_t;
 
 typedef struct umf_tracking_memory_provider_t umf_tracking_memory_provider_t;
@@ -429,10 +432,39 @@ static void check_if_tracker_is_empty(umf_memory_tracker_handle_t hTracker,
 }
 #endif /* NDEBUG */
 
+static void clear_tracker_for_the_pool(umf_memory_tracker_handle_t hTracker,
+                                       umf_memory_pool_handle_t pool) {
+    uintptr_t rkey;
+    void *rvalue;
+    uintptr_t last_key = 0;
+
+    while (1 == critnib_find((critnib *)hTracker->map, last_key, FIND_G, &rkey,
+                             &rvalue)) {
+        tracker_value_t *value = (tracker_value_t *)rvalue;
+        if (value->pool != pool) {
+            last_key = rkey;
+            continue;
+        }
+
+        void *removed_value = critnib_remove(hTracker->map, rkey);
+        assert(removed_value == rvalue);
+
+        umf_ba_free(hTracker->tracker_allocator, removed_value);
+
+        last_key = rkey;
+    }
+}
+
 static void trackingFinalize(void *provider) {
     umf_tracking_memory_provider_t *p =
         (umf_tracking_memory_provider_t *)provider;
+
     critnib_delete(p->ipcCache);
+
+    if (p->upstreamDoesNotFree) {
+        clear_tracker_for_the_pool(p->hTracker, p->pool);
+    }
+
 #ifndef NDEBUG
     check_if_tracker_is_empty(p->hTracker, p->pool);
 #endif /* NDEBUG */
@@ -661,10 +693,11 @@ umf_memory_provider_ops_t UMF_TRACKING_MEMORY_PROVIDER_OPS = {
 
 umf_result_t umfTrackingMemoryProviderCreate(
     umf_memory_provider_handle_t hUpstream, umf_memory_pool_handle_t hPool,
-    umf_memory_provider_handle_t *hTrackingProvider) {
+    umf_memory_provider_handle_t *hTrackingProvider, bool upstreamDoesNotFree) {
 
     umf_tracking_memory_provider_t params;
     params.hUpstream = hUpstream;
+    params.upstreamDoesNotFree = upstreamDoesNotFree;
     params.hTracker = TRACKER;
     if (!params.hTracker) {
         LOG_ERR("failed, TRACKER is NULL");

src/provider/provider_tracking.h

@@ -11,6 +11,7 @@
 #define UMF_MEMORY_TRACKER_INTERNAL_H 1
 
 #include <assert.h>
+#include <stdbool.h>
 #include <stdlib.h>
 
 #include <umf/base.h>
@@ -53,7 +54,7 @@ umf_result_t umfMemoryTrackerGetAllocInfo(const void *ptr,
 // forwards all requests to hUpstream memory Provider. hUpstream lifetime should be managed by the user of this function.
 umf_result_t umfTrackingMemoryProviderCreate(
     umf_memory_provider_handle_t hUpstream, umf_memory_pool_handle_t hPool,
-    umf_memory_provider_handle_t *hTrackingProvider);
+    umf_memory_provider_handle_t *hTrackingProvider, bool upstreamDoesNotFree);
 
 void umfTrackingMemoryProviderGetUpstreamProvider(
     umf_memory_provider_handle_t hTrackingProvider,

test/memoryPoolAPI.cpp

@@ -139,7 +139,8 @@ TEST_P(umfPoolWithCreateFlagsTest, memoryPoolWithCustomProvider) {
 }
 
 TEST_F(test, retrieveMemoryProvider) {
-    umf_memory_provider_handle_t provider = (umf_memory_provider_handle_t)0x1;
+    auto nullProvider = umf_test::wrapProviderUnique(nullProviderCreate());
+    umf_memory_provider_handle_t provider = nullProvider.get();
 
     auto pool =
         wrapPoolUnique(createPoolChecked(umfProxyPoolOps(), provider, nullptr));
@@ -258,7 +259,8 @@ TEST_P(poolInitializeTest, errorPropagation) {
 }
 
 TEST_F(test, retrieveMemoryProvidersError) {
-    umf_memory_provider_handle_t provider = (umf_memory_provider_handle_t)0x1;
+    auto nullProvider = umf_test::wrapProviderUnique(nullProviderCreate());
+    umf_memory_provider_handle_t provider = nullProvider.get();
 
     auto pool =
         wrapPoolUnique(createPoolChecked(umfProxyPoolOps(), provider, nullptr));

test/pools/disjoint_pool.cpp

@@ -82,7 +82,10 @@ TEST_F(test, sharedLimits) {
         }
         umf_result_t free(void *ptr, [[maybe_unused]] size_t size) noexcept {
             ::free(ptr);
-            numFrees++;
+            // umfMemoryProviderFree(provider, NULL, 0) is called inside umfPoolCreateInternal()
+            if (ptr != NULL && size != 0) {
+                numFrees++;
+            }
             return UMF_RESULT_SUCCESS;
         }
     };