mirrors/mimalloc
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

reorganize primitives for process initialization; use special data segment on Windows for thread termination by default on Windows now (issue #869)

Browse Source
This commit is contained in:
Daan 2024-10-21 22:56:59 -07:00
parent 34e66778ec
commit aa881733d7
4 changed files with 209 additions and 151 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
include/mimalloc/internal.h
View File

@ -77,6 +77,11 @@ static inline uintptr_t _mi_random_shuffle(uintptr_t x);
// init.c
extern mi_decl_cache_align mi_stats_t _mi_stats_main;
extern mi_decl_cache_align const mi_page_t _mi_page_empty;
void _mi_process_load(void);
void mi_cdecl _mi_process_done(void);
bool _mi_is_redirected(void);
bool _mi_allocator_init(const char** message);
void _mi_allocator_done(void);
bool _mi_is_main_thread(void);
size_t _mi_current_thread_count(void);
bool _mi_preloading(void); // true while the C runtime is not initialized yet

114
src/init.c
View File

@ -459,10 +459,6 @@ void mi_thread_init(void) mi_attr_noexcept
//_mi_verbose_message("thread init: 0x%zx\n", _mi_thread_id());
}
void mi_thread_done(void) mi_attr_noexcept {
_mi_thread_done(NULL);
}
void _mi_thread_done(mi_heap_t* heap)
{
// calling with NULL implies using the default heap
@ -508,54 +504,15 @@ void _mi_heap_set_default_direct(mi_heap_t* heap) {
// --------------------------------------------------------
// Run functions on process init/done, and thread init/done
// --------------------------------------------------------
static void mi_cdecl mi_process_done(void);
static bool os_preloading = true; // true until this module is initialized
static bool mi_redirected = false; // true if malloc redirects to mi_malloc
// Returns true if this module has not been initialized; Don't use C runtime routines until it returns false.
bool mi_decl_noinline _mi_preloading(void) {
return os_preloading;
}
mi_decl_nodiscard bool mi_is_redirected(void) mi_attr_noexcept {
return mi_redirected;
}
// Communicate with the redirection module on Windows
#if defined(_WIN32) && defined(MI_SHARED_LIB) && !defined(MI_WIN_NOREDIRECT)
#ifdef __cplusplus
extern "C" {
#endif
mi_decl_export void _mi_redirect_entry(DWORD reason) {
// called on redirection; careful as this may be called before DllMain
if (reason == DLL_PROCESS_ATTACH) {
mi_redirected = true;
}
else if (reason == DLL_PROCESS_DETACH) {
mi_redirected = false;
}
else if (reason == DLL_THREAD_DETACH) {
mi_thread_done();
}
}
__declspec(dllimport) bool mi_cdecl mi_allocator_init(const char** message);
__declspec(dllimport) void mi_cdecl mi_allocator_done(void);
#ifdef __cplusplus
}
#endif
#else
static bool mi_allocator_init(const char** message) {
if (message != NULL) *message = NULL;
return true;
}
static void mi_allocator_done(void) {
// nothing to do
}
#endif
// Called once by the process loader
static void mi_process_load(void) {
// Called once by the process loader from `src/prim/prim.c`
void _mi_process_load(void) {
mi_heap_main_init();
#if defined(__APPLE__) || defined(MI_TLS_RECURSE_GUARD)
volatile mi_heap_t* dummy = _mi_heap_default; // access TLS to allocate it before setting tls_initialized to true;
@ -563,17 +520,14 @@ static void mi_process_load(void) {
#endif
os_preloading = false;
mi_assert_internal(_mi_is_main_thread());
#if !(defined(_WIN32) && defined(MI_SHARED_LIB)) // use Dll process detach (see below) instead of atexit (issue #521)
atexit(&mi_process_done);
#endif
_mi_options_init();
mi_process_setup_auto_thread_done();
mi_process_init();
if (mi_redirected) _mi_verbose_message("malloc is redirected.\n");
if (_mi_is_redirected()) _mi_verbose_message("malloc is redirected.\n");
// show message from the redirector (if present)
const char* msg = NULL;
mi_allocator_init(&msg);
_mi_allocator_init(&msg);
if (msg != NULL && (mi_option_is_enabled(mi_option_verbose) || mi_option_is_enabled(mi_option_show_errors))) {
_mi_fputs(NULL,NULL,NULL,msg);
}
@ -651,7 +605,7 @@ void mi_process_init(void) mi_attr_noexcept {
}
// Called when the process is done (through `at_exit`)
static void mi_cdecl mi_process_done(void) {
void mi_cdecl _mi_process_done(void) {
// only shutdown if we were initialized
if (!_mi_process_is_initialized) return;
// ensure we are called once
@ -683,64 +637,8 @@ static void mi_cdecl mi_process_done(void) {
if (mi_option_is_enabled(mi_option_show_stats) || mi_option_is_enabled(mi_option_verbose)) {
mi_stats_print(NULL);
}
mi_allocator_done();
_mi_allocator_done();
_mi_verbose_message("process done: 0x%zx\n", _mi_heap_main.thread_id);
os_preloading = true; // don't call the C runtime anymore
}
#if defined(_WIN32) && defined(MI_SHARED_LIB)
// Windows DLL: easy to hook into process_init and thread_done
__declspec(dllexport) BOOL WINAPI DllMain(HINSTANCE inst, DWORD reason, LPVOID reserved) {
MI_UNUSED(reserved);
MI_UNUSED(inst);
if (reason==DLL_PROCESS_ATTACH) {
mi_process_load();
}
else if (reason==DLL_PROCESS_DETACH) {
mi_process_done();
}
else if (reason==DLL_THREAD_DETACH) {
if (!mi_is_redirected()) {
mi_thread_done();
}
}
return TRUE;
}
#elif defined(_MSC_VER)
// MSVC: use data section magic for static libraries
// See <https://www.codeguru.com/cpp/misc/misc/applicationcontrol/article.php/c6945/Running-Code-Before-and-After-Main.htm>
static int _mi_process_init(void) {
mi_process_load();
return 0;
}
typedef int(*_mi_crt_callback_t)(void);
#if defined(_M_X64) || defined(_M_ARM64)
__pragma(comment(linker, "/include:" "_mi_msvc_initu"))
#pragma section(".CRT$XIU", long, read)
#else
__pragma(comment(linker, "/include:" "__mi_msvc_initu"))
#endif
#pragma data_seg(".CRT$XIU")
mi_decl_externc _mi_crt_callback_t _mi_msvc_initu[] = { &_mi_process_init };
#pragma data_seg()
#elif defined(__cplusplus)
// C++: use static initialization to detect process start
static bool _mi_process_init(void) {
mi_process_load();
return (_mi_heap_main.thread_id != 0);
}
static bool mi_initialized = _mi_process_init();
#elif defined(__GNUC__) || defined(__clang__)
// GCC,Clang: use the constructor attribute
static void __attribute__((constructor)) _mi_process_init(void) {
mi_process_load();
}
#else
#pragma message("define a way to call mi_process_load on your platform")
#endif

44
src/prim/prim.c
View File

@ -25,3 +25,47 @@ terms of the MIT license. A copy of the license can be found in the file
#include "unix/prim.c" // mmap() (Linux, macOSX, BSD, Illumnos, Haiku, DragonFly, etc.)
#endif
// Generic process initialization
#ifndef MI_PRIM_HAS_PROCESS_ATTACH
#if defined(__GNUC__) || defined(__clang__)
// GCC,Clang: use the constructor attribute
#if defined(__clang__)
#define mi_attr_constructor __attribute__((constructor(101)))
#define mi_attr_destructor __attribute__((destructor(101)))
#else
#define mi_attr_constructor __attribute__((constructor))
#define mi_attr_destructor __attribute__((destructor))
#endif
static void mi_attr_constructor mi_process_attach(void) {
_mi_process_load();
}
static void mi_attr_destructor mi_process_detach(void) {
_mi_process_done();
}
#elif defined(__cplusplus)
// C++: use static initialization to detect process start
static bool mi_process_attach(void) {
_mi_process_load();
atexit(&_mi_process_done);
return (_mi_heap_main.thread_id != 0);
}
static bool mi_initialized = mi_process_attach();
#else
#pragma message("define a way to call _mi_process_load/done on your platform")
#endif
#endif
// Generic allocator init/done callback
#ifndef MI_PRIM_HAS_ALLOCATOR_INIT
bool _mi_is_redirected(void) {
return false;
}
bool _mi_allocator_init(const char** message) {
if (message != NULL) *message = NULL;
return true;
}
void _mi_allocator_done(void) {
// nothing to do
}
#endif

197
src/prim/windows/prim.c
View File

@ -499,8 +499,7 @@ void _mi_prim_process_info(mi_process_info_t* pinfo)
}
// get process info
PROCESS_MEMORY_COUNTERS info;
memset(&info, 0, sizeof(info));
PROCESS_MEMORY_COUNTERS info; _mi_memzero_var(info);
if (pGetProcessMemoryInfo != NULL) {
pGetProcessMemoryInfo(GetCurrentProcess(), &info, sizeof(info));
}
@ -602,60 +601,172 @@ bool _mi_prim_random_buf(void* buf, size_t buf_len) {
#endif // MI_USE_RTLGENRANDOM
//----------------------------------------------------------------
// Thread init/done
// Process & Thread Init/Done
//----------------------------------------------------------------
#if !defined(MI_SHARED_LIB)
// use thread local storage keys to detect thread ending
// note: another design could be to use special linker sections (see issue #869)
#include <fibersapi.h>
#if (_WIN32_WINNT < 0x600) // before Windows Vista
WINBASEAPI DWORD WINAPI FlsAlloc( _In_opt_ PFLS_CALLBACK_FUNCTION lpCallback );
WINBASEAPI PVOID WINAPI FlsGetValue( _In_ DWORD dwFlsIndex );
WINBASEAPI BOOL WINAPI FlsSetValue( _In_ DWORD dwFlsIndex, _In_opt_ PVOID lpFlsData );
WINBASEAPI BOOL WINAPI FlsFree(_In_ DWORD dwFlsIndex);
#endif
static DWORD mi_fls_key = (DWORD)(-1);
static void NTAPI mi_fls_done(PVOID value) {
mi_heap_t* heap = (mi_heap_t*)value;
if (heap != NULL) {
_mi_thread_done(heap);
FlsSetValue(mi_fls_key, NULL); // prevent recursion as _mi_thread_done may set it back to the main heap, issue #672
static void NTAPI mi_win_main(PVOID module, DWORD reason, LPVOID reserved) {
MI_UNUSED(reserved);
MI_UNUSED(module);
if (reason==DLL_PROCESS_ATTACH) {
_mi_process_load();
}
else if (reason==DLL_PROCESS_DETACH) {
_mi_process_done();
}
else if (reason==DLL_THREAD_DETACH && !_mi_is_redirected()) {
_mi_thread_done(NULL);
}
}
void _mi_prim_thread_init_auto_done(void) {
mi_fls_key = FlsAlloc(&mi_fls_done);
}
void _mi_prim_thread_done_auto_done(void) {
// call thread-done on all threads (except the main thread) to prevent
// dangling callback pointer if statically linked with a DLL; Issue #208
FlsFree(mi_fls_key);
}
#if defined(MI_SHARED_LIB)
#define MI_PRIM_HAS_PROCESS_INIT 1
void _mi_prim_thread_associate_default_heap(mi_heap_t* heap) {
mi_assert_internal(mi_fls_key != (DWORD)(-1));
FlsSetValue(mi_fls_key, heap);
}
// Windows DLL: easy to hook into process_init and thread_done
__declspec(dllexport) BOOL WINAPI DllMain(HINSTANCE inst, DWORD reason, LPVOID reserved) {
win_main((PVOID)inst,reason,reserved);
return TRUE;
}
#else
// nothing to do since `_mi_thread_done` is handled through the DLL_THREAD_DETACH event.
void _mi_prim_thread_init_auto_done(void) { }
void _mi_prim_thread_done_auto_done(void) { }
void _mi_prim_thread_associate_default_heap(mi_heap_t* heap) {
MI_UNUSED(heap);
}
// Dll; nothing to do as in that case thread_done is handled through the DLL_THREAD_DETACH event.
#elif !defined(MI_WIN_USE_FLS)
#define MI_PRIM_HAS_PROCESS_INIT 1
void _mi_prim_thread_init_auto_done(void) {
}
// Set up TLS callbacks in a statically linked library by using special data sections.
// See <https://stackoverflow.com/questions/14538159/tls-callback-in-windows>
// We may use ".CRT$XLY" instead of "B" -- see also issue #869.
#if defined(__cplusplus)
extern "C" {
#endif
void _mi_prim_thread_done_auto_done(void) {
}
#if defined(_WIN64)
#pragma comment(linker, "/INCLUDE:_tls_used")
#pragma comment(linker, "/INCLUDE:_mi_tls_callback")
#pragma const_seg(".CRT$XLB")
extern const PIMAGE_TLS_CALLBACK _mi_tls_callback[];
const PIMAGE_TLS_CALLBACK _mi_tls_callback[] = { &mi_win_main };
#pragma const_seg()
#else
#pragma comment(linker, "/INCLUDE:__tls_used")
#pragma comment(linker, "/INCLUDE:__mi_tls_callback")
#pragma data_seg(".CRT$XLB")
const PIMAGE_TLS_CALLBACK _mi_tls_callback[] = { &mi_win_main };
#pragma data_seg()
#endif
void _mi_prim_thread_associate_default_heap(mi_heap_t* heap) {
MI_UNUSED(heap);
}
#if defined(__cplusplus)
}
#endif
// nothing to do since `_mi_thread_done` is handled through the DLL_THREAD_DETACH event.
void _mi_prim_thread_init_auto_done(void) { }
void _mi_prim_thread_done_auto_done(void) { }
void _mi_prim_thread_associate_default_heap(mi_heap_t* heap) {
MI_UNUSED(heap);
}
#else // statically linked, use fiber api
#if defined(_MSC_VER) // on clang/gcc use the constructor attribute (in `src/prim/prim.c`)
// MSVC: use data section magic for static libraries
// See <https://www.codeguru.com/cpp/misc/misc/applicationcontrol/article.php/c6945/Running-Code-Before-and-After-Main.htm>
#define MI_PRIM_HAS_PROCESS_INIT 1
static int mi_process_attach(void) {
mi_win_main(NULL,DLL_PROCESS_ATTACH,NULL);
atexit(&_mi_process_done);
return 0;
}
typedef int(*_mi_crt_callback_t)(void);
#if defined(_M_X64) || defined(_M_ARM64)
__pragma(comment(linker, "/include:" "_mi_msvc_initu"))
#pragma section(".CRT$XIU", long, read)
#else
__pragma(comment(linker, "/include:" "__mi_msvc_initu"))
#endif
#pragma data_seg(".CRT$XIU")
mi_decl_externc _mi_crt_callback_t _mi_msvc_initu[] = { &mi_process_attach };
#pragma data_seg()
#endif
// use the fiber api for calling `_mi_thread_done`.
#include <fibersapi.h>
#if (_WIN32_WINNT < 0x600) // before Windows Vista
WINBASEAPI DWORD WINAPI FlsAlloc( _In_opt_ PFLS_CALLBACK_FUNCTION lpCallback );
WINBASEAPI PVOID WINAPI FlsGetValue( _In_ DWORD dwFlsIndex );
WINBASEAPI BOOL WINAPI FlsSetValue( _In_ DWORD dwFlsIndex, _In_opt_ PVOID lpFlsData );
WINBASEAPI BOOL WINAPI FlsFree(_In_ DWORD dwFlsIndex);
#endif
static DWORD mi_fls_key = (DWORD)(-1);
static void NTAPI mi_fls_done(PVOID value) {
mi_heap_t* heap = (mi_heap_t*)value;
if (heap != NULL) {
_mi_thread_done(heap);
FlsSetValue(mi_fls_key, NULL); // prevent recursion as _mi_thread_done may set it back to the main heap, issue #672
}
}
void _mi_prim_thread_init_auto_done(void) {
mi_fls_key = FlsAlloc(&mi_fls_done);
}
void _mi_prim_thread_done_auto_done(void) {
// call thread-done on all threads (except the main thread) to prevent
// dangling callback pointer if statically linked with a DLL; Issue #208
FlsFree(mi_fls_key);
}
void _mi_prim_thread_associate_default_heap(mi_heap_t* heap) {
mi_assert_internal(mi_fls_key != (DWORD)(-1));
FlsSetValue(mi_fls_key, heap);
}
#endif
// ----------------------------------------------------
// Communicate with the redirection module on Windows
// ----------------------------------------------------
#if defined(MI_SHARED_LIB) && !defined(MI_WIN_NOREDIRECT)
static bool mi_redirected = false; // true if malloc redirects to mi_malloc
bool _mi_is_redirected(void) {
return mi_redirected;
}
#ifdef __cplusplus
extern "C" {
#endif
mi_decl_export void _mi_redirect_entry(DWORD reason) {
// called on redirection; careful as this may be called before DllMain
if (reason == DLL_PROCESS_ATTACH) {
mi_redirected = true;
}
else if (reason == DLL_PROCESS_DETACH) {
mi_redirected = false;
}
else if (reason == DLL_THREAD_DETACH) {
_mi_thread_done(NULL);
}
}
__declspec(dllimport) bool mi_cdecl mi_allocator_init(const char** message);
__declspec(dllimport) void mi_cdecl mi_allocator_done(void);
#ifdef __cplusplus
}
#endif
bool _mi_allocator_init(const char** message) {
return mi_allocator_init(message);
}
void _mi_allocator_done(void) {
mi_allocator_done();
}
#endif