[WIP!] critnib: add refcount WIP!

ldorau · ldorau · commit dfc62fd6e1d7 · 2025-05-06T13:13:31.000+02:00
Signed-off-by: Lukasz Dorau &lt;lukasz.dorau@intel.com&gt;
diff --git a/src/critnib/critnib.c b/src/critnib/critnib.c
@@ -116,6 +116,7 @@ struct critnib_node {
 struct critnib_leaf {
     word key;
     void *value;
+    uint64_t ref_count;
 };
 
 struct critnib {
@@ -289,6 +290,11 @@ static void free_leaf(struct critnib *__restrict c,
     }
 
     if (c->cb_free_leaf && k && k->value) {
+        if (k->ref_count > 0) {
+            fprintf(stderr, "free_leaf(): k->ref_count = %llu of k->key = %p\n",
+                    (unsigned long long)k->ref_count, (void *)k->key);
+        }
+        assert(k->ref_count == 0); // TODO: check ref_count
         c->cb_free_leaf(c->leaf_allocator, (void *)k->value);
     }
 
@@ -336,6 +342,8 @@ int critnib_insert(struct critnib *c, word key, void *value, int update) {
 
     utils_atomic_store_release_ptr((void **)&k->key, (void *)key);
     utils_atomic_store_release_ptr((void **)&k->value, value);
+    // mark the leaf as valid (ref_count == 1)
+    utils_atomic_store_release_u64(&k->ref_count, 1ULL);
 
     struct critnib_node *kn = (void *)((word)k | 1);
 
@@ -490,6 +498,11 @@ void *critnib_remove(struct critnib *c, word key) {
     c->pending_del_nodes[del] = n;
 
 del_leaf:
+    uint64_t ref_count = utils_atomic_decrement_u64(&k->ref_count);
+    fprintf(stderr,
+            "critnib_remove(): k->ref_count = %llu of k->key = %p\n",
+            (unsigned long long)ref_count, (void *)k->key);
+
     value = k->value;
     c->pending_del_leaves[del] = k;
 
@@ -498,6 +511,28 @@ void *critnib_remove(struct critnib *c, word key) {
     return value;
 }
 
+/*
+ * critnib_release -- release a reference to a key
+ */
+int critnib_release(struct critnib *c, void *ref) {
+    if (!c || !ref) {
+        return -1;
+    }
+
+    struct critnib_leaf *k = (struct critnib_leaf *)ref;
+
+    // (TODO HERE...)
+
+    /* decrement the reference count */
+    if (utils_atomic_decrement_u64(&k->ref_count) == 0) {
+        fprintf(stderr, "critnib_release(): k->ref_count = %llu of k->key = %p\n",
+                (unsigned long long)k->ref_count, (void *)k->key);
+        free_leaf(c, k);
+    }
+
+    return 0;
+}
+
 /*
  * critnib_get -- query for a key ("==" match), returns value or NULL
  *
@@ -508,13 +543,17 @@ void *critnib_remove(struct critnib *c, word key) {
  * Counterintuitively, it's pointless to return the most current answer,
  * we need only one that was valid at any point after the call started.
  */
-void *critnib_get(struct critnib *c, word key) {
+void *critnib_get(struct critnib *c, word key, void **ref) {
+    struct critnib_leaf *k;
+    struct critnib_node *n;
     uint64_t wrs1, wrs2;
-    void *res;
+    void *res = NULL;
 
-    do {
-        struct critnib_node *n;
+    if (!ref) {
+        return NULL;
+    }
 
+    do {
         utils_atomic_load_acquire_u64(&c->remove_count, &wrs1);
         utils_atomic_load_acquire_ptr((void **)&c->root, (void **)&n);
 
@@ -529,11 +568,16 @@ void *critnib_get(struct critnib *c, word key) {
         }
 
         /* ... as we check it at the end. */
-        struct critnib_leaf *k = to_leaf(n);
+        k = to_leaf(n);
         res = (n && k->key == key) ? k->value : NULL;
         utils_atomic_load_acquire_u64(&c->remove_count, &wrs2);
     } while (wrs1 + DELETED_LIFE <= wrs2);
 
+    if (res) {
+        utils_atomic_increment_u64(&k->ref_count);
+        *ref = k;
+    }
+
     return res;
 }
 
@@ -645,19 +689,29 @@ static struct critnib_leaf *find_le(struct critnib_node *__restrict n,
  *
  * Same guarantees as critnib_get().
  */
-void *critnib_find_le(struct critnib *c, word key) {
+void *critnib_find_le(struct critnib *c, word key, void **ref) {
+    struct critnib_leaf *k;
     uint64_t wrs1, wrs2;
     void *res;
 
+    if (!ref) {
+        return NULL;
+    }
+
     do {
         utils_atomic_load_acquire_u64(&c->remove_count, &wrs1);
         struct critnib_node *n; /* avoid a subtle TOCTOU */
         utils_atomic_load_acquire_ptr((void **)&c->root, (void **)&n);
-        struct critnib_leaf *k = n ? find_le(n, key) : NULL;
+        k = n ? find_le(n, key) : NULL;
         res = k ? k->value : NULL;
         utils_atomic_load_acquire_u64(&c->remove_count, &wrs2);
     } while (wrs1 + DELETED_LIFE <= wrs2);
 
+    if (res) {
+        utils_atomic_increment_u64(&k->ref_count);
+        *ref = k;
+    }
+
     return res;
 }
 
@@ -743,12 +797,16 @@ static struct critnib_leaf *find_ge(struct critnib_node *__restrict n,
  * critnib_find -- parametrized query, returns 1 if found
  */
 int critnib_find(struct critnib *c, uintptr_t key, enum find_dir_t dir,
-                 uintptr_t *rkey, void **rvalue) {
+                 uintptr_t *rkey, void **rvalue, void **ref) {
     uint64_t wrs1, wrs2;
     struct critnib_leaf *k;
     uintptr_t _rkey = (uintptr_t)0x0;
     void **_rvalue = NULL;
 
+    if (!ref) {
+        return 0;
+    }
+
     /* <42 ≡ ≤41 */
     if (dir < -1) {
         if (!key) {
@@ -790,6 +848,9 @@ int critnib_find(struct critnib *c, uintptr_t key, enum find_dir_t dir,
     } while (wrs1 + DELETED_LIFE <= wrs2);
 
     if (k) {
+        utils_atomic_increment_u64(&k->ref_count);
+        *ref = k;
+
         if (rkey) {
             *rkey = _rkey;
         }
diff --git a/src/critnib/critnib.h b/src/critnib/critnib.h
@@ -33,13 +33,14 @@ void critnib_delete(critnib *c);
 
 int critnib_insert(critnib *c, uintptr_t key, void *value, int update);
 void *critnib_remove(critnib *c, uintptr_t key);
-void *critnib_get(critnib *c, uintptr_t key);
-void *critnib_find_le(critnib *c, uintptr_t key);
+void *critnib_get(critnib *c, uintptr_t key, void **ref);
+void *critnib_find_le(critnib *c, uintptr_t key, void **ref);
 int critnib_find(critnib *c, uintptr_t key, enum find_dir_t dir,
-                 uintptr_t *rkey, void **rvalue);
+                 uintptr_t *rkey, void **rvalue, void **ref);
 void critnib_iter(critnib *c, uintptr_t min, uintptr_t max,
                   int (*func)(uintptr_t key, void *value, void *privdata),
                   void *privdata);
+int critnib_release(struct critnib *c, void *ref);
 
 #ifdef __cplusplus
 }
diff --git a/src/pool/pool_disjoint.c b/src/pool/pool_disjoint.c
@@ -812,8 +812,9 @@ size_t disjoint_pool_malloc_usable_size(void *pool, const void *ptr) {
     }
 
     // check if given pointer is allocated inside any Disjoint Pool slab
-    slab_t *slab =
-        (slab_t *)critnib_find_le(disjoint_pool->known_slabs, (uintptr_t)ptr);
+    void *ref_slab = NULL;
+    slab_t *slab = (slab_t *)critnib_find_le(disjoint_pool->known_slabs,
+                                             (uintptr_t)ptr, &ref_slab);
     if (slab == NULL || ptr >= slab_get_end(slab)) {
         // memory comes directly from the provider
         umf_alloc_info_t allocInfo = {NULL, 0, NULL};
@@ -831,7 +832,12 @@ size_t disjoint_pool_malloc_usable_size(void *pool, const void *ptr) {
 
     ptrdiff_t diff = (ptrdiff_t)ptr - (ptrdiff_t)unaligned_ptr;
 
-    return slab->bucket->size - diff;
+    size_t size = slab->bucket->size - diff;
+
+    assert(ref_slab);
+    critnib_release(disjoint_pool->known_slabs, ref_slab);
+
+    return size;
 }
 
 umf_result_t disjoint_pool_free(void *pool, void *ptr) {
@@ -841,8 +847,9 @@ umf_result_t disjoint_pool_free(void *pool, void *ptr) {
     }
 
     // check if given pointer is allocated inside any Disjoint Pool slab
-    slab_t *slab =
-        (slab_t *)critnib_find_le(disjoint_pool->known_slabs, (uintptr_t)ptr);
+    void *ref_slab = NULL;
+    slab_t *slab = (slab_t *)critnib_find_le(disjoint_pool->known_slabs,
+                                             (uintptr_t)ptr, &ref_slab);
 
     if (slab == NULL || ptr >= slab_get_end(slab)) {
 
@@ -889,6 +896,9 @@ umf_result_t disjoint_pool_free(void *pool, void *ptr) {
     utils_annotate_memory_inaccessible(unaligned_ptr, bucket->size);
     bucket_free_chunk(bucket, unaligned_ptr, slab, &to_pool);
 
+    assert(ref_slab);
+    critnib_release(disjoint_pool->known_slabs, ref_slab);
+
     if (disjoint_pool->params.pool_trace > 1) {
         bucket->free_count++;
     }
diff --git a/src/provider/provider_file_memory.c b/src/provider/provider_file_memory.c
@@ -311,9 +311,11 @@ static void file_finalize(void *provider) {
     uintptr_t key = 0;
     uintptr_t rkey = 0;
     void *rvalue = NULL;
-    while (1 ==
-           critnib_find(file_provider->mmaps, key, FIND_G, &rkey, &rvalue)) {
+    void *ref = NULL;
+    while (1 == critnib_find(file_provider->mmaps, key, FIND_G, &rkey, &rvalue,
+                             &ref)) {
         utils_munmap((void *)rkey, (size_t)rvalue);
+        critnib_release(file_provider->mmaps, ref);
         critnib_remove(file_provider->mmaps, rkey);
         key = rkey;
     }
@@ -644,7 +646,9 @@ static umf_result_t file_allocation_split_cb(void *provider, void *ptr,
         return UMF_RESULT_SUCCESS;
     }
 
-    void *value = critnib_get(file_provider->fd_offset_map, (uintptr_t)ptr);
+    void *ref_value = NULL;
+    void *value =
+        critnib_get(file_provider->fd_offset_map, (uintptr_t)ptr, &ref_value);
     if (value == NULL) {
         LOG_ERR("getting a value from the file descriptor offset map failed "
                 "(addr=%p)",
@@ -658,6 +662,7 @@ static umf_result_t file_allocation_split_cb(void *provider, void *ptr,
 
     uintptr_t new_key = (uintptr_t)ptr + firstSize;
     void *new_value = (void *)((uintptr_t)value + firstSize);
+    critnib_release(file_provider->fd_offset_map, ref_value);
     int ret = critnib_insert(file_provider->fd_offset_map, new_key, new_value,
                              0 /* update */);
     if (ret) {
@@ -734,14 +739,17 @@ static umf_result_t file_get_ipc_handle(void *provider, const void *ptr,
         return UMF_RESULT_ERROR_INVALID_ARGUMENT;
     }
 
-    void *value = critnib_get(file_provider->fd_offset_map, (uintptr_t)ptr);
+    void *ref_value = NULL;
+    void *value =
+        critnib_get(file_provider->fd_offset_map, (uintptr_t)ptr, &ref_value);
     if (value == NULL) {
         LOG_ERR("getting a value from the IPC cache failed (addr=%p)", ptr);
         return UMF_RESULT_ERROR_INVALID_ARGUMENT;
     }
 
     file_ipc_data_t *file_ipc_data = (file_ipc_data_t *)providerIpcData;
     file_ipc_data->offset_fd = (size_t)value - 1;
+    critnib_release(file_provider->fd_offset_map, ref_value);
     file_ipc_data->size = size;
     strncpy(file_ipc_data->path, file_provider->path, PATH_MAX - 1);
     file_ipc_data->path[PATH_MAX - 1] = '\0';
diff --git a/src/provider/provider_os_memory.c b/src/provider/provider_os_memory.c
@@ -1245,7 +1245,9 @@ static umf_result_t os_allocation_split(void *provider, void *ptr,
         return UMF_RESULT_SUCCESS;
     }
 
-    void *value = critnib_get(os_provider->fd_offset_map, (uintptr_t)ptr);
+    void *ref_value = NULL;
+    void *value =
+        critnib_get(os_provider->fd_offset_map, (uintptr_t)ptr, &ref_value);
     if (value == NULL) {
         LOG_ERR("os_allocation_split(): getting a value from the file "
                 "descriptor offset map failed (addr=%p)",
@@ -1255,6 +1257,7 @@ static umf_result_t os_allocation_split(void *provider, void *ptr,
 
     uintptr_t new_key = (uintptr_t)ptr + firstSize;
     void *new_value = (void *)((uintptr_t)value + firstSize);
+    critnib_release(os_provider->fd_offset_map, ref_value);
     int ret = critnib_insert(os_provider->fd_offset_map, new_key, new_value,
                              0 /* update */);
     if (ret) {
@@ -1324,7 +1327,9 @@ static umf_result_t os_get_ipc_handle(void *provider, const void *ptr,
         return UMF_RESULT_ERROR_INVALID_ARGUMENT;
     }
 
-    void *value = critnib_get(os_provider->fd_offset_map, (uintptr_t)ptr);
+    void *ref_value = NULL;
+    void *value =
+        critnib_get(os_provider->fd_offset_map, (uintptr_t)ptr, &ref_value);
     if (value == NULL) {
         LOG_ERR("getting a value from the IPC cache failed (addr=%p)", ptr);
         return UMF_RESULT_ERROR_INVALID_ARGUMENT;
@@ -1333,6 +1338,7 @@ static umf_result_t os_get_ipc_handle(void *provider, const void *ptr,
     os_ipc_data_t *os_ipc_data = (os_ipc_data_t *)providerIpcData;
     os_ipc_data->pid = utils_getpid();
     os_ipc_data->fd_offset = (size_t)value - 1;
+    critnib_release(os_provider->fd_offset_map, ref_value);
     os_ipc_data->size = size;
     os_ipc_data->protection = os_provider->protection;
     os_ipc_data->visibility = os_provider->visibility;
diff --git a/src/provider/provider_tracking.c b/src/provider/provider_tracking.c
