merge from 'dev-exp'; add support for MIMALLOC_OVERRIDE environment variable

src/alloc-override-win.c

@@ -15,6 +15,9 @@ terms of the MIT license. A copy of the license can be found in the file
 #include <windows.h>
 #include <psapi.h>
 
+#include <stdlib.h> // getenv
+#include <string.h> // strstr
+
 
 /*
 To override the C runtime `malloc` on Windows we need to patch the allocation
@@ -143,12 +146,6 @@ static size_t mi__msize_term(void* p) {
 }
 
 
-// Debug versions, forward to base versions in ucrtbase (that get patched)
-void* _malloc_base(size_t size);
-void* _calloc_base(size_t size, size_t count);
-void* _realloc_base(void* p, size_t new_size);
-void  _free_base(void* p);
-
 static void* mi__malloc_dbg(size_t size, int block_type, const char* fname, int line) {
   UNUSED(block_type); UNUSED(fname); UNUSED(line);
   return _malloc_base(size);
@@ -578,7 +575,7 @@ static void mi_module_resolve(const char* fname, HMODULE mod, int priority) {
         if (addr != NULL) {
           // found it! set the address
           patch->originals[i] = addr;          
-          _mi_trace_message("  override %s at %s!%p (entry %i)\n", patch->name, fname, addr, i);
+          _mi_trace_message("  found %s at %s!%p (entry %i)\n", patch->name, fname, addr, i);
         }
       }
     }
@@ -605,7 +602,6 @@ static bool mi_patches_resolve(void) {
   int ucrtbase_index = 0;
   int mimalloc_index = 0;
   // iterate through the loaded modules
-  _mi_trace_message("overriding malloc dynamically...\n");
   for (int i = 0; i < count;  i++) {
     HMODULE mod = modules[i];
     char filename[MAX_PATH] = { 0 };
@@ -679,27 +675,39 @@ __declspec(dllexport) BOOL WINAPI DllEntry(HINSTANCE inst, DWORD reason, LPVOID
     mi_patches_enable_term();
   }
   // C runtime main
-  BOOL ok = _DllMainCRTStartup(inst, reason, reserved);
+  BOOL ok = _DllMainCRTStartup(inst, reason, reserved);  
   if (reason == DLL_PROCESS_ATTACH && ok) {
+    // initialize at exit lists
+    mi_initialize_atexit();
+
     // Now resolve patches
     ok = mi_patches_resolve();
     if (ok) {
-      // and register our unwind entry (this must be after resolving due to possible delayed DLL initialization from GetProcAddress)
+      // check if patching is not disabled
-      mi_fls_unwind_entry = FlsAlloc(&mi_fls_unwind);
+      #pragma warning(suppress:4996)
-      if (mi_fls_unwind_entry != FLS_OUT_OF_INDEXES) {
+      const char* s = getenv("MIMALLOC_OVERRIDE");
-        FlsSetValue(mi_fls_unwind_entry, (void*)1);
+      bool enabled = (s == NULL || strstr("1;TRUE;YES;ON", s) != NULL);
+      if (!enabled) {
+        _mi_verbose_message("override is disabled\n");
       }
+      else {        
+        // and register our unwind entry (this must be after resolving due to possible delayed DLL initialization from GetProcAddress)
+        mi_fls_unwind_entry = FlsAlloc(&mi_fls_unwind);
+        if (mi_fls_unwind_entry != FLS_OUT_OF_INDEXES) {
+          FlsSetValue(mi_fls_unwind_entry, (void*)1);
+        }
 
-      // register our patch disabler in the global exit list
+        // register our patch disabler in the global exit list
-      mi_initialize_atexit();
+        if (crt_atexit != NULL)        (*crt_atexit)(&mi_patches_atexit);
-      if (crt_atexit != NULL)        (*crt_atexit)(&mi_patches_atexit);
+        if (crt_at_quick_exit != NULL) (*crt_at_quick_exit)(&mi_patches_at_quick_exit);
-      if (crt_at_quick_exit != NULL) (*crt_at_quick_exit)(&mi_patches_at_quick_exit);
 
-      // and patch !  this also redirects the `atexit` handling for the global exit list
+        // and patch !  this also redirects the `atexit` handling for the global exit list
-      mi_patches_enable();
+        mi_patches_enable();
+        _mi_verbose_message("override is enabled\n");
 
-      // hide internal allocation
+        // hide internal allocation
-      mi_stats_reset();
+        mi_stats_reset();
+      }
     }
   }
   return ok;

src/init.c

@@ -385,7 +385,7 @@ bool _mi_preloading() {
 }
 
 // Communicate with the redirection module on Windows
-#if defined(_WIN32) && defined(XMI_SHARED_LIB) 
+#if 0
 #ifdef __cplusplus
 extern "C" {
 #endif

src/options.c

@@ -60,7 +60,7 @@ static mi_option_desc_t options[_mi_option_last] =
   { 0, UNINIT, "large_os_pages" },      // use large OS pages, use only with eager commit to prevent fragmentation of VMA's
   { 0, UNINIT, "page_reset" },
   { 0, UNINIT, "cache_reset" },
-  { 0, UNINIT, "reset_decommits" },     // note: cannot enable this if secure is on 
+  { 0, UNINIT, "reset_decommits" },     // note: cannot enable this if secure is on
   { 0, UNINIT, "reset_discards" }       // note: cannot enable this if secure is on
 };
 
@@ -213,12 +213,13 @@ static const char* mi_getenv(const char* name) {
   #pragma warning(suppress:4996)
   const char* s = getenv(name);
   if (s == NULL) {
-    char buf[64];
+    char buf[64+1];
-    mi_strlcpy(buf,name,sizeof(buf));
+    size_t len = strlen(name);
-    size_t buf_size = strlen(buf);
+    if (len >= sizeof(buf)) len = sizeof(buf) - 1;
-    for (size_t i = 0; i < buf_size; i++) {
+    for (size_t i = 0; i < len; i++) {
-      buf[i] = toupper(buf[i]);
+      buf[i] = toupper(name[i]);
     }
+    buf[len] = 0;
     #pragma warning(suppress:4996)
     s = getenv(buf);
   }
@@ -228,16 +229,17 @@ static const char* mi_getenv(const char* name) {
 static void mi_option_init(mi_option_desc_t* desc) {
   if (!_mi_preloading()) desc->init = DEFAULTED;
   // Read option value from the environment
-  char buf[64];
+  char buf[64+1];
   mi_strlcpy(buf, "mimalloc_", sizeof(buf));
   mi_strlcat(buf, desc->name, sizeof(buf));
-  const char* s = mi_getenv(buf);  
+  const char* s = mi_getenv(buf);
   if (s != NULL) {
-    mi_strlcpy(buf, s, sizeof(buf));
+    size_t len = strlen(s);
-    size_t buf_size = strlen(buf); // TODO: use strnlen?
+    if (len >= sizeof(buf)) len = sizeof(buf) - 1;
-    for (size_t i = 0; i < buf_size; i++) {
+    for (size_t i = 0; i < len; i++) {
-      buf[i] = toupper(buf[i]);
+      buf[i] = toupper(s[i]);
     }
+    buf[len] = 0;
     if (buf[0]==0 || strstr("1;TRUE;YES;ON", buf) != NULL) {
       desc->value = 1;
       desc->init = INITIALIZED;

src/os.c

@@ -204,7 +204,7 @@ static void* mi_win_virtual_alloc(void* addr, size_t size, size_t try_alignment,
   if (use_large_os_page(size, try_alignment)) {
     uintptr_t try_ok = mi_atomic_read(&large_page_try_ok);
     if (try_ok > 0) {
-      // if a large page page allocation fails, it seems the calls to VirtualAlloc get very expensive.
+      // if a large page allocation fails, it seems the calls to VirtualAlloc get very expensive.
       // therefore, once a large page allocation failed, we don't try again for `large_page_try_ok` times.
       mi_atomic_compare_exchange(&large_page_try_ok, try_ok - 1, try_ok);
     }

test/main-override.cpp

@@ -30,12 +30,12 @@ int main() {
   free(p1);
   p1 = malloc(8);
   char* s = mi_strdup("hello\n");
-  free(p2);
+  mi_free(p2);
   p2 = malloc(16);
   p1 = realloc(p1, 32);
   free(p1);
   free(p2);
-  free(s);
+  mi_free(s);
   Test* t = new Test(42);
   delete t;
   t = new (std::nothrow) Test(42);