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: #1233

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

Author: ldorau

src/provider/provider_tracking.c

@@ -271,8 +271,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;
 }
 
@@ -664,8 +662,6 @@ static umf_result_t trackingAllocationMerge(void *hProvider, void *lowPtr,
               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 +1152,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 +1182,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;
         }