@@ -116,6 +116,7 @@ struct critnib_node {
116116struct critnib_leaf {
117117 word key ;
118118 void * value ;
119+ void * to_be_freed ;
119120 uint64_t ref_count ;
120121};
121122
@@ -195,7 +196,12 @@ struct critnib *critnib_new(free_leaf_t cb_free_leaf, void *leaf_allocator) {
195196static void delete_node (struct critnib * c , struct critnib_node * __restrict n ) {
196197 if (is_leaf (n )) {
197198 if (c -> cb_free_leaf && to_leaf (n )) {
198- c -> cb_free_leaf (c -> leaf_allocator , (void * )to_leaf (n )-> value );
199+ if (to_leaf (n )-> value ) {
200+ c -> cb_free_leaf (c -> leaf_allocator , (void * )to_leaf (n )-> value );
201+ } else if (to_leaf (n )-> to_be_freed ) {
202+ c -> cb_free_leaf (c -> leaf_allocator ,
203+ (void * )to_leaf (n )-> to_be_freed );
204+ }
199205 }
200206 umf_ba_global_free (to_leaf (n ));
201207 } else {
@@ -234,8 +240,13 @@ void critnib_delete(struct critnib *c) {
234240 for (int i = 0 ; i < DELETED_LIFE ; i ++ ) {
235241 umf_ba_global_free (c -> pending_del_nodes [i ]);
236242 if (c -> cb_free_leaf && c -> pending_del_leaves [i ]) {
237- c -> cb_free_leaf (c -> leaf_allocator ,
238- (void * )c -> pending_del_leaves [i ]-> value );
243+ if (c -> pending_del_leaves [i ]-> value ) {
244+ c -> cb_free_leaf (c -> leaf_allocator ,
245+ (void * )c -> pending_del_leaves [i ]-> value );
246+ } else if (c -> pending_del_leaves [i ]-> to_be_freed ) {
247+ c -> cb_free_leaf (c -> leaf_allocator ,
248+ (void * )c -> pending_del_leaves [i ]-> to_be_freed );
249+ }
239250 }
240251 umf_ba_global_free (c -> pending_del_leaves [i ]);
241252 }
@@ -289,9 +300,8 @@ static void free_leaf(struct critnib *__restrict c,
289300 return ;
290301 }
291302
292- if (c -> cb_free_leaf && k && k -> value ) {
293- assert (k -> ref_count == 0 ); // TODO: check ref_count
294- c -> cb_free_leaf (c -> leaf_allocator , (void * )k -> value );
303+ if (c -> cb_free_leaf && k && k -> to_be_freed ) {
304+ c -> cb_free_leaf (c -> leaf_allocator , (void * )k -> to_be_freed );
295305 }
296306
297307 utils_atomic_store_release_ptr ((void * * )& k -> value , c -> deleted_leaf );
@@ -336,6 +346,7 @@ int critnib_insert(struct critnib *c, word key, void *value, int update) {
336346
337347 utils_annotate_memory_no_check (k , sizeof (struct critnib_leaf ));
338348
349+ utils_atomic_store_release_ptr (& k -> to_be_freed , 0 );
339350 utils_atomic_store_release_ptr ((void * * )& k -> key , (void * )key );
340351 utils_atomic_store_release_ptr ((void * * )& k -> value , value );
341352 // mark the leaf as valid (ref_count == 1)
@@ -422,10 +433,14 @@ int critnib_insert(struct critnib *c, word key, void *value, int update) {
422433/*
423434 * critnib_remove -- delete a key from the critnib structure, return its value
424435 */
425- void * critnib_remove (struct critnib * c , word key ) {
436+ void * critnib_remove (struct critnib * c , word key , void * * ref ) {
426437 struct critnib_leaf * k ;
427438 void * value = NULL ;
428439
440+ if (!ref ) {
441+ return NULL ;
442+ }
443+
429444 utils_mutex_lock (& c -> mutex );
430445
431446 struct critnib_node * n = c -> root ;
@@ -503,14 +518,9 @@ void *critnib_remove(struct critnib *c, word key) {
503518
504519del_leaf :
505520 value = k -> value ;
506- uint64_t ref_count ;
507- if ((ref_count = utils_atomic_decrement_u64 (& k -> ref_count )) == 0 ) {
508- utils_atomic_store_release_ptr ((void * * )& k -> value , NULL );
509- if (value && c -> cb_free_leaf ) {
510- c -> cb_free_leaf (c -> leaf_allocator , (void * )value );
511- }
512- }
513-
521+ * ref = k ;
522+ utils_atomic_store_release_ptr (& k -> to_be_freed , value );
523+ utils_atomic_store_release_ptr ((void * * )& k -> value , NULL );
514524 c -> pending_del_leaves [del ] = k ;
515525
516526not_found :
@@ -528,12 +538,11 @@ int critnib_release(struct critnib *c, void *ref) {
528538
529539 struct critnib_leaf * k = (struct critnib_leaf * )ref ;
530540
531- // (TODO HERE...)
532-
533541 /* decrement the reference count */
534542 if (utils_atomic_decrement_u64 (& k -> ref_count ) == 0 ) {
535- void * value = k -> value ;
536- utils_atomic_store_release_ptr ((void * * )& k -> value , NULL );
543+ void * value = k -> to_be_freed ;
544+ utils_atomic_store_release_ptr (& k -> to_be_freed , NULL );
545+ utils_atomic_store_release_u64 (& k -> key , 0 );
537546 if (value && c -> cb_free_leaf ) {
538547 c -> cb_free_leaf (c -> leaf_allocator , value );
539548 }
0 commit comments