Skip to content

Commit e422121

Browse files
committed
new changes
1 parent dfc62fd commit e422121

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
@@ -437,8 +437,25 @@ void *critnib_remove(struct critnib *c, word key) {
437437
goto not_found;
438438
}
439439

440-
word del =
441-
(utils_atomic_increment_u64(&c->remove_count) - 1) % DELETED_LIFE;
440+
word del;
441+
int i = 0;
442+
do {
443+
del = (utils_atomic_increment_u64(&c->remove_count) - 1) % DELETED_LIFE;
444+
if (c->pending_del_leaves[del] &&
445+
(c->pending_del_leaves[del]->ref_count > 0)) {
446+
fprintf(stderr,
447+
"critnib_remove(): pending_del_leaves[%llu]->key = %p, "
448+
"ref_count=%llu \n",
449+
(unsigned long long)del,
450+
(void *)c->pending_del_leaves[del]->key,
451+
(unsigned long long)c->pending_del_leaves[del]->ref_count);
452+
if (i++ == DELETED_LIFE) {
453+
break;
454+
}
455+
}
456+
} while (c->pending_del_leaves[del] &&
457+
(c->pending_del_leaves[del]->ref_count > 0));
458+
442459
free_node(c, c->pending_del_nodes[del]);
443460
free_leaf(c, c->pending_del_leaves[del]);
444461
c->pending_del_nodes[del] = NULL;
@@ -498,12 +515,18 @@ void *critnib_remove(struct critnib *c, word key) {
498515
c->pending_del_nodes[del] = n;
499516

500517
del_leaf:
501-
uint64_t ref_count = utils_atomic_decrement_u64(&k->ref_count);
502-
fprintf(stderr,
503-
"critnib_remove(): k->ref_count = %llu of k->key = %p\n",
504-
(unsigned long long)ref_count, (void *)k->key);
505-
506518
value = k->value;
519+
uint64_t ref_count;
520+
if ((ref_count = utils_atomic_decrement_u64(&k->ref_count)) == 0) {
521+
utils_atomic_store_release_ptr((void **)&k->value, NULL);
522+
if (value && c->cb_free_leaf) {
523+
c->cb_free_leaf(c->leaf_allocator, (void *)value);
524+
}
525+
}
526+
527+
fprintf(stderr, "critnib_remove(): k->ref_count = %llu of k->key = %p\n",
528+
(unsigned long long)k->ref_count, (void *)k->key);
529+
507530
c->pending_del_leaves[del] = k;
508531

509532
not_found:
@@ -525,9 +548,14 @@ int critnib_release(struct critnib *c, void *ref) {
525548

526549
/* decrement the reference count */
527550
if (utils_atomic_decrement_u64(&k->ref_count) == 0) {
528-
fprintf(stderr, "critnib_release(): k->ref_count = %llu of k->key = %p\n",
529-
(unsigned long long)k->ref_count, (void *)k->key);
530-
free_leaf(c, k);
551+
void *value = k->value;
552+
utils_atomic_store_release_ptr((void **)&k->value, NULL);
553+
if (value && c->cb_free_leaf) {
554+
c->cb_free_leaf(c->leaf_allocator, value);
555+
}
556+
557+
fprintf(stderr, "critnib_release(): k->ref_count = 0 of k->key = %p\n",
558+
(void *)k->key);
531559
}
532560

533561
return 0;
@@ -574,7 +602,19 @@ void *critnib_get(struct critnib *c, word key, void **ref) {
574602
} while (wrs1 + DELETED_LIFE <= wrs2);
575603

576604
if (res) {
577-
utils_atomic_increment_u64(&k->ref_count);
605+
uint64_t ref_count;
606+
utils_atomic_load_acquire_u64(&k->ref_count, &ref_count);
607+
if (ref_count == 0) {
608+
return NULL;
609+
}
610+
if ((ref_count = utils_atomic_increment_u64(&k->ref_count)) == 1) {
611+
utils_atomic_decrement_u64(&k->ref_count);
612+
return NULL;
613+
}
614+
615+
fprintf(stderr, "critnib_get(): k->ref_count = %llu of k->key = %p\n",
616+
(unsigned long long)ref_count, (void *)k->key);
617+
578618
*ref = k;
579619
}
580620

@@ -708,7 +748,20 @@ void *critnib_find_le(struct critnib *c, word key, void **ref) {
708748
} while (wrs1 + DELETED_LIFE <= wrs2);
709749

710750
if (res) {
711-
utils_atomic_increment_u64(&k->ref_count);
751+
uint64_t ref_count;
752+
utils_atomic_load_acquire_u64(&k->ref_count, &ref_count);
753+
if (ref_count == 0) {
754+
return NULL;
755+
}
756+
if ((ref_count = utils_atomic_increment_u64(&k->ref_count)) == 1) {
757+
utils_atomic_decrement_u64(&k->ref_count);
758+
return NULL;
759+
}
760+
761+
fprintf(stderr,
762+
"critnib_find_le(): k->ref_count = %llu of k->key = %p\n",
763+
(unsigned long long)ref_count, (void *)k->key);
764+
712765
*ref = k;
713766
}
714767

@@ -848,7 +901,19 @@ int critnib_find(struct critnib *c, uintptr_t key, enum find_dir_t dir,
848901
} while (wrs1 + DELETED_LIFE <= wrs2);
849902

850903
if (k) {
851-
utils_atomic_increment_u64(&k->ref_count);
904+
uint64_t ref_count;
905+
utils_atomic_load_acquire_u64(&k->ref_count, &ref_count);
906+
if (ref_count == 0) {
907+
return 0;
908+
}
909+
if ((ref_count = utils_atomic_increment_u64(&k->ref_count)) == 1) {
910+
utils_atomic_decrement_u64(&k->ref_count);
911+
return 0;
912+
}
913+
914+
fprintf(stderr, "critnib_find(): k->ref_count = %llu of k->key = %p\n",
915+
(unsigned long long)ref_count, (void *)k->key);
916+
852917
*ref = k;
853918

854919
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);
@@ -1173,7 +1241,7 @@ static umf_result_t trackingCloseIpcHandle(void *provider, void *ptr,
11731241
umfIpcHandleMappedCacheRelease(trackerIpcInfo->ipc_cache_value);
11741242

11751243
assert(ref_value);
1176-
critnib_release(TRACKER->ipc_segments_map, &ref_value);
1244+
critnib_release(TRACKER->ipc_segments_map, ref_value);
11771245

11781246
return umf_result;
11791247
}

0 commit comments

Comments
 (0)