Postpone freeing a tracker entry until it is really removed from tracker

Postpone freeing a tracker entry until it is really removed from tracker.

Ref: oneapi-src#1233

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

Author: ldorau

src/provider/provider_tracking.c

@@ -201,6 +201,19 @@ static umf_result_t umfMemoryTrackerAdd(umf_memory_tracker_handle_t hTracker,
 
         utils_atomic_load_acquire_u64((uint64_t *)&rvalue->size, &rsize);
 
+        // size == 0 means that the entry was removed
+        if (rsize == 0) {
+            // restart the search
+            parent_value = NULL;
+            rvalue = NULL;
+            parent_key = 0;
+            rkey = 0;
+            rsize = 0;
+            level = 0;
+            found = 0;
+            continue;
+        }
+
         if ((uintptr_t)ptr < rkey + rsize) {
             if (level == MAX_LEVELS_OF_ALLOC_SEGMENT_MAP - 1) {
                 // TODO: we need to support an arbitrary amount of layers in the future
@@ -254,13 +267,16 @@ static umf_result_t umfMemoryTrackerRemove(umf_memory_tracker_handle_t hTracker,
         return UMF_RESULT_ERROR_UNKNOWN;
     }
 
+    LOG_DEBUG("memory region removed: tracker=%p, level=%i, pool=%p, ptr=%p, "
+              "size=%zu",
+              (void *)hTracker, level, (void *)value->pool, ptr, value->size);
+
     assert(level < MAX_LEVELS_OF_ALLOC_SEGMENT_MAP);
     value = critnib_remove(hTracker->alloc_segments_map[level], (uintptr_t)ptr);
     assert(value);
 
-    LOG_DEBUG("memory region removed: tracker=%p, level=%i, pool=%p, ptr=%p, "
-              "size=%zu",
-              (void *)hTracker, level, (void *)value->pool, ptr, value->size);
+    // size == 0 means that the entry was removed
+    utils_atomic_store_release_u64((uint64_t *)&value->size, 0);
 
     if (parent_value) {
         LOG_DEBUG(
@@ -271,8 +287,6 @@ static umf_result_t umfMemoryTrackerRemove(umf_memory_tracker_handle_t hTracker,
         parent_value->n_children--;
     }
 
-    umf_ba_free(hTracker->alloc_info_allocator, value);
-
     return UMF_RESULT_SUCCESS;
 }
 
@@ -657,15 +671,16 @@ static umf_result_t trackingAllocationMerge(void *hProvider, void *lowPtr,
     assert(erasedhighValue == highValue);
     (void)erasedhighValue; // unused in the Release build
 
+    // size == 0 means that the entry was removed
+    utils_atomic_store_release_u64((uint64_t *)&highValue->size, 0);
+
     utils_mutex_unlock(&provider->hTracker->splitMergeMutex);
 
     LOG_DEBUG("merged memory regions (level=%i): lowPtr=%p (child=%zu), "
               "highPtr=%p (child=%zu), totalSize=%zu",
               lowLevel, lowPtr, lowValue->n_children, highPtr,
               highValue->n_children, totalSize);
 
-    umf_ba_free(provider->hTracker->alloc_info_allocator, highValue);
-
     return UMF_RESULT_SUCCESS;
 
 err_fatal:
@@ -1156,6 +1171,12 @@ void umfTrackingMemoryProviderGetUpstreamProvider(
     *hUpstream = p->hUpstream;
 }
 
+static void free_leaf(void *leaf_allocator, void *ptr) {
+    if (ptr) {
+        umf_ba_free(leaf_allocator, ptr);
+    }
+}
+
 umf_memory_tracker_handle_t umfMemoryTrackerCreate(void) {
     umf_memory_tracker_handle_t handle =
         umf_ba_global_alloc(sizeof(struct umf_memory_tracker_t));
@@ -1180,7 +1201,8 @@ umf_memory_tracker_handle_t umfMemoryTrackerCreate(void) {
 
     int i;
     for (i = 0; i < MAX_LEVELS_OF_ALLOC_SEGMENT_MAP; i++) {
-        handle->alloc_segments_map[i] = critnib_new(NULL, NULL);
+        handle->alloc_segments_map[i] =
+            critnib_new(alloc_info_allocator, free_leaf);
         if (!handle->alloc_segments_map[i]) {
             goto err_destroy_alloc_segments_map;
         }