add thread-id check for thread local FLS callbacks on Windows with static linking; found by @jasongibson

src/init.c

@@ -216,7 +216,8 @@ static bool _mi_heap_done(mi_heap_t* heap) {
   heap = heap->tld->heap_backing;
   if (!mi_heap_is_initialized(heap)) return false;
 
-
+  // check thread-id as on Windows shutdown with FLS the main (exit) thread may call this on thread-local heaps
+  if (heap->thread_id == _mi_thread_id()) {  
     // delete all non-backing heaps in this thread
     mi_heap_t* curr = heap->tld->heaps;
     while (curr != NULL) {
@@ -234,13 +235,14 @@ static bool _mi_heap_done(mi_heap_t* heap) {
     if (heap != &_mi_heap_main) {
       _mi_heap_collect_abandon(heap);
     }
+  }
 
   // merge stats
   _mi_stats_done(&heap->tld->stats);
 
   // free if not the main thread
   if (heap != &_mi_heap_main) {
-    mi_assert_internal(heap->tld->segments.count == 0);
+    mi_assert_internal(heap->tld->segments.count == 0 || heap->thread_id != _mi_thread_id());
     _mi_os_free(heap, sizeof(mi_thread_data_t), &_mi_stats_main);
   }
 #if 0  
@@ -294,6 +296,10 @@ static void _mi_thread_done(mi_heap_t* default_heap);
   #endif
   static DWORD mi_fls_key = (DWORD)(-1);
   static void NTAPI mi_fls_done(PVOID value) {
+    if (mi_fls_key != -1 && _mi_is_main_thread()) {
+      FlsSetValue(mi_fls_key, NULL); // null out once to prevent recursion on the main thread
+      mi_fls_key = -1;
+    }
     if (value!=NULL) _mi_thread_done((mi_heap_t*)value);
   }
 #elif defined(MI_USE_PTHREADS)

test/main-override.cpp

@@ -7,11 +7,15 @@
 #include <mimalloc.h>
 #include <new>
 #include <vector>
+#include <future>
+#include <iostream>
 
 #include <thread>
 #include <mimalloc.h>
 #include <assert.h>
 
+#include <mimalloc-new-delete.h>
+
 #ifdef _WIN32
 #include <windows.h>
 static void msleep(unsigned long msecs) { Sleep(msecs); }
@@ -25,6 +29,7 @@ void heap_no_delete();         // issue #202
 void heap_late_free();         // issue #204
 void padding_shrink();         // issue #209
 void various_tests();
+void test_mt_shutdown();
 
 int main() {
   mi_stats_reset();  // ignore earlier allocations
@@ -33,6 +38,7 @@ int main() {
   heap_late_free();
   padding_shrink();
   various_tests();
+  //test_mt_shutdown();
   mi_stats_print(NULL);
   return 0;
 }
@@ -173,3 +179,29 @@ void heap_thread_free_large() {
     t1.join();
   }
 }
+
+
+
+void test_mt_shutdown()
+{
+  const int threads = 5;
+  std::vector< std::future< std::vector< char* > > > ts;
+
+  auto fn = [&]()
+  {
+    std::vector< char* > ps;
+    ps.reserve(1000);
+    for (int i = 0; i < 1000; i++)
+      ps.emplace_back(new char[1]);
+    return ps;
+  };
+
+  for (int i = 0; i < threads; i++)
+    ts.emplace_back(std::async(std::launch::async, fn));
+
+  for (auto& f : ts)
+    for (auto& p : f.get())
+      delete[] p;
+
+  std::cout << "done" << std::endl;
+}