add support to visit _all_ abandoned segment blocks per sub-process, upstream for python/cpython#114133

This commit is contained in:
daanx 2024-06-02 09:41:12 -07:00
parent 8f874555d5
commit 855e3b2549
3 changed files with 121 additions and 45 deletions

View file

@ -404,7 +404,7 @@ typedef struct mi_segment_s {
bool was_reclaimed; // true if it was reclaimed (used to limit on-free reclamation)
size_t abandoned; // abandoned pages (i.e. the original owning thread stopped) (`abandoned <= used`)
size_t abandoned_visits; // count how often this segment is visited in the abandoned list (to force reclaim if it is too long)
size_t abandoned_visits; // count how often this segment is visited for reclaiming (to force reclaim if it is too long)
size_t used; // count of pages in use (`used <= capacity`)
size_t capacity; // count of available pages (`#free + used`)
@ -412,6 +412,9 @@ typedef struct mi_segment_s {
uintptr_t cookie; // verify addresses in secure mode: `_mi_ptr_cookie(segment) == segment->cookie`
mi_subproc_t* subproc; // segment belongs to sub process
struct mi_segment_s* abandoned_os_next; // only used for abandoned segments outside arena's, and only if `mi_option_visit_abandoned` is enabled
struct mi_segment_s* abandoned_os_prev;
// layout like this to optimize access in `mi_free`
_Atomic(mi_threadid_t) thread_id; // unique id of the thread owning this segment
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);
struct mi_subproc_s {
_Atomic(size_t) abandoned_count; // count of abandoned segments for this sup-process
mi_lock_t abandoned_os_lock; // lock for the abandoned segments outside of arena's
mi_segment_t* abandoned_os_list; // doubly-linked list of abandoned segments outside of arena's (in OS allocated memory)
mi_memid_t memid; // provenance
};

View file

@ -757,17 +757,34 @@ bool _mi_arena_contains(const void* p) {
// sets the thread_id.
bool _mi_arena_segment_clear_abandoned(mi_segment_t* segment )
{
if (segment->memid.memkind != MI_MEM_ARENA) {
// not in an arena, consider it un-abandoned now.
// but we need to still claim it atomically -- we use the thread_id for that.
if mi_unlikely(segment->memid.memkind != MI_MEM_ARENA) {
// not in an arena
// if abandoned visiting is allowed, we need to take a lock on the abandoned os list
bool has_lock = false;
if (mi_option_is_enabled(mi_option_visit_abandoned)) {
has_lock = mi_lock_try_acquire(&segment->subproc->abandoned_os_lock);
if (!has_lock) {
return false; // failed to acquire the lock, we just give up
}
}
// abandon it, but we need to still claim it atomically -- we use the thread_id for that.
bool reclaimed = false;
size_t expected = 0;
if (mi_atomic_cas_strong_acq_rel(&segment->thread_id, &expected, _mi_thread_id())) {
// reclaim
mi_atomic_decrement_relaxed(&segment->subproc->abandoned_count);
return true;
}
else {
return false;
reclaimed = true;
// and remove from the abandoned os list (if needed)
mi_segment_t* const next = segment->abandoned_os_next;
mi_segment_t* const prev = segment->abandoned_os_prev;
if (prev != NULL) { prev->abandoned_os_next = next; }
else { segment->subproc->abandoned_os_list = next; }
if (next != NULL) { next->abandoned_os_prev = prev; }
segment->abandoned_os_next = NULL;
segment->abandoned_os_prev = NULL;
}
if (has_lock) { mi_lock_release(&segment->subproc->abandoned_os_lock); }
return reclaimed;
}
// arena segment: use the blocks_abandoned bitmap.
size_t arena_idx;
@ -794,12 +811,30 @@ void _mi_arena_segment_mark_abandoned(mi_segment_t* segment)
{
mi_atomic_store_release(&segment->thread_id, 0);
mi_assert_internal(segment->used == segment->abandoned);
if (segment->memid.memkind != MI_MEM_ARENA) {
// not in an arena; count it as abandoned and return
if mi_unlikely(segment->memid.memkind != MI_MEM_ARENA) {
// not in an arena; count it as abandoned and return (these can be reclaimed on a `free`)
mi_atomic_increment_relaxed(&segment->subproc->abandoned_count);
// if abandoned visiting is allowed, we need to take a lock on the abandoned os list to insert it
if (mi_option_is_enabled(mi_option_visit_abandoned)) {
if (!mi_lock_acquire(&segment->subproc->abandoned_os_lock)) {
_mi_error_message(EFAULT, "internal error: failed to acquire the abandoned (os) segment lock to mark abandonment");
}
else {
// push on the front of the list
mi_segment_t* next = segment->subproc->abandoned_os_list;
mi_assert_internal(next == NULL || next->abandoned_os_prev == NULL);
mi_assert_internal(segment->abandoned_os_prev == NULL);
mi_assert_internal(segment->abandoned_os_next == NULL);
if (next != NULL) { next->abandoned_os_prev = segment; }
segment->abandoned_os_prev = NULL;
segment->abandoned_os_next = next;
segment->subproc->abandoned_os_list = segment;
mi_lock_release(&segment->subproc->abandoned_os_lock);
}
}
return;
}
// segment is in an arena
// segment is in an arena, mark it in the arena `blocks_abandoned` bitmap
size_t arena_idx;
size_t bitmap_idx;
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
current->subproc = subproc;
}
static mi_segment_t* mi_arena_segment_clear_abandoned_at(mi_arena_t* arena, mi_subproc_t* subproc, mi_bitmap_index_t bitmap_idx) {
// try to reclaim an abandoned segment in the arena atomically
if (!_mi_bitmap_unclaim(arena->blocks_abandoned, arena->field_count, 1, bitmap_idx)) return NULL;
mi_assert_internal(_mi_bitmap_is_claimed(arena->blocks_inuse, arena->field_count, 1, bitmap_idx));
mi_segment_t* segment = (mi_segment_t*)mi_arena_block_start(arena, bitmap_idx);
mi_assert_internal(mi_atomic_load_relaxed(&segment->thread_id) == 0);
// check that the segment belongs to our sub-process
// note: this is the reason we need a lock in the case abandoned visiting is enabled.
// without the lock an abandoned visit may otherwise fail to visit all segments.
// for regular reclaim it is fine to miss one sometimes so without abandoned visiting we don't need the arena lock.
if (segment->subproc != subproc) {
// it is from another subprocess, re-mark it and continue searching
const bool was_zero = _mi_bitmap_claim(arena->blocks_abandoned, arena->field_count, 1, bitmap_idx, NULL);
mi_assert_internal(was_zero); MI_UNUSED(was_zero);
return NULL;
}
else {
// success, we unabandoned a segment in our sub-process
mi_atomic_decrement_relaxed(&subproc->abandoned_count);
return segment;
}
}
// reclaim abandoned segments
// this does not set the thread id (so it appears as still abandoned)
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
has_lock = (visit_all ? mi_lock_acquire(&arena->abandoned_visit_lock) : mi_lock_try_acquire(&arena->abandoned_visit_lock));
if (!has_lock) {
if (visit_all) {
_mi_error_message(EINVAL, "failed to visit all abandoned segments due to failure to acquire the visitor lock");
_mi_error_message(EFAULT, "internal error: failed to visit all abandoned segments due to failure to acquire the visitor lock");
}
// skip to next arena
break;
@ -860,27 +918,11 @@ mi_segment_t* _mi_arena_segment_clear_abandoned_next(mi_arena_field_cursor_t* pr
// pre-check if the bit is set
size_t mask = ((size_t)1 << bit_idx);
if mi_unlikely((field & mask) == mask) {
mi_bitmap_index_t bitmap_idx = mi_bitmap_index_create(field_idx, bit_idx);
// try to reclaim it atomically
if (_mi_bitmap_unclaim(arena->blocks_abandoned, arena->field_count, 1, bitmap_idx)) {
mi_assert_internal(_mi_bitmap_is_claimed(arena->blocks_inuse, arena->field_count, 1, bitmap_idx));
mi_segment_t* segment = (mi_segment_t*)mi_arena_block_start(arena, bitmap_idx);
mi_assert_internal(mi_atomic_load_relaxed(&segment->thread_id) == 0);
// check that the segment belongs to our sub-process
// note: this is the reason we need a lock in the case abandoned visiting is enabled.
// without the lock an abandoned visit may otherwise fail to visit all segments.
// for regular reclaim it is fine to miss one sometimes so without abandoned visiting we don't need the arena lock.
if (segment->subproc != previous->subproc) {
// it is from another subprocess, re-mark it and continue searching
const bool was_zero = _mi_bitmap_claim(arena->blocks_abandoned, arena->field_count, 1, bitmap_idx, NULL);
mi_assert_internal(was_zero);
}
else {
// success, we unabandoned a segment in our sub-process
mi_atomic_decrement_relaxed(&previous->subproc->abandoned_count);
const mi_bitmap_index_t bitmap_idx = mi_bitmap_index_create(field_idx, bit_idx);
mi_segment_t* const segment = mi_arena_segment_clear_abandoned_at(arena, previous->subproc, bitmap_idx);
if (segment != NULL) {
previous->bitmap_idx = bitmap_idx;
previous->count = count;
//mi_assert_internal(arena->blocks_committed == NULL || _mi_bitmap_is_claimed(arena->blocks_committed, arena->field_count, 1, bitmap_idx));
if (has_lock) { mi_lock_release(&arena->abandoned_visit_lock); }
return segment;
@ -889,7 +931,6 @@ mi_segment_t* _mi_arena_segment_clear_abandoned_next(mi_arena_field_cursor_t* pr
}
}
}
}
if (has_lock) { mi_lock_release(&arena->abandoned_visit_lock); }
}
}
@ -910,16 +951,35 @@ static bool mi_arena_visit_abandoned_blocks(mi_subproc_t* subproc, int heap_tag,
return true;
}
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) {
if (!mi_lock_acquire(&subproc->abandoned_os_lock)) {
_mi_error_message(EFAULT, "internal error: failed to acquire abandoned (OS) segment lock");
return false;
}
bool all_visited = true;
for (mi_segment_t* segment = subproc->abandoned_os_list; segment != NULL; segment = segment->abandoned_os_next) {
if (!_mi_segment_visit_blocks(segment, heap_tag, visit_blocks, visitor, arg)) {
all_visited = false;
break;
}
}
mi_lock_release(&subproc->abandoned_os_lock);
return all_visited;
}
bool mi_abandoned_visit_blocks(mi_subproc_id_t subproc_id, int heap_tag, bool visit_blocks, mi_block_visit_fun* visitor, void* arg) {
// (unfortunately) the visit_abandoned option must be enabled from the start.
// This is to avoid taking locks if abandoned list visiting is not required (as for most programs)
if (!mi_option_is_enabled(mi_option_visit_abandoned)) {
mi_assert(false);
_mi_error_message(EINVAL, "internal error: can only visit abandoned blocks when MIMALLOC_VISIT_ABANDONED=ON");
_mi_error_message(EFAULT, "internal error: can only visit abandoned blocks when MIMALLOC_VISIT_ABANDONED=ON");
return false;
}
mi_subproc_t* const subproc = _mi_subproc_from_id(subproc_id);
// visit abandoned segments in the arena's
return mi_arena_visit_abandoned_blocks(_mi_subproc_from_id(subproc_id), heap_tag, visit_blocks, visitor, arg);
if (!mi_arena_visit_abandoned_blocks(subproc, heap_tag, visit_blocks, visitor, arg)) return false;
// and visit abandoned segments outside arena's (in OS allocated memory)
if (!mi_subproc_visit_abandoned_os_blocks(subproc, heap_tag, visit_blocks, visitor, arg)) return false;
return true;
}

View file

@ -172,6 +172,7 @@ static void mi_heap_main_init(void) {
_mi_heap_main.cookie = _mi_heap_random_next(&_mi_heap_main);
_mi_heap_main.keys[0] = _mi_heap_random_next(&_mi_heap_main);
_mi_heap_main.keys[1] = _mi_heap_random_next(&_mi_heap_main);
mi_lock_init(&mi_subproc_default.abandoned_os_lock);
}
}
@ -185,8 +186,6 @@ mi_heap_t* _mi_heap_main_get(void) {
Sub process
----------------------------------------------------------- */
static mi_decl_cache_align _Atomic(uintptr_t) mi_subproc_count;
mi_subproc_id_t mi_subproc_main(void) {
return NULL;
}
@ -195,8 +194,9 @@ mi_subproc_id_t mi_subproc_new(void) {
mi_memid_t memid = _mi_memid_none();
mi_subproc_t* subproc = (mi_subproc_t*)_mi_arena_meta_zalloc(sizeof(mi_subproc_t), &memid);
if (subproc == NULL) return NULL;
mi_atomic_increment_relaxed(&mi_subproc_count);
subproc->memid = memid;
subproc->abandoned_os_list = NULL;
mi_lock_init(&subproc->abandoned_os_lock);
return subproc;
}
@ -207,8 +207,19 @@ mi_subproc_t* _mi_subproc_from_id(mi_subproc_id_t subproc_id) {
void mi_subproc_delete(mi_subproc_id_t subproc_id) {
if (subproc_id == NULL) return;
mi_subproc_t* subproc = _mi_subproc_from_id(subproc_id);
// check if there are no abandoned segments still..
bool safe_to_delete = false;
if (mi_lock_acquire(&subproc->abandoned_os_lock)) {
if (subproc->abandoned_os_list == NULL) {
safe_to_delete = true;
}
mi_lock_release(&subproc->abandoned_os_lock);
}
if (!safe_to_delete) return;
// safe to release
// todo: should we refcount subprocesses?
mi_lock_done(&subproc->abandoned_os_lock);
_mi_arena_meta_free(subproc, subproc->memid, sizeof(mi_subproc_t));
mi_atomic_decrement_relaxed(&mi_subproc_count);
}
void mi_subproc_add_current_thread(mi_subproc_id_t subproc_id) {