From a278db940bb25c294c60cd258b4051df7ae5187f Mon Sep 17 00:00:00 2001
From: daan <daanl@outlook.com>
Date: Tue, 5 May 2020 10:10:46 -0700
Subject: [PATCH] improve thread-id check for thread local FLS callbacks on
 Windows with static linking

---
 src/init.c | 44 +++++++++++++++++++++-----------------------
 1 file changed, 21 insertions(+), 23 deletions(-)

diff --git a/src/init.c b/src/init.c
index 2ccc0d83..132043e8 100644
--- a/src/init.c
+++ b/src/init.c
@@ -216,26 +216,24 @@ 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) {
-      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;
-    }
-    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);
+  // 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;
   }
+  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);
+  }
+  
 
   // merge stats
   _mi_stats_done(&heap->tld->stats);
@@ -296,10 +294,6 @@ 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)
@@ -361,12 +355,16 @@ void mi_thread_done(void) mi_attr_noexcept {
 }
 
 static void _mi_thread_done(mi_heap_t* heap) {
+  // 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()) return;
+
   // stats
   if (!_mi_is_main_thread() && mi_heap_is_initialized(heap))  {
     _mi_stat_decrease(&heap->tld->stats.threads, 1);
   }
+
   // abandon the thread local heap
-  if (_mi_heap_done(heap)) return; // returns true if already ran
+  if (_mi_heap_done(heap)) return;  // returns true if already ran
 }
 
 void _mi_heap_set_default_direct(mi_heap_t* heap)  {