update tests on windows

Author: daanx

src/prim/windows/prim.c

@@ -746,7 +746,96 @@ static void NTAPI mi_win_main(PVOID module, DWORD reason, LPVOID reserved) {
 }
 
 
-#if defined(MI_SHARED_LIB)
+#if 1  /* we use a combination of _pRawDllMain and TLS init/fini for both static and dynamic linkage */
+  #define MI_PRIM_HAS_PROCESS_ATTACH  1
+  // 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_theap(mi_theap_t* theap) {
+    MI_UNUSED(theap);
+  }
+
+  // If linked into a DLL module, this raw entry is called before the CRT attach and 
+  // after the CRT detach through the CRT _pRawDllMain pointer.
+  static BOOL NTAPI mi_dll_main_raw(PVOID module, DWORD reason, LPVOID reserved) {
+    //if (reason == DLL_PROCESS_ATTACH)      { mi_debug_out("dll process attach\n"); }
+    //else if (reason == DLL_PROCESS_DETACH) { mi_debug_out("dll process detach\n"); }
+    //else if (reason == DLL_THREAD_ATTACH)  { mi_debug_out("dll thread attach\n"); }
+    //else if (reason == DLL_THREAD_DETACH)  { mi_debug_out("dll thread detach\n"); }
+    mi_win_main(module, reason, reserved);
+    return TRUE;
+  }
+
+  // Set the value of the CRT _pRawDllMain pointer
+  #if defined(__cplusplus)
+  extern "C"
+  #endif
+  PVOID _pRawDllMain = &mi_dll_main_raw;
+
+  // We also hook into the Windows loader TLS initialization and finalization.
+  // If we are linked into an EXE module we rely on these as `mi_dll_main_raw`
+  // is not called (and otherwise we ignore the TLS callbacks by checking if we are in a DLL).
+  static bool mi_module_is_dll(PVOID mod) {
+    if (mod==NULL) return false;
+    PIMAGE_DOS_HEADER imageDosHeader = (PIMAGE_DOS_HEADER)mod;
+    PIMAGE_NT_HEADERS imageNtHeaders = (PIMAGE_NT_HEADERS)((unsigned char*)imageDosHeader + imageDosHeader->e_lfanew);
+    return ((imageNtHeaders->FileHeader.Characteristics & IMAGE_FILE_DLL) == IMAGE_FILE_DLL);
+  }
+
+  static void NTAPI mi_tls_attach(PVOID module, DWORD reason, LPVOID reserved) {
+    if (reason == DLL_PROCESS_ATTACH || reason == DLL_THREAD_ATTACH) {
+      if (!mi_module_is_dll(module)) {
+        // mi_debug_out(reason==DLL_PROCESS_ATTACH ? "exe process attach\n" : "exe thread attach\n");
+        mi_win_main(module, reason, reserved);
+      }
+    }
+  }
+  static void NTAPI mi_tls_detach(PVOID module, DWORD reason, LPVOID reserved) {
+    if (reason == DLL_PROCESS_DETACH || reason == DLL_THREAD_DETACH) {
+      if (!mi_module_is_dll(module)) {
+        // mi_debug_out(reason==DLL_PROCESS_DETACH ? "exe process detach\n" : "exe thread detach\n");
+        mi_win_main(module, reason, reserved);
+      }
+    }
+  }
+
+  // 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 use 2 entries to ensure we call attach events before constructors
+  // are called, and detach events after destructors are called.
+  #if defined(__cplusplus)
+  extern "C" {
+  #endif
+
+  #if defined(_WIN64)
+    #pragma comment(linker, "/INCLUDE:_tls_used")
+    #pragma comment(linker, "/INCLUDE:_mi_tls_callback_pre")
+    #pragma comment(linker, "/INCLUDE:_mi_tls_callback_post")
+    #pragma const_seg(".CRT$XLB")
+      extern const PIMAGE_TLS_CALLBACK _mi_tls_callback_pre[];
+      const PIMAGE_TLS_CALLBACK _mi_tls_callback_pre[] = { &mi_tls_attach };
+    #pragma const_seg()
+    #pragma const_seg(".CRT$XLY")
+      extern const PIMAGE_TLS_CALLBACK _mi_tls_callback_post[];
+      const PIMAGE_TLS_CALLBACK _mi_tls_callback_post[] = { &mi_tls_detach };
+    #pragma const_seg()
+  #else
+    #pragma comment(linker, "/INCLUDE:__tls_used")
+    #pragma comment(linker, "/INCLUDE:__mi_tls_callback_pre")
+    #pragma comment(linker, "/INCLUDE:__mi_tls_callback_post")
+    #pragma data_seg(".CRT$XLB")
+      PIMAGE_TLS_CALLBACK _mi_tls_callback_pre[] = { &mi_tls_attach };
+    #pragma data_seg()
+    #pragma data_seg(".CRT$XIY")
+      PIMAGE_TLS_CALLBACK _mi_tls_callback_post[] = { &mi_tls_detach };
+    #pragma data_seg()
+  #endif
+
+  #if defined(__cplusplus)
+  }
+  #endif
+
+#elif defined(MI_SHARED_LIB)
   #define MI_PRIM_HAS_PROCESS_ATTACH  1
 
   // Windows DLL: easy to hook into process_init and thread_done
