check heaptag on abandonded page allocation

This commit is contained in:
daanx 2024-12-07 15:13:17 -08:00
parent 6b52b19e3b
commit bf42759d97
4 changed files with 24 additions and 14 deletions

View file

@ -237,6 +237,8 @@ typedef uintptr_t mi_thread_free_t;
// Sub processes are used to keep memory separate between them (e.g. multiple interpreters in CPython)
typedef struct mi_subproc_s mi_subproc_t;
// A heap can serve only specific objects signified by its heap tag (e.g. various object types in CPython)
typedef uint8_t mi_heaptag_t;
// A page contains blocks of one specific size (`block_size`).
// Each page has three list of free blocks:
@ -280,7 +282,7 @@ typedef struct mi_page_s {
size_t block_size; // size available in each block (always `>0`)
uint8_t* page_start; // start of the blocks
uint8_t heap_tag; // tag of the owning heap, used to separate heaps by object type
mi_heaptag_t heap_tag; // tag of the owning heap, used to separate heaps by object type
bool free_is_zero; // `true` if the blocks in the free list are zero initialized
// padding
#if (MI_ENCODE_FREELIST || MI_PADDING)
@ -411,7 +413,16 @@ struct mi_heap_s {
mi_page_queue_t pages[MI_BIN_FULL + 1]; // queue of pages for each size class (or "bin")
};
// ------------------------------------------------------
// Arena's
// These are large reserved areas of memory allocated from
// the OS that are managed by mimalloc to efficiently
// allocate MI_SLICE_SIZE slices of memory for the
// mimalloc pages.
// ------------------------------------------------------
// A large memory arena where pages are allocated in.
typedef struct mi_arena_s mi_arena_t;
// ------------------------------------------------------
// Debug

View file

@ -479,11 +479,9 @@ void* _mi_arena_alloc(size_t size, bool commit, bool allow_large, mi_arena_id_t
Arena page allocation
----------------------------------------------------------- */
static bool mi_arena_try_claim_abandoned(size_t slice_index, void* arg1, void* arg2, bool* keep_abandoned) {
static bool mi_arena_try_claim_abandoned(size_t slice_index, mi_arena_t* arena, mi_subproc_t* subproc, mi_heaptag_t heap_tag, bool* keep_abandoned) {
// found an abandoned page of the right size
mi_arena_t* const arena = (mi_arena_t*)arg1;
mi_subproc_t* const subproc = (mi_subproc_t*)arg2;
mi_page_t* const page = (mi_page_t*)mi_arena_slice_start(arena, slice_index);
mi_page_t* const page = (mi_page_t*)mi_arena_slice_start(arena, slice_index);
// can we claim ownership?
if (!mi_page_try_claim_ownership(page)) {
// there was a concurrent free ..
@ -493,8 +491,9 @@ static bool mi_arena_try_claim_abandoned(size_t slice_index, void* arg1, void* a
*keep_abandoned = true;
return false;
}
if (subproc != page->subproc) {
// wrong sub-process.. we need to unown again
if (subproc != page->subproc || heap_tag != page->heap_tag) {
// wrong sub-process or heap_tag.. we need to unown again
// note: this normally never happens unless subprocesses/heaptags are actually used.
// (an unown might free the page, and depending on that we can keep it in the abandoned map or not)
// note: a minor wrinkle: the page will still be mapped but the abandoned map entry is (temporarily) clear at this point.
// so we cannot check in `mi_arena_free` for this invariant to hold.
@ -507,7 +506,7 @@ static bool mi_arena_try_claim_abandoned(size_t slice_index, void* arg1, void* a
return true;
}
static mi_page_t* mi_arena_page_try_find_abandoned(size_t slice_count, size_t block_size, mi_arena_id_t req_arena_id, mi_tld_t* tld)
static mi_page_t* mi_arena_page_try_find_abandoned(size_t slice_count, size_t block_size, mi_arena_id_t req_arena_id, mi_heaptag_t heaptag, mi_tld_t* tld)
{
MI_UNUSED(slice_count);
const size_t bin = _mi_bin(block_size);
@ -525,7 +524,7 @@ static mi_page_t* mi_arena_page_try_find_abandoned(size_t slice_count, size_t bl
size_t slice_index;
mi_bitmap_t* const bitmap = arena->pages_abandoned[bin];
if (mi_bitmap_try_find_and_claim(bitmap, tseq, &slice_index, &mi_arena_try_claim_abandoned, arena, subproc)) {
if (mi_bitmap_try_find_and_claim(bitmap, tseq, &slice_index, &mi_arena_try_claim_abandoned, arena, subproc, heaptag)) {
// found an abandoned page of the right size
// and claimed ownership.
mi_page_t* page = (mi_page_t*)mi_arena_slice_start(arena, slice_index);
@ -632,7 +631,7 @@ static mi_page_t* mi_arena_page_allocN(mi_heap_t* heap, size_t slice_count, size
mi_tld_t* const tld = heap->tld;
// 1. look for an abandoned page
mi_page_t* page = mi_arena_page_try_find_abandoned(slice_count, block_size, req_arena_id, tld);
mi_page_t* page = mi_arena_page_try_find_abandoned(slice_count, block_size, req_arena_id, heap->tag, tld);
if (page != NULL) {
return page; // return as abandoned
}

View file

@ -1165,7 +1165,7 @@ mi_decl_nodiscard bool mi_bitmap_try_find_and_clearN(mi_bitmap_t* bitmap, size_t
// Find a set bit in the bitmap and try to atomically clear it and claim it.
// (Used to find pages in the pages_abandoned bitmaps.)
mi_decl_nodiscard bool mi_bitmap_try_find_and_claim(mi_bitmap_t* bitmap, size_t tseq, size_t* pidx,
mi_claim_fun_t* claim, void* arg1, void* arg2)
mi_claim_fun_t* claim, mi_arena_t* arena, mi_subproc_t* subproc, mi_heaptag_t heap_tag )
{
mi_bitmap_forall_chunks(bitmap, tseq, chunk_idx)
{
@ -1174,7 +1174,7 @@ mi_decl_nodiscard bool mi_bitmap_try_find_and_claim(mi_bitmap_t* bitmap, size_t
const size_t slice_index = (chunk_idx * MI_BCHUNK_BITS) + cidx;
mi_assert_internal(slice_index < mi_bitmap_max_bits(bitmap));
bool keep_set = true;
if ((*claim)(slice_index, arg1, arg2, &keep_set)) {
if ((*claim)(slice_index, arena, subproc, heap_tag, &keep_set)) {
// success!
mi_assert_internal(!keep_set);
*pidx = slice_index;

View file

@ -185,10 +185,10 @@ static inline bool mi_bitmap_try_clearN(mi_bitmap_t* bitmap, size_t idx, size_t
// Returns true on success, and in that case sets the index: `0 <= *pidx <= MI_BITMAP_MAX_BITS-n`.
mi_decl_nodiscard bool mi_bitmap_try_find_and_clearN(mi_bitmap_t* bitmap, size_t n, size_t tseq, size_t* pidx);
typedef bool (mi_claim_fun_t)(size_t slice_index, void* arg1, void* arg2, bool* keep_set);
typedef bool (mi_claim_fun_t)(size_t slice_index, mi_arena_t* arena, mi_subproc_t* subproc, mi_heaptag_t heap_tag, bool* keep_set);
mi_decl_nodiscard bool mi_bitmap_try_find_and_claim(mi_bitmap_t* bitmap, size_t tseq, size_t* pidx,
mi_claim_fun_t* claim, void* arg1, void* arg2);
mi_claim_fun_t* claim, mi_arena_t* arena, mi_subproc_t* subproc, mi_heaptag_t heap_tag );
void mi_bitmap_clear_once_set(mi_bitmap_t* bitmap, size_t idx);