Merge pull request #2 from hardened-steel/feature/add-cleanup-mem-function

Feature/add cleanup mem function

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)

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

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;

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

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