diff --git a/CMakeLists.txt b/CMakeLists.txt index e16830aa..4f335805 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -13,6 +13,7 @@ option(MI_INTERPOSE "Use interpose to override standard malloc on macOS" option(MI_OSX_ZONE "Use malloc zone to override standard malloc on macOS" OFF) # enables interpose as well option(MI_LOCAL_DYNAMIC_TLS "Use slightly slower, dlopen-compatible TLS mechanism (Unix)" OFF) option(MI_BUILD_TESTS "Build test executables" ON) +option(MI_USER_CLEANUP "Enable user defined functionality for custom memory clean function" OFF) option(MI_CHECK_FULL "Use full internal invariant checking in DEBUG mode (deprecated, use MI_DEBUG_FULL instead)" OFF) include("cmake/mimalloc-config-version.cmake") @@ -84,6 +85,11 @@ if(MI_SECURE MATCHES "ON") list(APPEND mi_defines MI_SECURE=4) endif() +if(MI_USER_CLEANUP MATCHES "ON") + message(STATUS "Enable mi_register_user_cleanup functionality (MI_USER_CLEANUP=ON)") + list(APPEND mi_defines MI_USER_CLEANUP=1) +endif() + if(MI_SEE_ASM MATCHES "ON") message(STATUS "Generate assembly listings (MI_SEE_ASM=ON)") list(APPEND mi_cflags -save-temps) diff --git a/include/mimalloc-types.h b/include/mimalloc-types.h index dc85bbcd..13720ad6 100644 --- a/include/mimalloc-types.h +++ b/include/mimalloc-types.h @@ -36,6 +36,10 @@ terms of the MIT license. A copy of the license can be found in the file #define MI_SECURE 0 #endif +#if !defined(MI_USER_CLEANUP) +#define MI_USER_CLEANUP 0 +#endif + // Define MI_DEBUG for debug mode // #define MI_DEBUG 1 // basic assertion checks and statistics, check double free, corrupted free list, and invalid pointer free. // #define MI_DEBUG 2 // + internal assertion checks diff --git a/include/mimalloc.h b/include/mimalloc.h index 85f25ffb..2165333a 100644 --- a/include/mimalloc.h +++ b/include/mimalloc.h @@ -136,6 +136,9 @@ mi_decl_export void mi_register_output(mi_output_fun* out, void* arg) mi_attr_no typedef void (mi_cdecl mi_error_fun)(int err, void* arg); mi_decl_export void mi_register_error(mi_error_fun* fun, void* arg); +typedef void (mi_cleanup_fun)(void* user_data, void* p, size_t size); +mi_decl_export void mi_register_user_cleanup(mi_cleanup_fun* cleanup, void* user_data) mi_attr_noexcept; + mi_decl_export void mi_collect(bool force) mi_attr_noexcept; mi_decl_export int mi_version(void) mi_attr_noexcept; mi_decl_export void mi_stats_reset(void) mi_attr_noexcept; diff --git a/readme.md b/readme.md index 423c91b9..9e6cce15 100644 --- a/readme.md +++ b/readme.md @@ -39,6 +39,8 @@ It also has an easy way to override the allocator in [Windows](#override_on_wind randomized allocation, encrypted free lists, etc. to protect against various heap vulnerabilities. The performance penalty is usually around 10% on average over our benchmarks. +- __user function for clean up memory__: _mimalloc_ can be built with MI_USER_CLEANUP=ON flag. This mode + allows setup user function for memory clean up before it returned to system. - __first-class heaps__: efficiently create and use multiple heaps to allocate across different regions. A heap can be destroyed at once instead of deallocating each object separately. - __bounded__: it does not suffer from _blowup_ \[1\], has bounded worst-case allocation diff --git a/src/os.c b/src/os.c index 0aa85bd6..c03c555a 100644 --- a/src/os.c +++ b/src/os.c @@ -173,6 +173,39 @@ void _mi_os_init() { } #endif +#if MI_USER_CLEANUP +static mi_cleanup_fun* user_cleanup = NULL; +static void* cleanup_udata = NULL; +static void _mi_call_user_cleanup(void* p, size_t size) +{ + if(user_cleanup) + user_cleanup(cleanup_udata, p, size); +#if (MI_DEBUG>1) + else + memset(p, 0, size); +#endif +} +#else +static void _mi_call_user_cleanup(void* p, size_t size) +{ + (void)p; + (void)size; +#if (MI_DEBUG>1) + if (MI_SECURE==0) + memset(p, 0, size); // pretend it is eagerly reset +#endif +} +#endif + +void mi_register_user_cleanup(mi_cleanup_fun* cleanup, void* user_data) mi_attr_noexcept +{ + (void)cleanup; + (void)user_data; +#if MI_USER_CLEANUP + user_cleanup = cleanup; + cleanup_udata = user_data; +#endif +} /* ----------------------------------------------------------- Raw allocation on Windows (VirtualAlloc) and Unix's (mmap). @@ -182,6 +215,8 @@ static bool mi_os_mem_free(void* addr, size_t size, bool was_committed, mi_stats { if (addr == NULL || size == 0) return true; // || _mi_os_is_huge_reserved(addr) bool err = false; + if(was_committed) + _mi_call_user_cleanup(addr, size); #if defined(_WIN32) err = (VirtualFree(addr, 0, MEM_RELEASE) == 0); #elif defined(__wasi__) @@ -689,11 +724,7 @@ static bool mi_os_resetx(void* addr, size_t size, bool reset, mi_stats_t* stats) else _mi_stat_decrease(&stats->reset, csize); if (!reset) return true; // nothing to do on unreset! - #if (MI_DEBUG>1) - if (MI_SECURE==0) { - memset(start, 0, csize); // pretend it is eagerly reset - } - #endif + _mi_call_user_cleanup(start, csize); #if defined(_WIN32) // Testing shows that for us (on `malloc-large`) MEM_RESET is 2x faster than DiscardVirtualMemory