critnib: add callback for freeing a leaf

Signed-off-by: Lukasz Dorau <lukasz.dorau@intel.com>

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;
@@ -161,8 +163,12 @@ static inline unsigned slice_index(word key, sh_t shift) {
 
 /*
  * critnib_new -- allocates a new critnib structure
+ *
+ * Arguments:
+ * - cb_free_leaf - callback for freeing a leaf (can be NULL)
+ * - leaf_allocator - handle of allocator for leaves (can be NULL)
  */
-struct critnib *critnib_new(void) {
+struct critnib *critnib_new(free_leaf_t cb_free_leaf, void *leaf_allocator) {
     struct critnib *c = umf_ba_global_alloc(sizeof(struct critnib));
     if (!c) {
         return NULL;
@@ -175,6 +181,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 +197,10 @@ struct critnib *critnib_new(void) {
  */
 static void delete_node(struct critnib *c, struct critnib_node *__restrict n) {
     if (is_leaf(n)) {
+        // call the callback freeing the leaf
+        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 +237,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 +293,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(free_leaf_t cb_free_leaf, void *leaf_allocator);
 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

@@ -1131,7 +1131,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;
@@ -1183,7 +1183,7 @@ 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(NULL, NULL);
         if (!handle->alloc_segments_map[i]) {
             goto err_destroy_alloc_segments_map;
         }
@@ -1195,7 +1195,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;
     }