Skip to content

Commit 21aa64b

Browse files
ldorauldorau
committed
new changes
1 parent 6c334bf commit 21aa64b

File tree

2 files changed

+150
-17
lines changed

2 files changed

+150
-17
lines changed

‎src/critnib/critnib.c

Lines changed: 78 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -433,8 +433,25 @@ void *critnib_remove(struct critnib *c, word key) {
433433
goto not_found;
434434
}
435435

436-
word del =
437-
(utils_atomic_increment_u64(&c->remove_count) - 1) % DELETED_LIFE;
436+
word del;
437+
int i = 0;
438+
do {
439+
del = (utils_atomic_increment_u64(&c->remove_count) - 1) % DELETED_LIFE;
440+
if (c->pending_del_leaves[del] &&
441+
(c->pending_del_leaves[del]->ref_count > 0)) {
442+
fprintf(stderr,
443+
"critnib_remove(): pending_del_leaves[%llu]->key = %p, "
444+
"ref_count=%llu \n",
445+
(unsigned long long)del,
446+
(void *)c->pending_del_leaves[del]->key,
447+
(unsigned long long)c->pending_del_leaves[del]->ref_count);
448+
if (i++ == DELETED_LIFE) {
449+
break;
450+
}
451+
}
452+
} while (c->pending_del_leaves[del] &&
453+
(c->pending_del_leaves[del]->ref_count > 0));
454+
438455
free_node(c, c->pending_del_nodes[del]);
439456
free_leaf(c, c->pending_del_leaves[del]);
440457
c->pending_del_nodes[del] = NULL;
@@ -494,12 +511,18 @@ void *critnib_remove(struct critnib *c, word key) {
494511
c->pending_del_nodes[del] = n;
495512

496513
del_leaf:
497-
uint64_t ref_count = utils_atomic_decrement_u64(&k->ref_count);
498-
fprintf(stderr,
499-
"critnib_remove(): k->ref_count = %llu of k->key = %p\n",
500-
(unsigned long long)ref_count, (void *)k->key);
501-
502514
value = k->value;
515+
uint64_t ref_count;
516+
if ((ref_count = utils_atomic_decrement_u64(&k->ref_count)) == 0) {
517+
utils_atomic_store_release_ptr((void **)&k->value, NULL);
518+
if (value && c->cb_free_leaf) {
519+
c->cb_free_leaf(c->leaf_allocator, (void *)value);
520+
}
521+
}
522+
523+
fprintf(stderr, "critnib_remove(): k->ref_count = %llu of k->key = %p\n",
524+
(unsigned long long)k->ref_count, (void *)k->key);
525+
503526
c->pending_del_leaves[del] = k;
504527

505528
not_found:
@@ -521,9 +544,14 @@ int critnib_release(struct critnib *c, void *ref) {
521544

522545
/* decrement the reference count */
523546
if (utils_atomic_decrement_u64(&k->ref_count) == 0) {
524-
fprintf(stderr, "critnib_release(): k->ref_count = %llu of k->key = %p\n",
525-
(unsigned long long)k->ref_count, (void *)k->key);
526-
free_leaf(c, k);
547+
void *value = k->value;
548+
utils_atomic_store_release_ptr((void **)&k->value, NULL);
549+
if (value && c->cb_free_leaf) {
550+
c->cb_free_leaf(c->leaf_allocator, value);
551+
}
552+
553+
fprintf(stderr, "critnib_release(): k->ref_count = 0 of k->key = %p\n",
554+
(void *)k->key);
527555
}
528556

529557
return 0;
@@ -570,7 +598,19 @@ void *critnib_get(struct critnib *c, word key, void **ref) {
570598
} while (wrs1 + DELETED_LIFE <= wrs2);
571599

572600
if (res) {
573-
utils_atomic_increment_u64(&k->ref_count);
601+
uint64_t ref_count;
602+
utils_atomic_load_acquire_u64(&k->ref_count, &ref_count);
603+
if (ref_count == 0) {
604+
return NULL;
605+
}
606+
if ((ref_count = utils_atomic_increment_u64(&k->ref_count)) == 1) {
607+
utils_atomic_decrement_u64(&k->ref_count);
608+
return NULL;
609+
}
610+
611+
fprintf(stderr, "critnib_get(): k->ref_count = %llu of k->key = %p\n",
612+
(unsigned long long)ref_count, (void *)k->key);
613+
574614
*ref = k;
575615
}
576616

@@ -704,7 +744,20 @@ void *critnib_find_le(struct critnib *c, word key, void **ref) {
704744
} while (wrs1 + DELETED_LIFE <= wrs2);
705745

706746
if (res) {
707-
utils_atomic_increment_u64(&k->ref_count);
747+
uint64_t ref_count;
748+
utils_atomic_load_acquire_u64(&k->ref_count, &ref_count);
749+
if (ref_count == 0) {
750+
return NULL;
751+
}
752+
if ((ref_count = utils_atomic_increment_u64(&k->ref_count)) == 1) {
753+
utils_atomic_decrement_u64(&k->ref_count);
754+
return NULL;
755+
}
756+
757+
fprintf(stderr,
758+
"critnib_find_le(): k->ref_count = %llu of k->key = %p\n",
759+
(unsigned long long)ref_count, (void *)k->key);
760+
708761
*ref = k;
709762
}
710763

@@ -844,7 +897,19 @@ int critnib_find(struct critnib *c, uintptr_t key, enum find_dir_t dir,
844897
} while (wrs1 + DELETED_LIFE <= wrs2);
845898

846899
if (k) {
847-
utils_atomic_increment_u64(&k->ref_count);
900+
uint64_t ref_count;
901+
utils_atomic_load_acquire_u64(&k->ref_count, &ref_count);
902+
if (ref_count == 0) {
903+
return 0;
904+
}
905+
if ((ref_count = utils_atomic_increment_u64(&k->ref_count)) == 1) {
906+
utils_atomic_decrement_u64(&k->ref_count);
907+
return 0;
908+
}
909+
910+
fprintf(stderr, "critnib_find(): k->ref_count = %llu of k->key = %p\n",
911+
(unsigned long long)ref_count, (void *)k->key);
912+
848913
*ref = k;
849914

850915
if (rkey) {

‎src/provider/provider_tracking.c

Lines changed: 72 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -88,10 +88,26 @@ static tracker_alloc_info_t *get_most_nested_alloc_segment(
8888
found =
8989
critnib_find(hTracker->alloc_segments_map[level], (uintptr_t)ptr,
9090
FIND_LE, (void *)&rkey, (void **)&rvalue, &ref_value);
91-
if (!found || !rvalue) {
91+
if (!found) {
92+
assert(ref_value == NULL);
9293
break;
9394
}
9495

96+
if (rvalue == NULL) {
97+
if (ref_value) {
98+
critnib_release(hTracker->alloc_segments_map[level], ref_value);
99+
}
100+
parent_value = NULL;
101+
rvalue = NULL;
102+
parent_key = 0;
103+
rkey = 0;
104+
level = 0;
105+
found = 0;
106+
ref_value = NULL;
107+
ref_parent_value = NULL;
108+
continue;
109+
}
110+
95111
if (found && ((uintptr_t)ptr < rkey + rvalue->size) &&
96112
rvalue->n_children) {
97113
if (level == MAX_LEVELS_OF_ALLOC_SEGMENT_MAP - 1) {
@@ -226,10 +242,26 @@ static umf_result_t umfMemoryTrackerAdd(umf_memory_tracker_handle_t hTracker,
226242
found =
227243
critnib_find(hTracker->alloc_segments_map[level], (uintptr_t)ptr,
228244
FIND_LE, (void *)&rkey, (void **)&rvalue, &ref_value);
229-
if (!found || !rvalue) {
245+
if (!found) {
246+
assert(ref_value == NULL);
230247
break;
231248
}
232249

250+
if (!rvalue) {
251+
if (ref_value) {
252+
critnib_release(hTracker->alloc_segments_map[level], ref_value);
253+
}
254+
parent_value = NULL;
255+
rvalue = NULL;
256+
parent_key = 0;
257+
rkey = 0;
258+
level = 0;
259+
found = 0;
260+
ref_value = NULL;
261+
ref_parent_value = NULL;
262+
continue;
263+
}
264+
233265
if ((uintptr_t)ptr < rkey + rvalue->size) {
234266
if (level == MAX_LEVELS_OF_ALLOC_SEGMENT_MAP - 1) {
235267
// TODO: we need to support an arbitrary amount of layers in the future
@@ -259,6 +291,10 @@ static umf_result_t umfMemoryTrackerAdd(umf_memory_tracker_handle_t hTracker,
259291
} while (found && ((uintptr_t)ptr < rkey + rvalue->size) &&
260292
rvalue->n_children);
261293

294+
if (ref_value && ref_value != ref_parent_value) {
295+
critnib_release(hTracker->alloc_segments_map[level], ref_value);
296+
}
297+
262298
return umfMemoryTrackerAddAtLevel(hTracker, level, pool, ptr, size,
263299
parent_key, parent_value,
264300
ref_parent_value);
@@ -426,6 +462,27 @@ umf_result_t umfMemoryTrackerGetAllocInfo(const void *ptr,
426462
found =
427463
critnib_find(TRACKER->alloc_segments_map[level], (uintptr_t)ptr,
428464
FIND_LE, (void *)&rkey, (void **)&rvalue, &ref_value);
465+
if (!found) {
466+
assert(ref_value == NULL);
467+
break;
468+
}
469+
470+
if (!rvalue) {
471+
if (ref_value) {
472+
critnib_release(TRACKER->alloc_segments_map[level], ref_value);
473+
}
474+
top_most_value = NULL;
475+
rvalue = NULL;
476+
top_most_key = 0;
477+
rkey = 0;
478+
level = 0;
479+
found = 0;
480+
ref_value = NULL;
481+
ref_top_most_value = NULL;
482+
ref_level = 0;
483+
continue;
484+
}
485+
429486
if (found && (uintptr_t)ptr < rkey + rvalue->size) {
430487
top_most_key = rkey;
431488
top_most_value = rvalue;
@@ -483,11 +540,22 @@ umf_result_t umfMemoryTrackerGetIpcInfo(const void *ptr,
483540
return UMF_RESULT_ERROR_NOT_SUPPORTED;
484541
}
485542

486-
uintptr_t rkey;
543+
uintptr_t rkey = 0;
487544
tracker_ipc_info_t *rvalue = NULL;
488545
void *ref_value = NULL;
489546
int found = critnib_find(TRACKER->ipc_segments_map, (uintptr_t)ptr, FIND_LE,
490547
(void *)&rkey, (void **)&rvalue, &ref_value);
548+
549+
if (found && !rvalue) {
550+
found = 0;
551+
if (ref_value) {
552+
critnib_release(TRACKER->ipc_segments_map, ref_value);
553+
}
554+
rkey = 0;
555+
rvalue = NULL;
556+
ref_value = NULL;
557+
}
558+
491559
if (!found || (uintptr_t)ptr >= rkey + rvalue->size) {
492560
LOG_DEBUG("pointer %p not found in the tracker, TRACKER=%p", ptr,
493561
(void *)TRACKER);
@@ -1174,7 +1242,7 @@ static umf_result_t trackingCloseIpcHandle(void *provider, void *ptr,
11741242
umfIpcHandleMappedCacheRelease(trackerIpcInfo->ipc_cache_value);
11751243

11761244
assert(ref_value);
1177-
critnib_release(TRACKER->ipc_segments_map, &ref_value);
1245+
critnib_release(TRACKER->ipc_segments_map, ref_value);
11781246

11791247
return umf_result;
11801248
}

0 commit comments

Comments
 (0)