Skip to content

Commit a64c7f6

Browse files
committed
Fixes
1 parent 9469eed commit a64c7f6

File tree

6 files changed

+93
-32
lines changed

6 files changed

+93
-32
lines changed

src/critnib/critnib.c

Lines changed: 28 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -116,6 +116,7 @@ struct critnib_node {
116116
struct 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) {
195196
static 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

504519
del_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

516526
not_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
}

src/critnib/critnib.h

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@ critnib *critnib_new(free_leaf_t cb_free_leaf, void *leaf_allocator);
3232
void critnib_delete(critnib *c);
3333

3434
int critnib_insert(critnib *c, uintptr_t key, void *value, int update);
35-
void *critnib_remove(critnib *c, uintptr_t key);
35+
void *critnib_remove(critnib *c, uintptr_t key, void **ref);
3636
void *critnib_get(critnib *c, uintptr_t key, void **ref);
3737
void *critnib_find_le(critnib *c, uintptr_t key, void **ref);
3838
int critnib_find(critnib *c, uintptr_t key, enum find_dir_t dir,

src/pool/pool_disjoint.c

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -231,7 +231,11 @@ static umf_result_t pool_unregister_slab(disjoint_pool_t *pool, slab_t *slab) {
231231
// TODO ASSERT_IS_ALIGNED((uintptr_t)slab_addr, bucket->size);
232232
LOG_DEBUG("slab: %p, start: %p", (void *)slab, slab_addr);
233233

234-
critnib_remove(slabs, (uintptr_t)slab_addr);
234+
void *ref_slabs = NULL;
235+
critnib_remove(slabs, (uintptr_t)slab_addr, &ref_slabs);
236+
if (ref_slabs) {
237+
critnib_release(slabs, ref_slabs);
238+
}
235239

236240
return UMF_RESULT_SUCCESS;
237241
}

src/provider/provider_file_memory.c

Lines changed: 20 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -315,8 +315,12 @@ static void file_finalize(void *provider) {
315315
while (1 == critnib_find(file_provider->mmaps, key, FIND_G, &rkey, &rvalue,
316316
&ref)) {
317317
utils_munmap((void *)rkey, (size_t)rvalue);
318+
assert(ref);
318319
critnib_release(file_provider->mmaps, ref);
319-
critnib_remove(file_provider->mmaps, rkey);
320+
critnib_remove(file_provider->mmaps, rkey, &ref);
321+
if (ref) {
322+
critnib_release(file_provider->mmaps, ref);
323+
}
320324
key = rkey;
321325
}
322326

@@ -695,8 +699,12 @@ static umf_result_t file_allocation_merge_cb(void *provider, void *lowPtr,
695699
return UMF_RESULT_SUCCESS;
696700
}
697701

698-
void *value =
699-
critnib_remove(file_provider->fd_offset_map, (uintptr_t)highPtr);
702+
void *ref_value = NULL;
703+
void *value = critnib_remove(file_provider->fd_offset_map,
704+
(uintptr_t)highPtr, &ref_value);
705+
if (ref_value) {
706+
critnib_release(file_provider->fd_offset_map, ref_value);
707+
}
700708
if (value == NULL) {
701709
LOG_ERR("removing a value from the file descriptor offset map failed "
702710
"(addr=%p)",
@@ -768,6 +776,15 @@ static umf_result_t file_put_ipc_handle(void *provider, void *providerIpcData) {
768776

769777
file_ipc_data_t *file_ipc_data = (file_ipc_data_t *)providerIpcData;
770778

779+
LOG_ERR("file_ipc_data = %p", file_ipc_data);
780+
if (file_ipc_data) {
781+
LOG_ERR("file_ipc_data->path = %s", file_ipc_data->path);
782+
}
783+
LOG_ERR("file_provider = %p", file_provider);
784+
if (file_provider) {
785+
LOG_ERR("file_provider->path = %s", file_provider->path);
786+
}
787+
771788
if (strncmp(file_ipc_data->path, file_provider->path, PATH_MAX)) {
772789
return UMF_RESULT_ERROR_INVALID_ARGUMENT;
773790
}

src/provider/provider_os_memory.c

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1133,7 +1133,11 @@ static umf_result_t os_free(void *provider, void *ptr, size_t size) {
11331133
os_memory_provider_t *os_provider = (os_memory_provider_t *)provider;
11341134

11351135
if (os_provider->fd > 0) {
1136-
critnib_remove(os_provider->fd_offset_map, (uintptr_t)ptr);
1136+
void *ref_value = NULL;
1137+
critnib_remove(os_provider->fd_offset_map, (uintptr_t)ptr, &ref_value);
1138+
if (ref_value) {
1139+
critnib_release(os_provider->fd_offset_map, ref_value);
1140+
}
11371141
}
11381142

11391143
errno = 0;
@@ -1281,8 +1285,12 @@ static umf_result_t os_allocation_merge(void *provider, void *lowPtr,
12811285
return UMF_RESULT_SUCCESS;
12821286
}
12831287

1284-
void *value =
1285-
critnib_remove(os_provider->fd_offset_map, (uintptr_t)highPtr);
1288+
void *ref_value = NULL;
1289+
void *value = critnib_remove(os_provider->fd_offset_map, (uintptr_t)highPtr,
1290+
&ref_value);
1291+
if (ref_value) {
1292+
critnib_release(os_provider->fd_offset_map, ref_value);
1293+
}
12861294
if (value == NULL) {
12871295
LOG_ERR("os_allocation_merge(): removing a value from the file "
12881296
"descriptor offset map failed (addr=%p)",

src/provider/provider_tracking.c

Lines changed: 28 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -324,14 +324,19 @@ static umf_result_t umfMemoryTrackerRemove(umf_memory_tracker_handle_t hTracker,
324324
}
325325

326326
assert(level < MAX_LEVELS_OF_ALLOC_SEGMENT_MAP);
327+
assert(ref_value);
327328
critnib_release(hTracker->alloc_segments_map[level], ref_value);
328-
value = critnib_remove(hTracker->alloc_segments_map[level], (uintptr_t)ptr);
329+
value = critnib_remove(hTracker->alloc_segments_map[level], (uintptr_t)ptr,
330+
&ref_value);
329331
assert(value);
330332

331333
LOG_DEBUG("memory region removed: tracker=%p, level=%i, pool=%p, ptr=%p, "
332334
"size=%zu",
333335
(void *)hTracker, level, (void *)value->pool, ptr, value->size);
334336

337+
assert(ref_value);
338+
critnib_release(hTracker->alloc_segments_map[level], ref_value);
339+
335340
if (parent_value) {
336341
LOG_DEBUG(
337342
"child #%zu removed from memory region: tracker=%p, level=%i, "
@@ -399,7 +404,9 @@ umfMemoryTrackerRemoveIpcSegment(umf_memory_tracker_handle_t hTracker,
399404
const void *ptr) {
400405
assert(ptr);
401406

402-
void *value = critnib_remove(hTracker->ipc_segments_map, (uintptr_t)ptr);
407+
void *ref_value = NULL;
408+
void *value =
409+
critnib_remove(hTracker->ipc_segments_map, (uintptr_t)ptr, &ref_value);
403410

404411
if (!value) {
405412
LOG_ERR("pointer %p not found in the ipc_segments_map", ptr);
@@ -413,6 +420,9 @@ umfMemoryTrackerRemoveIpcSegment(umf_memory_tracker_handle_t hTracker,
413420
(void *)hTracker, ptr, v->size, (void *)v->provider,
414421
(void *)v->ipc_cache_value);
415422

423+
assert(ref_value);
424+
critnib_release(hTracker->ipc_segments_map, ref_value);
425+
416426
return UMF_RESULT_SUCCESS;
417427
}
418428

@@ -788,8 +798,9 @@ static umf_result_t trackingAllocationMerge(void *hProvider, void *lowPtr,
788798
critnib_release(provider->hTracker->alloc_segments_map[highLevel],
789799
ref_highValue);
790800

791-
void *erasedhighValue = critnib_remove(
792-
provider->hTracker->alloc_segments_map[highLevel], (uintptr_t)highPtr);
801+
void *erasedhighValue =
802+
critnib_remove(provider->hTracker->alloc_segments_map[highLevel],
803+
(uintptr_t)highPtr, &ref_highValue);
793804
assert(erasedhighValue == highValue);
794805
(void)erasedhighValue; // unused in the Release build
795806

@@ -798,6 +809,9 @@ static umf_result_t trackingAllocationMerge(void *hProvider, void *lowPtr,
798809
lowLevel, lowPtr, low_children, highPtr, high_children,
799810
totalSize);
800811

812+
critnib_release(provider->hTracker->alloc_segments_map[highLevel],
813+
ref_highValue);
814+
801815
utils_mutex_unlock(&provider->hTracker->splitMergeMutex);
802816

803817
return UMF_RESULT_SUCCESS;
@@ -839,7 +853,8 @@ static umf_result_t trackingFree(void *hProvider, void *ptr, size_t size) {
839853
}
840854
}
841855

842-
void *value = critnib_remove(p->ipcCache, (uintptr_t)ptr);
856+
void *ref_value = NULL;
857+
void *value = critnib_remove(p->ipcCache, (uintptr_t)ptr, &ref_value);
843858
if (value) {
844859
ipc_cache_value_t *cache_value = (ipc_cache_value_t *)value;
845860
ret = umfMemoryProviderPutIPCHandle(p->hUpstream,
@@ -849,6 +864,8 @@ static umf_result_t trackingFree(void *hProvider, void *ptr, size_t size) {
849864
"size=%zu, ret = %d",
850865
ptr, size, ret);
851866
}
867+
assert(ref_value);
868+
critnib_release(p->ipcCache, ref_value);
852869
}
853870

854871
ret = umfMemoryProviderFree(p->hUpstream, ptr, size);
@@ -1067,6 +1084,12 @@ static umf_result_t trackingGetIpcHandle(void *provider, const void *ptr,
10671084
}
10681085
}
10691086
}
1087+
1088+
if (!cached && ref_value) {
1089+
critnib_release(p->ipcCache, &ref_value);
1090+
ref_value = NULL;
1091+
}
1092+
10701093
} while (!cached);
10711094

10721095
memcpy(providerIpcData, cache_value->providerIpcData,

0 commit comments

Comments
 (0)