Add callback for freeing a leaf

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

Author: ldorau

src/critnib/critnib.c

@@ -120,6 +120,8 @@ struct critnib_leaf {
 
 struct critnib {
     struct critnib_node *root;
+    free_leaf_t cb_free_leaf; // callback for freeing a leaf
+    void *leaf_allocator;     // handle of allocator for leaves
 
     /* pool of freed nodes: singly linked list, next at child[0] */
     struct critnib_node *deleted_node;
@@ -162,7 +164,7 @@ static inline unsigned slice_index(word key, sh_t shift) {
 /*
  * critnib_new -- allocates a new critnib structure
  */
-struct critnib *critnib_new(void) {
+struct critnib *critnib_new(void *leaf_allocator, free_leaf_t cb_free_leaf) {
     struct critnib *c = umf_ba_global_alloc(sizeof(struct critnib));
     if (!c) {
         return NULL;
@@ -175,6 +177,8 @@ struct critnib *critnib_new(void) {
         goto err_free_critnib;
     }
 
+    c->leaf_allocator = leaf_allocator;
+    c->cb_free_leaf = cb_free_leaf;
     utils_annotate_memory_no_check(&c->root, sizeof(c->root));
     utils_annotate_memory_no_check(&c->remove_count, sizeof(c->remove_count));
 
@@ -189,6 +193,9 @@ struct critnib *critnib_new(void) {
  */
 static void delete_node(struct critnib *c, struct critnib_node *__restrict n) {
     if (is_leaf(n)) {
+        if (c->cb_free_leaf && to_leaf(n)) {
+            c->cb_free_leaf(c->leaf_allocator, (void *)to_leaf(n)->value);
+        }
         umf_ba_global_free(to_leaf(n));
     } else {
         for (int i = 0; i < SLNODES; i++) {
@@ -225,6 +232,10 @@ void critnib_delete(struct critnib *c) {
 
     for (int i = 0; i < DELETED_LIFE; i++) {
         umf_ba_global_free(c->pending_del_nodes[i]);
+        if (c->cb_free_leaf && c->pending_del_leaves[i]) {
+            c->cb_free_leaf(c->leaf_allocator,
+                            (void *)c->pending_del_leaves[i]->value);
+        }
         umf_ba_global_free(c->pending_del_leaves[i]);
     }
 
@@ -277,6 +288,10 @@ static void free_leaf(struct critnib *__restrict c,
         return;
     }
 
+    if (c->cb_free_leaf && k) {
+        c->cb_free_leaf(c->leaf_allocator, (void *)k->value);
+    }
+
     utils_atomic_store_release_ptr((void **)&k->value, c->deleted_leaf);
     utils_atomic_store_release_ptr((void **)&c->deleted_leaf, k);
 }

src/critnib/critnib.h

@@ -1,6 +1,6 @@
 /*
  *
- * Copyright (C) 2023-2024 Intel Corporation
+ * Copyright (C) 2023-2025 Intel Corporation
  *
  * Under the Apache License v2.0 with LLVM Exceptions. See LICENSE.TXT.
  * SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
@@ -18,6 +18,7 @@ extern "C" {
 
 struct critnib;
 typedef struct critnib critnib;
+typedef void (*free_leaf_t)(void *leaf_allocator, void *ptr);
 
 enum find_dir_t {
     FIND_L = -2,
@@ -27,7 +28,7 @@ enum find_dir_t {
     FIND_G = +2,
 };
 
-critnib *critnib_new(void);
+critnib *critnib_new(void *leaf_allocator, free_leaf_t cb_free_leaf);
 void critnib_delete(critnib *c);
 
 int critnib_insert(critnib *c, uintptr_t key, void *value, int update);

src/pool/pool_disjoint.c

@@ -597,7 +597,7 @@ umf_result_t disjoint_pool_initialize(umf_memory_provider_handle_t provider,
     disjoint_pool->provider = provider;
     disjoint_pool->params = *dp_params;
 
-    disjoint_pool->known_slabs = critnib_new();
+    disjoint_pool->known_slabs = critnib_new(NULL, NULL);
     if (disjoint_pool->known_slabs == NULL) {
         goto err_free_disjoint_pool;
     }

src/provider/provider_file_memory.c

@@ -274,14 +274,14 @@ static umf_result_t file_initialize(const void *params, void **provider) {
         goto err_coarse_delete;
     }
 
-    file_provider->fd_offset_map = critnib_new();
+    file_provider->fd_offset_map = critnib_new(NULL, NULL);
     if (!file_provider->fd_offset_map) {
         LOG_ERR("creating the map of file descriptor offsets failed");
         ret = UMF_RESULT_ERROR_OUT_OF_HOST_MEMORY;
         goto err_mutex_destroy_not_free;
     }
 
-    file_provider->mmaps = critnib_new();
+    file_provider->mmaps = critnib_new(NULL, NULL);
     if (!file_provider->mmaps) {
         LOG_ERR("creating the map of memory mappings failed");
         ret = UMF_RESULT_ERROR_OUT_OF_HOST_MEMORY;

src/provider/provider_os_memory.c

@@ -647,7 +647,7 @@ static umf_result_t os_initialize(const void *params, void **provider) {
         goto err_destroy_hwloc_topology;
     }
 
-    os_provider->fd_offset_map = critnib_new();
+    os_provider->fd_offset_map = critnib_new(NULL, NULL);
     if (!os_provider->fd_offset_map) {
         LOG_ERR("creating file descriptor offset map failed");
         ret = UMF_RESULT_ERROR_OUT_OF_HOST_MEMORY;

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,14 +267,17 @@ static umf_result_t umfMemoryTrackerRemove(umf_memory_tracker_handle_t hTracker,
         return UMF_RESULT_ERROR_UNKNOWN;
     }
 
-    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);
+
+    assert(level < MAX_LEVELS_OF_ALLOC_SEGMENT_MAP);
+    value = critnib_remove(hTracker->alloc_segments_map[level], (uintptr_t)ptr);
+    assert(value);
+
     if (parent_value) {
         LOG_DEBUG(
             "child #%zu removed from memory region: tracker=%p, level=%i, "
@@ -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;
 }
 
@@ -652,6 +666,9 @@ static umf_result_t trackingAllocationMerge(void *hProvider, void *lowPtr,
     // we only need to update the size of the first part
     utils_atomic_store_release_u64((uint64_t *)&lowValue->size, totalSize);
 
+    // size == 0 means that the entry was removed
+    utils_atomic_store_release_u64((uint64_t *)&highValue->size, 0);
+
     void *erasedhighValue = critnib_remove(
         provider->hTracker->alloc_segments_map[highLevel], (uintptr_t)highPtr);
     assert(erasedhighValue == highValue);
@@ -664,8 +681,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:
@@ -1128,7 +1143,7 @@ umf_result_t umfTrackingMemoryProviderCreate(
         return UMF_RESULT_ERROR_UNKNOWN;
     }
     params.pool = hPool;
-    params.ipcCache = critnib_new();
+    params.ipcCache = critnib_new(NULL, NULL);
     if (!params.ipcCache) {
         LOG_ERR("failed to create IPC cache");
         return UMF_RESULT_ERROR_OUT_OF_HOST_MEMORY;
@@ -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();
+        handle->alloc_segments_map[i] =
+            critnib_new(alloc_info_allocator, free_leaf);
         if (!handle->alloc_segments_map[i]) {
             goto err_destroy_alloc_segments_map;
         }
@@ -1192,7 +1214,7 @@ umf_memory_tracker_handle_t umfMemoryTrackerCreate(void) {
         goto err_destroy_alloc_segments_map;
     }
 
-    handle->ipc_segments_map = critnib_new();
+    handle->ipc_segments_map = critnib_new(NULL, NULL);
     if (!handle->ipc_segments_map) {
         goto err_destroy_ipc_info_allocator;
     }