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) // Sub processes are used to keep memory separate between them (e.g. multiple interpreters in CPython)
typedef struct mi_subproc_s mi_subproc_t; 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`). // A page contains blocks of one specific size (`block_size`).
// Each page has three list of free blocks: // 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`) size_t block_size; // size available in each block (always `>0`)
uint8_t* page_start; // start of the blocks 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 bool free_is_zero; // `true` if the blocks in the free list are zero initialized
// padding // padding
#if (MI_ENCODE_FREELIST || MI_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") 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 // 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 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 // found an abandoned page of the right size
mi_arena_t* const arena = (mi_arena_t*)arg1; mi_page_t* const page = (mi_page_t*)mi_arena_slice_start(arena, slice_index);
mi_subproc_t* const subproc = (mi_subproc_t*)arg2;
mi_page_t* const page = (mi_page_t*)mi_arena_slice_start(arena, slice_index);
// can we claim ownership? // can we claim ownership?
if (!mi_page_try_claim_ownership(page)) { if (!mi_page_try_claim_ownership(page)) {
// there was a concurrent free .. // 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; *keep_abandoned = true;
return false; return false;
} }
if (subproc != page->subproc) { if (subproc != page->subproc || heap_tag != page->heap_tag) {
// wrong sub-process.. we need to unown again // 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) // (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. // 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. // 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; 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); MI_UNUSED(slice_count);
const size_t bin = _mi_bin(block_size); 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; size_t slice_index;
mi_bitmap_t* const bitmap = arena->pages_abandoned[bin]; 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 // found an abandoned page of the right size
// and claimed ownership. // and claimed ownership.
mi_page_t* page = (mi_page_t*)mi_arena_slice_start(arena, slice_index); 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; mi_tld_t* const tld = heap->tld;
// 1. look for an abandoned page // 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) { if (page != NULL) {
return page; // return as abandoned 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. // 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.) // (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_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) 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; const size_t slice_index = (chunk_idx * MI_BCHUNK_BITS) + cidx;
mi_assert_internal(slice_index < mi_bitmap_max_bits(bitmap)); mi_assert_internal(slice_index < mi_bitmap_max_bits(bitmap));
bool keep_set = true; bool keep_set = true;
if ((*claim)(slice_index, arg1, arg2, &keep_set)) { if ((*claim)(slice_index, arena, subproc, heap_tag, &keep_set)) {
// success! // success!
mi_assert_internal(!keep_set); mi_assert_internal(!keep_set);
*pidx = slice_index; *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`. // 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); 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_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); void mi_bitmap_clear_once_set(mi_bitmap_t* bitmap, size_t idx);