Enable the lock in umfMemoryTrackerAdd/Remove only if there is level #1

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

Author: ldorau

src/provider/provider_tracking.c

@@ -38,6 +38,8 @@ struct umf_memory_tracker_t {
     // for another memory pool (nested memory pooling).
     critnib *alloc_segments_map[MAX_LEVELS_OF_ALLOC_SEGMENT_MAP];
     utils_mutex_t mutex;
+    // number of memory regions at levels > 0
+    uint64_t n_at_higher_levels;
 };
 
 typedef struct tracker_alloc_info_t {
@@ -131,9 +133,14 @@ umfMemoryTrackerAddAtLevel(umf_memory_tracker_handle_t hTracker, int level,
     int ret = critnib_insert(hTracker->alloc_segments_map[level],
                              (uintptr_t)ptr, value, 0);
     if (ret == 0) {
+        if (level > 0) {
+            utils_fetch_and_add64(&hTracker->n_at_higher_levels, 1);
+        }
+
         LOG_DEBUG("memory region is added, tracker=%p, level=%i, pool=%p, "
                   "ptr=%p, size=%zu",
                   (void *)hTracker, level, (void *)pool, ptr, size);
+
         if (parent_value) {
             parent_value->n_children++;
             LOG_DEBUG(
@@ -158,9 +165,10 @@ umfMemoryTrackerAddAtLevel(umf_memory_tracker_handle_t hTracker, int level,
     return umf_result;
 }
 
-static umf_result_t umfMemoryTrackerAdd(umf_memory_tracker_handle_t hTracker,
-                                        umf_memory_pool_handle_t pool,
-                                        const void *ptr, size_t size) {
+static umf_result_t
+umfMemoryTrackerAddLock(umf_memory_tracker_handle_t hTracker,
+                        umf_memory_pool_handle_t pool, const void *ptr,
+                        size_t size, int lock) {
     assert(ptr);
 
     umf_result_t umf_result = UMF_RESULT_ERROR_UNKNOWN;
@@ -170,10 +178,13 @@ static umf_result_t umfMemoryTrackerAdd(umf_memory_tracker_handle_t hTracker,
     uintptr_t rkey = 0;
     int level = 0;
     int found = 0;
+    int ret;
 
-    int ret = utils_mutex_lock(&hTracker->mutex);
-    if (ret) {
-        return UMF_RESULT_ERROR_UNKNOWN;
+    if (lock) {
+        ret = utils_mutex_lock(&hTracker->mutex);
+        if (ret) {
+            return UMF_RESULT_ERROR_UNKNOWN;
+        }
     }
 
     // Find the most nested (in the highest level) entry
@@ -217,13 +228,25 @@ static umf_result_t umfMemoryTrackerAdd(umf_memory_tracker_handle_t hTracker,
     umf_result = UMF_RESULT_SUCCESS;
 
 err_unlock:
-    utils_mutex_unlock(&hTracker->mutex);
+    if (lock) {
+        utils_mutex_unlock(&hTracker->mutex);
+    }
 
     return umf_result;
 }
 
-static umf_result_t umfMemoryTrackerRemove(umf_memory_tracker_handle_t hTracker,
-                                           const void *ptr) {
+static umf_result_t umfMemoryTrackerAdd(umf_memory_tracker_handle_t hTracker,
+                                        umf_memory_pool_handle_t pool,
+                                        const void *ptr, size_t size) {
+
+    return umfMemoryTrackerAddLock(
+        hTracker, pool, ptr, size,
+        (utils_fetch_and_add64(&hTracker->n_at_higher_levels, 0) > 0));
+}
+
+static umf_result_t
+umfMemoryTrackerRemoveLock(umf_memory_tracker_handle_t hTracker,
+                           const void *ptr, int lock) {
     assert(ptr);
 
     // TODO: there is no support for removing partial ranges (or multiple entries
@@ -235,10 +258,13 @@ static umf_result_t umfMemoryTrackerRemove(umf_memory_tracker_handle_t hTracker,
     tracker_alloc_info_t *parent_value = NULL;
     uintptr_t parent_key = 0;
     int level = 0;
+    int ret;
 
-    int ret = utils_mutex_lock(&hTracker->mutex);
-    if (ret) {
-        return UMF_RESULT_ERROR_UNKNOWN;
+    if (lock) {
+        ret = utils_mutex_lock(&hTracker->mutex);
+        if (ret) {
+            return UMF_RESULT_ERROR_UNKNOWN;
+        }
     }
 
     // Find the most nested (on the highest level) entry in the map
@@ -254,6 +280,10 @@ static umf_result_t umfMemoryTrackerRemove(umf_memory_tracker_handle_t hTracker,
     value = critnib_remove(hTracker->alloc_segments_map[level], (uintptr_t)ptr);
     assert(value);
 
+    if (level > 0) {
+        utils_fetch_and_add64(&hTracker->n_at_higher_levels, -1);
+    }
+
     LOG_DEBUG("memory region removed: tracker=%p, level=%i, pool=%p, ptr=%p, "
               "size=%zu",
               (void *)hTracker, level, value->pool, ptr, value->size);
@@ -272,11 +302,20 @@ static umf_result_t umfMemoryTrackerRemove(umf_memory_tracker_handle_t hTracker,
     umf_result = UMF_RESULT_SUCCESS;
 
 err_unlock:
-    utils_mutex_unlock(&hTracker->mutex);
+    if (lock) {
+        utils_mutex_unlock(&hTracker->mutex);
+    }
 
     return umf_result;
 }
 
+static umf_result_t umfMemoryTrackerRemove(umf_memory_tracker_handle_t hTracker,
+                                           const void *ptr) {
+    return umfMemoryTrackerRemoveLock(
+        hTracker, ptr,
+        (utils_fetch_and_add64(&hTracker->n_at_higher_levels, 0) > 0));
+}
+
 umf_memory_pool_handle_t umfMemoryTrackerGetPool(const void *ptr) {
     umf_alloc_info_t allocInfo = {NULL, 0, NULL};
     umf_result_t ret = umfMemoryTrackerGetAllocInfo(ptr, &allocInfo);
@@ -1084,6 +1123,7 @@ umf_memory_tracker_handle_t umfMemoryTrackerCreate(void) {
     }
 
     handle->alloc_info_allocator = alloc_info_allocator;
+    handle->n_at_higher_levels = 0;
 
     void *mutex_ptr = utils_mutex_init(&handle->mutex);
     if (!mutex_ptr) {