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

src/init.c

@@ -216,23 +216,25 @@ static bool _mi_heap_done(mi_heap_t* heap) {
   heap = heap->tld->heap_backing;
   if (!mi_heap_is_initialized(heap)) return false;
 
-
-  // delete all non-backing heaps in this thread
-  mi_heap_t* curr = heap->tld->heaps;
-  while (curr != NULL) {
-    mi_heap_t* next = curr->next; // save `next` as `curr` will be freed
-    if (curr != heap) {
-      mi_assert_internal(!mi_heap_is_backing(curr));
-      mi_heap_delete(curr);
+  // 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) {
+      mi_heap_t* next = curr->next; // save `next` as `curr` will be freed
+      if (curr != heap) {
+        mi_assert_internal(!mi_heap_is_backing(curr));
+        mi_heap_delete(curr);
+      }
+      curr = next;
     }
-    curr = next;
-  }
-  mi_assert_internal(heap->tld->heaps == heap && heap->next == NULL);
-  mi_assert_internal(mi_heap_is_backing(heap));
+    mi_assert_internal(heap->tld->heaps == heap && heap->next == NULL);
+    mi_assert_internal(mi_heap_is_backing(heap));
 
-  // collect if not the main thread
-  if (heap != &_mi_heap_main) {
-    _mi_heap_collect_abandon(heap);
+    // collect if not the main thread
+    if (heap != &_mi_heap_main) {
+      _mi_heap_collect_abandon(heap);
+    }
   }
 
   // merge stats
@@ -240,7 +242,7 @@ static bool _mi_heap_done(mi_heap_t* heap) {
 
   // 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,14 +29,16 @@ 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
   heap_thread_free_large();
-  heap_no_delete();  
-  heap_late_free();  
-  padding_shrink();  
+  heap_no_delete();
+  heap_late_free();
+  padding_shrink();
   various_tests();
+  //test_mt_shutdown();
   mi_stats_print(NULL);
   return 0;
 }
@@ -53,20 +59,20 @@ public:
 };
 
 
-void various_tests() {  
+void various_tests() {
   atexit(free_p);
   void* p1 = malloc(78);
-  void* p2 = mi_malloc_aligned(16,24);
-  free(p1);  
+  void* p2 = mi_malloc_aligned(16, 24);
+  free(p1);
   p1 = malloc(8);
   char* s = mi_strdup("hello\n");
-  
+
   //char* s = _strdup("hello\n");
   //char* buf = NULL;
   //size_t len;
   //_dupenv_s(&buf,&len,"MIMALLOC_VERBOSE"); 
   //mi_free(buf);
-  
+
   mi_free(p2);
   p2 = malloc(16);
   p1 = realloc(p1, 32);
@@ -76,7 +82,7 @@ void various_tests() {
   Test* t = new Test(42);
   delete t;
   t = new (std::nothrow) Test(42);
-  delete t;  
+  delete t;
 }
 
 class Static {
@@ -105,7 +111,7 @@ bool test_stl_allocator1() {
 
 struct some_struct { int i; int j; double z; };
 
-bool test_stl_allocator2() {  
+bool test_stl_allocator2() {
   std::vector<some_struct, mi_stl_allocator<some_struct> > vec;
   vec.push_back(some_struct());
   vec.pop_back();
@@ -117,13 +123,13 @@ bool test_stl_allocator2() {
 // Issue #202
 void heap_no_delete_worker() {
   mi_heap_t* heap = mi_heap_new();
-  void* q = mi_heap_malloc(heap,1024);
+  void* q = mi_heap_malloc(heap, 1024);
   // mi_heap_delete(heap); // uncomment to prevent assertion
 }
 
 void heap_no_delete() {
   auto t1 = std::thread(heap_no_delete_worker);
-  t1.join();  
+  t1.join();
 }
 
 
@@ -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;
+}