[WIP!] critnib: add refcount WIP!

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

Author: ldorau

src/critnib/critnib.c

@@ -116,6 +116,8 @@ struct critnib_node {
 struct critnib_leaf {
     word key;
     void *value;
+    void *to_be_freed;
+    uint64_t ref_count;
 };
 
 struct critnib {
@@ -194,7 +196,12 @@ struct critnib *critnib_new(free_leaf_t cb_free_leaf, void *leaf_allocator) {
 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);
+            if (to_leaf(n)->value) {
+                c->cb_free_leaf(c->leaf_allocator, (void *)to_leaf(n)->value);
+            } else if (to_leaf(n)->to_be_freed) {
+                c->cb_free_leaf(c->leaf_allocator,
+                                (void *)to_leaf(n)->to_be_freed);
+            }
         }
         umf_ba_global_free(to_leaf(n));
     } else {
@@ -233,8 +240,13 @@ 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);
+            if (c->pending_del_leaves[i]->value) {
+                c->cb_free_leaf(c->leaf_allocator,
+                                (void *)c->pending_del_leaves[i]->value);
+            } else if (c->pending_del_leaves[i]->to_be_freed) {
+                c->cb_free_leaf(c->leaf_allocator,
+                                (void *)c->pending_del_leaves[i]->to_be_freed);
+            }
         }
         umf_ba_global_free(c->pending_del_leaves[i]);
     }
@@ -288,8 +300,8 @@ static void free_leaf(struct critnib *__restrict c,
         return;
     }
 
-    if (c->cb_free_leaf && k && k->value) {
-        c->cb_free_leaf(c->leaf_allocator, (void *)k->value);
+    if (c->cb_free_leaf && k && k->to_be_freed) {
+        c->cb_free_leaf(c->leaf_allocator, (void *)k->to_be_freed);
     }
 
     utils_atomic_store_release_ptr((void **)&k->value, c->deleted_leaf);
@@ -334,8 +346,11 @@ int critnib_insert(struct critnib *c, word key, void *value, int update) {
 
     utils_annotate_memory_no_check(k, sizeof(struct critnib_leaf));
 
+    utils_atomic_store_release_ptr(&k->to_be_freed, 0);
     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);
 
@@ -370,9 +385,9 @@ int critnib_insert(struct critnib *c, word key, void *value, int update) {
     word at = path ^ key;
     if (!at) {
         ASSERT(is_leaf(n));
-        if (to_leaf(kn)->value == value) {
+        if (to_leaf(kn)->to_be_freed == value) {
             // do not free the value
-            to_leaf(kn)->value = NULL;
+            to_leaf(kn)->to_be_freed = NULL;
         }
         free_leaf(c, to_leaf(kn));
 
@@ -418,19 +433,31 @@ int critnib_insert(struct critnib *c, word key, void *value, int update) {
 /*
  * critnib_remove -- delete a key from the critnib structure, return its value
  */
-void *critnib_remove(struct critnib *c, word key) {
+void *critnib_remove(struct critnib *c, word key, void **ref) {
     struct critnib_leaf *k;
     void *value = NULL;
 
+    if (!ref) {
+        return NULL;
+    }
+
     utils_mutex_lock(&c->mutex);
 
     struct critnib_node *n = c->root;
     if (!n) {
         goto not_found;
     }
 
-    word del =
-        (utils_atomic_increment_u64(&c->remove_count) - 1) % DELETED_LIFE;
+    word del;
+    int i_del = 0;
+    do {
+        del = (utils_atomic_increment_u64(&c->remove_count) - 1) % DELETED_LIFE;
+        k = c->pending_del_leaves[del];
+        if (i_del++ == DELETED_LIFE) {
+            break;
+        }
+    } while (k && (k->ref_count > 0));
+
     free_node(c, c->pending_del_nodes[del]);
     free_leaf(c, c->pending_del_leaves[del]);
     c->pending_del_nodes[del] = NULL;
@@ -491,13 +518,40 @@ void *critnib_remove(struct critnib *c, word key) {
 
 del_leaf:
     value = k->value;
+    *ref = k;
+    utils_atomic_store_release_ptr(&k->to_be_freed, value);
+    utils_atomic_store_release_ptr((void **)&k->value, NULL);
     c->pending_del_leaves[del] = k;
 
 not_found:
     utils_mutex_unlock(&c->mutex);
     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;
+
+    /* decrement the reference count */
+    if (k->ref_count && (utils_atomic_decrement_u64(&k->ref_count) == 0) &&
+        k->to_be_freed) {
+        void *value = k->to_be_freed;
+        utils_atomic_store_release_ptr(&k->to_be_freed, NULL);
+        utils_atomic_store_release_u64(&k->key, 0);
+        if (value && c->cb_free_leaf) {
+            c->cb_free_leaf(c->leaf_allocator, value);
+        }
+    }
+
+    return 0;
+}
+
 /*
  * critnib_get -- query for a key ("==" match), returns value or NULL
  *
@@ -508,13 +562,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 +587,25 @@ 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) {
+        uint64_t ref_count;
+        utils_atomic_load_acquire_u64(&k->ref_count, &ref_count);
+        if (ref_count == 0) {
+            return NULL;
+        }
+        if (utils_atomic_increment_u64(&k->ref_count) == 1) {
+            utils_atomic_decrement_u64(&k->ref_count);
+            return NULL;
+        }
+
+        *ref = k;
+    }
+
     return res;
 }
 
@@ -645,19 +717,38 @@ 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) {
+        uint64_t ref_count;
+        utils_atomic_load_acquire_u64(&k->ref_count, &ref_count);
+        if (ref_count == 0) {
+            return NULL;
+        }
+        if (utils_atomic_increment_u64(&k->ref_count) == 1) {
+            utils_atomic_decrement_u64(&k->ref_count);
+            return NULL;
+        }
+
+        *ref = k;
+    }
+
     return res;
 }
 
@@ -743,12 +834,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 +885,18 @@ int critnib_find(struct critnib *c, uintptr_t key, enum find_dir_t dir,
     } while (wrs1 + DELETED_LIFE <= wrs2);
 
     if (k) {
+        uint64_t ref_count;
+        utils_atomic_load_acquire_u64(&k->ref_count, &ref_count);
+        if (ref_count == 0) {
+            return 0;
+        }
+        if (utils_atomic_increment_u64(&k->ref_count) == 1) {
+            utils_atomic_decrement_u64(&k->ref_count);
+            return 0;
+        }
+
+        *ref = k;
+
         if (rkey) {
             *rkey = _rkey;
         }

src/critnib/critnib.h

@@ -32,14 +32,15 @@ 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);
-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_remove(critnib *c, uintptr_t key, void **ref);
+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
 }

src/pool/pool_disjoint.c

@@ -231,7 +231,11 @@ static umf_result_t pool_unregister_slab(disjoint_pool_t *pool, slab_t *slab) {
     // TODO ASSERT_IS_ALIGNED((uintptr_t)slab_addr, bucket->size);
     LOG_DEBUG("slab: %p, start: %p", (void *)slab, slab_addr);
 
-    critnib_remove(slabs, (uintptr_t)slab_addr);
+    void *ref_slabs = NULL;
+    critnib_remove(slabs, (uintptr_t)slab_addr, &ref_slabs);
+    if (ref_slabs) {
+        critnib_release(slabs, ref_slabs);
+    }
 
     return UMF_RESULT_SUCCESS;
 }
@@ -812,8 +816,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 +836,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 +851,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 +900,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++;
     }