@@ -775,7 +864,7 @@ static void NTAPI mi_win_main(PVOID module, DWORD reason, LPVOID reserved) {
       mi_win_main(module, reason, reserved);
     }
   }
-
+  
   // 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 use 2 entries to ensure we call attach events before constructors
@@ -795,15 +884,15 @@ static void NTAPI mi_win_main(PVOID module, DWORD reason, LPVOID reserved) {
     #pragma const_seg(".CRT$XLY")
     extern const PIMAGE_TLS_CALLBACK _mi_tls_callback_post[];
     const PIMAGE_TLS_CALLBACK _mi_tls_callback_post[] = { &mi_win_main_detach };
-    #pragma const_seg()
+    #pragma const_seg()    
   #else
     #pragma comment(linker, "/INCLUDE:__tls_used")
     #pragma comment(linker, "/INCLUDE:__mi_tls_callback_pre")
     #pragma comment(linker, "/INCLUDE:__mi_tls_callback_post")
     #pragma data_seg(".CRT$XLB")
     PIMAGE_TLS_CALLBACK _mi_tls_callback_pre[] = { &mi_win_main_attach };
     #pragma data_seg()
-    #pragma data_seg(".CRT$XLY")
+    #pragma data_seg(".CRT$XIY")
     PIMAGE_TLS_CALLBACK _mi_tls_callback_post[] = { &mi_win_main_detach };
     #pragma data_seg()
   #endif
@@ -903,6 +992,7 @@ static void NTAPI mi_win_main(PVOID module, DWORD reason, LPVOID reserved) {
       mi_redirected = false;
     }
     else if (reason == DLL_THREAD_DETACH) {
+      // mi_debug_out("redirect thread detach\n");
       _mi_thread_done(NULL);
     }
   }

test/main-override-static.c

@@ -287,7 +287,7 @@ static void test_manage_os_memory(void) {
 
 #if _WIN32
 
-static void test_dep(void) {
+static void call_library(void) {
   HMODULE dll = LoadLibraryA("mimalloc-test-static-dep.dll");
   if (dll != NULL) {
     TestFun fun = (TestFun)GetProcAddress(dll, "Test");
@@ -297,6 +297,11 @@ static void test_dep(void) {
     bool ok = FreeLibrary(dll);
   }
 }
+
+static void test_dep(void) {
+  call_library();
+  call_library();
+}
 #endif
 
 // Experiment with huge OS pages

test/main-static-dep.cpp

@@ -39,6 +39,7 @@ BOOL WINAPI DllMain(HINSTANCE module, DWORD reason, LPVOID reserved) {
     printf("static-dep: dll attach\n");
   }
   else if (reason==DLL_PROCESS_DETACH) {
+    mi_option_enable(mi_option_destroy_on_exit);
     printf("static-dep: dll detach\n");
   }
   return TRUE;