Fix race between critnib_release() and free_leaf() in critnib

Fix race between critnib_release() and free_leaf() in critnib:
critnib_release() decremented ref_count to 0
and (before it called c->cb_free_leaf(k->to_be_freed))
free_leaf() added this leaf to the c->deleted_leaf list
and alloc_leaf() reused it and zeroed k->to_be_freed
before it could be freed in critnib_release().
This patch fixes this issue.

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

Author: ldorau

src/critnib/critnib.c

@@ -89,6 +89,9 @@
 #define NIB ((1ULL << SLICE) - 1)
 #define SLNODES (1 << SLICE)
 
+// it has to be >= 2
+#define LEAF_VALID (2ULL)
+
 typedef uintptr_t word;
 typedef uint8_t sh_t;
 
@@ -331,18 +334,26 @@ static void free_leaf(struct critnib *__restrict c,
         return;
     }
 
+    // k should be added to the c->deleted_leaf list here
+    // or in critnib_release() when the reference count drops to 0.
+    utils_atomic_store_release_u8(&k->pending_deleted_leaf, 1);
+
     if (c->cb_free_leaf) {
         uint64_t ref_count;
         utils_atomic_load_acquire_u64(&k->ref_count, &ref_count);
         if (ref_count > 0) {
-            // k will be added to c->deleted_leaf in critnib_release()
+            // k will be added to the c->deleted_leaf list in critnib_release()
             // when the reference count drops to 0.
-            utils_atomic_store_release_u8(&k->pending_deleted_leaf, 1);
             return;
         }
     }
 
-    add_to_deleted_leaf_list(c, k);
+    uint8_t expected = 1;
+    uint8_t desired = 0;
+    if (utils_compare_exchange_u8(&k->pending_deleted_leaf, &expected,
+                                  &desired)) {
+        add_to_deleted_leaf_list(c, k);
+    }
 }
 
 /*
@@ -392,8 +403,8 @@ int critnib_insert(struct critnib *c, word key, void *value, int update) {
     utils_atomic_store_release_u8(&k->pending_deleted_leaf, 0);
 
     if (c->cb_free_leaf) {
-        // mark the leaf as valid (ref_count == 1)
-        utils_atomic_store_release_u64(&k->ref_count, 1ULL);
+        // mark the leaf as valid (ref_count == 2)
+        utils_atomic_store_release_u64(&k->ref_count, LEAF_VALID);
     } else {
         // the reference counter is not used in this case
         utils_atomic_store_release_u64(&k->ref_count, 0ULL);
@@ -602,36 +613,50 @@ int critnib_release(struct critnib *c, void *ref) {
     struct critnib_leaf *k = (struct critnib_leaf *)ref;
 
     uint64_t ref_count;
-    utils_atomic_load_acquire_u64(&k->ref_count, &ref_count);
-
-    if (ref_count == 0) {
-        return -1;
-    }
-
+    uint64_t ref_desired;
     /* decrement the reference count */
-    if (utils_atomic_decrement_u64(&k->ref_count) == 0) {
-        void *to_be_freed = NULL;
-        utils_atomic_load_acquire_ptr(&k->to_be_freed, &to_be_freed);
-        if (to_be_freed) {
-            utils_atomic_store_release_ptr(&k->to_be_freed, NULL);
-            c->cb_free_leaf(c->leaf_allocator, to_be_freed);
-        }
-        uint8_t pending_deleted_leaf;
-        utils_atomic_load_acquire_u8(&k->pending_deleted_leaf,
-                                     &pending_deleted_leaf);
-        if (pending_deleted_leaf) {
-            utils_atomic_store_release_u8(&k->pending_deleted_leaf, 0);
-            add_to_deleted_leaf_list(c, k);
+    do {
+        utils_atomic_load_acquire_u64(&k->ref_count, &ref_count);
+        if (ref_count < LEAF_VALID) {
+            LOG_FATAL("critnib_release() was called too many times (ref_count "
+                      "= %llu)\n",
+                      (unsigned long long)ref_count);
+            assert(ref_count >= LEAF_VALID);
+            return -1;
         }
+        ref_desired = ref_count - 1;
+    } while (
+        !utils_compare_exchange_u64(&k->ref_count, &ref_count, &ref_desired));
+
+    if (ref_desired >= LEAF_VALID) {
+        // ref_counter was decremented and it is still valid
+        return 0;
     }
 
+    /* ref_counter == (LEAF_VALID - 1)) - the leaf will be freed */
+    void *to_be_freed = NULL;
+    utils_atomic_load_acquire_ptr(&k->to_be_freed, &to_be_freed);
+    utils_atomic_store_release_ptr(&k->to_be_freed, NULL);
 #ifndef NDEBUG
-    // check if the reference count is overflowed
-    utils_atomic_load_acquire_u64(&k->ref_count, &ref_count);
-    assert((ref_count & (1ULL << 63)) == 0);
-    assert(ref_count != (uint64_t)(0 - 1ULL));
+    if (to_be_freed == NULL) {
+        LOG_FATAL("leaf will not be freed (to_be_freed == NULL, value = %p)\n",
+                  k->value);
+        assert(to_be_freed != NULL);
+    }
 #endif
 
+    // mark the leaf as not used (ref_count == 0)
+    utils_atomic_store_release_u64(&k->ref_count, 0ULL);
+
+    c->cb_free_leaf(c->leaf_allocator, to_be_freed);
+
+    uint8_t expected = 1;
+    uint8_t desired = 0;
+    if (utils_compare_exchange_u8(&k->pending_deleted_leaf, &expected,
+                                  &desired)) {
+        add_to_deleted_leaf_list(c, k);
+    }
+
     return 0;
 }
 
@@ -661,7 +686,7 @@ static inline int increment_ref_count(struct critnib_leaf *k) {
 
     do {
         utils_atomic_load_acquire_u64(&k->ref_count, &expected);
-        if (expected == 0) {
+        if (expected < LEAF_VALID) {
             return -1;
         }
         desired = expected + 1;