Skip to content

Commit 855e3b2

Browse files
committedJun 2, 2024
add support to visit _all_ abandoned segment blocks per sub-process, upstream for python/cpython#114133
1 parent 8f87455 commit 855e3b2

File tree

3 files changed

+121
-45
lines changed

3 files changed

+121
-45
lines changed
 

‎include/mimalloc/types.h

+6-1
Original file line numberDiff line numberDiff line change
@@ -404,14 +404,17 @@ typedef struct mi_segment_s {
404404
bool was_reclaimed; // true if it was reclaimed (used to limit on-free reclamation)
405405

406406
size_t abandoned; // abandoned pages (i.e. the original owning thread stopped) (`abandoned <= used`)
407-
size_t abandoned_visits; // count how often this segment is visited in the abandoned list (to force reclaim if it is too long)
407+
size_t abandoned_visits; // count how often this segment is visited for reclaiming (to force reclaim if it is too long)
408408

409409
size_t used; // count of pages in use (`used <= capacity`)
410410
size_t capacity; // count of available pages (`#free + used`)
411411
size_t segment_info_size;// space we are using from the first page for segment meta-data and possible guard pages.
412412
uintptr_t cookie; // verify addresses in secure mode: `_mi_ptr_cookie(segment) == segment->cookie`
413413
mi_subproc_t* subproc; // segment belongs to sub process
414414

415+
struct mi_segment_s* abandoned_os_next; // only used for abandoned segments outside arena's, and only if `mi_option_visit_abandoned` is enabled
416+
struct mi_segment_s* abandoned_os_prev;
417+
415418
// layout like this to optimize access in `mi_free`
416419
_Atomic(mi_threadid_t) thread_id; // unique id of the thread owning this segment
417420
size_t page_shift; // `1 << page_shift` == the page sizes == `page->block_size * page->reserved` (unless the first page, then `-segment_info_size`).
@@ -609,6 +612,8 @@ void _mi_stat_counter_increase(mi_stat_counter_t* stat, size_t amount);
609612

610613
struct mi_subproc_s {
611614
_Atomic(size_t) abandoned_count; // count of abandoned segments for this sup-process
615+
mi_lock_t abandoned_os_lock; // lock for the abandoned segments outside of arena's
616+
mi_segment_t* abandoned_os_list; // doubly-linked list of abandoned segments outside of arena's (in OS allocated memory)
612617
mi_memid_t memid; // provenance
613618
};
614619

‎src/arena.c

+99-39
Original file line numberDiff line numberDiff line change
@@ -757,17 +757,34 @@ bool _mi_arena_contains(const void* p) {
757757
// sets the thread_id.
758758
bool _mi_arena_segment_clear_abandoned(mi_segment_t* segment )
759759
{
760-
if (segment->memid.memkind != MI_MEM_ARENA) {
761-
// not in an arena, consider it un-abandoned now.
762-
// but we need to still claim it atomically -- we use the thread_id for that.
760+
if mi_unlikely(segment->memid.memkind != MI_MEM_ARENA) {
761+
// not in an arena
762+
// if abandoned visiting is allowed, we need to take a lock on the abandoned os list
763+
bool has_lock = false;
764+
if (mi_option_is_enabled(mi_option_visit_abandoned)) {
765+
has_lock = mi_lock_try_acquire(&segment->subproc->abandoned_os_lock);
766+
if (!has_lock) {
767+
return false; // failed to acquire the lock, we just give up
768+
}
769+
}
770+
// abandon it, but we need to still claim it atomically -- we use the thread_id for that.
771+
bool reclaimed = false;
763772
size_t expected = 0;
764773
if (mi_atomic_cas_strong_acq_rel(&segment->thread_id, &expected, _mi_thread_id())) {
774+
// reclaim
765775
mi_atomic_decrement_relaxed(&segment->subproc->abandoned_count);
766-
return true;
767-
}
768-
else {
769-
return false;
776+
reclaimed = true;
777+
// and remove from the abandoned os list (if needed)
778+
mi_segment_t* const next = segment->abandoned_os_next;
779+
mi_segment_t* const prev = segment->abandoned_os_prev;
780+
if (prev != NULL) { prev->abandoned_os_next = next; }
781+
else { segment->subproc->abandoned_os_list = next; }
782+
if (next != NULL) { next->abandoned_os_prev = prev; }
783+
segment->abandoned_os_next = NULL;
784+
segment->abandoned_os_prev = NULL;
770785
}
786+
if (has_lock) { mi_lock_release(&segment->subproc->abandoned_os_lock); }
787+
return reclaimed;
771788
}
772789
// arena segment: use the blocks_abandoned bitmap.
773790
size_t arena_idx;
@@ -794,12 +811,30 @@ void _mi_arena_segment_mark_abandoned(mi_segment_t* segment)
794811
{
795812
mi_atomic_store_release(&segment->thread_id, 0);
796813
mi_assert_internal(segment->used == segment->abandoned);
797-
if (segment->memid.memkind != MI_MEM_ARENA) {
798-
// not in an arena; count it as abandoned and return
814+
if mi_unlikely(segment->memid.memkind != MI_MEM_ARENA) {
815+
// not in an arena; count it as abandoned and return (these can be reclaimed on a `free`)
799816
mi_atomic_increment_relaxed(&segment->subproc->abandoned_count);
817+
// if abandoned visiting is allowed, we need to take a lock on the abandoned os list to insert it
818+
if (mi_option_is_enabled(mi_option_visit_abandoned)) {
819+
if (!mi_lock_acquire(&segment->subproc->abandoned_os_lock)) {
820+
_mi_error_message(EFAULT, "internal error: failed to acquire the abandoned (os) segment lock to mark abandonment");
821+
}
822+
else {
823+
// push on the front of the list
824+
mi_segment_t* next = segment->subproc->abandoned_os_list;
825+
mi_assert_internal(next == NULL || next->abandoned_os_prev == NULL);
826+
mi_assert_internal(segment->abandoned_os_prev == NULL);
827+
mi_assert_internal(segment->abandoned_os_next == NULL);
828+
if (next != NULL) { next->abandoned_os_prev = segment; }
829+
segment->abandoned_os_prev = NULL;
830+
segment->abandoned_os_next = next;
831+
segment->subproc->abandoned_os_list = segment;
832+
mi_lock_release(&segment->subproc->abandoned_os_lock);
833+
}
834+
}
800835
return;
801836
}
802-
// segment is in an arena
837+
// segment is in an arena, mark it in the arena `blocks_abandoned` bitmap
803838
size_t arena_idx;
804839
size_t bitmap_idx;
805840
mi_arena_memid_indices(segment->memid, &arena_idx, &bitmap_idx);
@@ -822,6 +857,29 @@ void _mi_arena_field_cursor_init(mi_heap_t* heap, mi_subproc_t* subproc, mi_aren
822857
current->subproc = subproc;
823858
}
824859

860+
static mi_segment_t* mi_arena_segment_clear_abandoned_at(mi_arena_t* arena, mi_subproc_t* subproc, mi_bitmap_index_t bitmap_idx) {
861+
// try to reclaim an abandoned segment in the arena atomically
862+
if (!_mi_bitmap_unclaim(arena->blocks_abandoned, arena->field_count, 1, bitmap_idx)) return NULL;
863+
mi_assert_internal(_mi_bitmap_is_claimed(arena->blocks_inuse, arena->field_count, 1, bitmap_idx));
864+
mi_segment_t* segment = (mi_segment_t*)mi_arena_block_start(arena, bitmap_idx);
865+
mi_assert_internal(mi_atomic_load_relaxed(&segment->thread_id) == 0);
866+
// check that the segment belongs to our sub-process
867+
// note: this is the reason we need a lock in the case abandoned visiting is enabled.
868+
// without the lock an abandoned visit may otherwise fail to visit all segments.
869+
// for regular reclaim it is fine to miss one sometimes so without abandoned visiting we don't need the arena lock.
870+
if (segment->subproc != subproc) {
871+
// it is from another subprocess, re-mark it and continue searching
872+
const bool was_zero = _mi_bitmap_claim(arena->blocks_abandoned, arena->field_count, 1, bitmap_idx, NULL);
873+
mi_assert_internal(was_zero); MI_UNUSED(was_zero);
874+
return NULL;
875+
}
876+
else {
877+
// success, we unabandoned a segment in our sub-process
878+
mi_atomic_decrement_relaxed(&subproc->abandoned_count);
879+
return segment;
880+
}
881+
}
882+
825883
// reclaim abandoned segments
826884
// this does not set the thread id (so it appears as still abandoned)
827885
mi_segment_t* _mi_arena_segment_clear_abandoned_next(mi_arena_field_cursor_t* previous, bool visit_all )
@@ -848,7 +906,7 @@ mi_segment_t* _mi_arena_segment_clear_abandoned_next(mi_arena_field_cursor_t* pr
848906
has_lock = (visit_all ? mi_lock_acquire(&arena->abandoned_visit_lock) : mi_lock_try_acquire(&arena->abandoned_visit_lock));
849907
if (!has_lock) {
850908
if (visit_all) {
851-
_mi_error_message(EINVAL, "failed to visit all abandoned segments due to failure to acquire the visitor lock");
909+
_mi_error_message(EFAULT, "internal error: failed to visit all abandoned segments due to failure to acquire the visitor lock");
852910
}
853911
// skip to next arena
854912
break;
@@ -860,31 +918,14 @@ mi_segment_t* _mi_arena_segment_clear_abandoned_next(mi_arena_field_cursor_t* pr
860918
// pre-check if the bit is set
861919
size_t mask = ((size_t)1 << bit_idx);
862920
if mi_unlikely((field & mask) == mask) {
863-
mi_bitmap_index_t bitmap_idx = mi_bitmap_index_create(field_idx, bit_idx);
864-
// try to reclaim it atomically
865-
if (_mi_bitmap_unclaim(arena->blocks_abandoned, arena->field_count, 1, bitmap_idx)) {
866-
mi_assert_internal(_mi_bitmap_is_claimed(arena->blocks_inuse, arena->field_count, 1, bitmap_idx));
867-
mi_segment_t* segment = (mi_segment_t*)mi_arena_block_start(arena, bitmap_idx);
868-
mi_assert_internal(mi_atomic_load_relaxed(&segment->thread_id) == 0);
869-
// check that the segment belongs to our sub-process
870-
// note: this is the reason we need a lock in the case abandoned visiting is enabled.
871-
// without the lock an abandoned visit may otherwise fail to visit all segments.
872-
// for regular reclaim it is fine to miss one sometimes so without abandoned visiting we don't need the arena lock.
873-
if (segment->subproc != previous->subproc) {
874-
// it is from another subprocess, re-mark it and continue searching
875-
const bool was_zero = _mi_bitmap_claim(arena->blocks_abandoned, arena->field_count, 1, bitmap_idx, NULL);
876-
mi_assert_internal(was_zero);
877-
}
878-
else {
879-
// success, we unabandoned a segment in our sub-process
880-
mi_atomic_decrement_relaxed(&previous->subproc->abandoned_count);
881-
previous->bitmap_idx = bitmap_idx;
882-
previous->count = count;
883-
884-
//mi_assert_internal(arena->blocks_committed == NULL || _mi_bitmap_is_claimed(arena->blocks_committed, arena->field_count, 1, bitmap_idx));
885-
if (has_lock) { mi_lock_release(&arena->abandoned_visit_lock); }
886-
return segment;
887-
}
921+
const mi_bitmap_index_t bitmap_idx = mi_bitmap_index_create(field_idx, bit_idx);
922+
mi_segment_t* const segment = mi_arena_segment_clear_abandoned_at(arena, previous->subproc, bitmap_idx);
923+
if (segment != NULL) {
924+
previous->bitmap_idx = bitmap_idx;
925+
previous->count = count;
926+
//mi_assert_internal(arena->blocks_committed == NULL || _mi_bitmap_is_claimed(arena->blocks_committed, arena->field_count, 1, bitmap_idx));
927+
if (has_lock) { mi_lock_release(&arena->abandoned_visit_lock); }
928+
return segment;
888929
}
889930
}
890931
}
@@ -910,16 +951,35 @@ static bool mi_arena_visit_abandoned_blocks(mi_subproc_t* subproc, int heap_tag,
910951
return true;
911952
}
912953

954+
static bool mi_subproc_visit_abandoned_os_blocks(mi_subproc_t* subproc, int heap_tag, bool visit_blocks, mi_block_visit_fun* visitor, void* arg) {
955+
if (!mi_lock_acquire(&subproc->abandoned_os_lock)) {
956+
_mi_error_message(EFAULT, "internal error: failed to acquire abandoned (OS) segment lock");
957+
return false;
958+
}
959+
bool all_visited = true;
960+
for (mi_segment_t* segment = subproc->abandoned_os_list; segment != NULL; segment = segment->abandoned_os_next) {
961+
if (!_mi_segment_visit_blocks(segment, heap_tag, visit_blocks, visitor, arg)) {
962+
all_visited = false;
963+
break;
964+
}
965+
}
966+
mi_lock_release(&subproc->abandoned_os_lock);
967+
return all_visited;
968+
}
969+
913970
bool mi_abandoned_visit_blocks(mi_subproc_id_t subproc_id, int heap_tag, bool visit_blocks, mi_block_visit_fun* visitor, void* arg) {
914971
// (unfortunately) the visit_abandoned option must be enabled from the start.
915972
// This is to avoid taking locks if abandoned list visiting is not required (as for most programs)
916973
if (!mi_option_is_enabled(mi_option_visit_abandoned)) {
917-
mi_assert(false);
918-
_mi_error_message(EINVAL, "internal error: can only visit abandoned blocks when MIMALLOC_VISIT_ABANDONED=ON");
974+
_mi_error_message(EFAULT, "internal error: can only visit abandoned blocks when MIMALLOC_VISIT_ABANDONED=ON");
919975
return false;
920976
}
977+
mi_subproc_t* const subproc = _mi_subproc_from_id(subproc_id);
921978
// visit abandoned segments in the arena's
922-
return mi_arena_visit_abandoned_blocks(_mi_subproc_from_id(subproc_id), heap_tag, visit_blocks, visitor, arg);
979+
if (!mi_arena_visit_abandoned_blocks(subproc, heap_tag, visit_blocks, visitor, arg)) return false;
980+
// and visit abandoned segments outside arena's (in OS allocated memory)
981+
if (!mi_subproc_visit_abandoned_os_blocks(subproc, heap_tag, visit_blocks, visitor, arg)) return false;
982+
return true;
923983
}
924984

925985

‎src/init.c

+16-5
Original file line numberDiff line numberDiff line change
@@ -171,7 +171,8 @@ static void mi_heap_main_init(void) {
171171
#endif
172172
_mi_heap_main.cookie = _mi_heap_random_next(&_mi_heap_main);
173173
_mi_heap_main.keys[0] = _mi_heap_random_next(&_mi_heap_main);
174-
_mi_heap_main.keys[1] = _mi_heap_random_next(&_mi_heap_main);
174+
_mi_heap_main.keys[1] = _mi_heap_random_next(&_mi_heap_main);
175+
mi_lock_init(&mi_subproc_default.abandoned_os_lock);
175176
}
176177
}
177178

@@ -185,8 +186,6 @@ mi_heap_t* _mi_heap_main_get(void) {
185186
Sub process
186187
----------------------------------------------------------- */
187188

188-
static mi_decl_cache_align _Atomic(uintptr_t) mi_subproc_count;
189-
190189
mi_subproc_id_t mi_subproc_main(void) {
191190
return NULL;
192191
}
@@ -195,8 +194,9 @@ mi_subproc_id_t mi_subproc_new(void) {
195194
mi_memid_t memid = _mi_memid_none();
196195
mi_subproc_t* subproc = (mi_subproc_t*)_mi_arena_meta_zalloc(sizeof(mi_subproc_t), &memid);
197196
if (subproc == NULL) return NULL;
198-
mi_atomic_increment_relaxed(&mi_subproc_count);
199197
subproc->memid = memid;
198+
subproc->abandoned_os_list = NULL;
199+
mi_lock_init(&subproc->abandoned_os_lock);
200200
return subproc;
201201
}
202202

@@ -207,8 +207,19 @@ mi_subproc_t* _mi_subproc_from_id(mi_subproc_id_t subproc_id) {
207207
void mi_subproc_delete(mi_subproc_id_t subproc_id) {
208208
if (subproc_id == NULL) return;
209209
mi_subproc_t* subproc = _mi_subproc_from_id(subproc_id);
210+
// check if there are no abandoned segments still..
211+
bool safe_to_delete = false;
212+
if (mi_lock_acquire(&subproc->abandoned_os_lock)) {
213+
if (subproc->abandoned_os_list == NULL) {
214+
safe_to_delete = true;
215+
}
216+
mi_lock_release(&subproc->abandoned_os_lock);
217+
}
218+
if (!safe_to_delete) return;
219+
// safe to release
220+
// todo: should we refcount subprocesses?
221+
mi_lock_done(&subproc->abandoned_os_lock);
210222
_mi_arena_meta_free(subproc, subproc->memid, sizeof(mi_subproc_t));
211-
mi_atomic_decrement_relaxed(&mi_subproc_count);
212223
}
213224

214225
void mi_subproc_add_current_thread(mi_subproc_id_t subproc_id) {

0 commit comments

Comments
 (0)