From 13ee94cef6900539ad5f4abf322efdfc3650cfc9 Mon Sep 17 00:00:00 2001
From: daanx <daan@effp.org>
Date: Sun, 15 Dec 2024 13:22:00 -0800
Subject: [PATCH] fix concurrent mi_tld access bug

---
 src/arena-meta.c |  4 ++--
 src/arena.c      |  1 +
 src/init.c       | 28 ++++++++++++++++++----------
 3 files changed, 21 insertions(+), 12 deletions(-)

diff --git a/src/arena-meta.c b/src/arena-meta.c
index bc98d3f9..ceda06ba 100644
--- a/src/arena-meta.c
+++ b/src/arena-meta.c
@@ -93,7 +93,7 @@ static mi_meta_page_t* mi_meta_page_zalloc(void) {
 
 
 // allocate meta-data
-void* _mi_meta_zalloc( size_t size, mi_memid_t* pmemid )
+mi_decl_noinline void* _mi_meta_zalloc( size_t size, mi_memid_t* pmemid )
 {
   mi_assert_internal(pmemid != NULL);
   size = _mi_align_up(size,MI_META_BLOCK_SIZE);
@@ -133,7 +133,7 @@ void* _mi_meta_zalloc( size_t size, mi_memid_t* pmemid )
 }
 
 // free meta-data
-void _mi_meta_free(void* p, size_t size, mi_memid_t memid) {
+mi_decl_noinline void _mi_meta_free(void* p, size_t size, mi_memid_t memid) {
   if (p==NULL) return;
   if (memid.memkind == MI_MEM_META) {
     mi_assert_internal(_mi_divide_up(size, MI_META_BLOCK_SIZE) == memid.mem.meta.block_count);
diff --git a/src/arena.c b/src/arena.c
index 7aec429e..d8b882d3 100644
--- a/src/arena.c
+++ b/src/arena.c
@@ -551,6 +551,7 @@ static mi_page_t* mi_arena_page_try_find_abandoned(size_t slice_count, size_t bl
 
   // any abandoned in our size class?
   mi_subproc_t* const subproc = tld->subproc;
+  mi_assert_internal(subproc != NULL);
   if (mi_atomic_load_relaxed(&subproc->abandoned_count[bin]) == 0) return NULL;
 
   // search arena's
diff --git a/src/init.c b/src/init.c
index 8f1449a3..c103f521 100644
--- a/src/init.c
+++ b/src/init.c
@@ -134,7 +134,8 @@ static mi_decl_cache_align mi_subproc_t mi_subproc_default;
 
 static mi_decl_cache_align mi_tld_t tld_main = {
   0,
-  &_mi_heap_main, &_mi_heap_main,
+  &_mi_heap_main,         // heap_backing
+  &_mi_heap_main,         // heaps list
   &mi_subproc_default,    // subproc
   0,                      // tseq
   MI_MEMID_STATIC,        // memid
@@ -271,10 +272,23 @@ static mi_tld_t* mi_tld_alloc(void) {
   }
 }
 
-mi_tld_t* _mi_tld(void) {
+#define MI_TLD_INVALID  ((mi_tld_t*)1)
+
+static mi_decl_noinline void mi_tld_free(void) {
+  mi_tld_t* tld = _mi_tld();
+  mi_tld = MI_TLD_INVALID;
+  _mi_meta_free(tld, sizeof(mi_tld_t), tld->memid);
+}
+
+mi_tld_t* mi_decl_noinline _mi_tld(void) {
+  if (mi_tld == MI_TLD_INVALID) {
+    _mi_error_message(EFAULT, "internal error: tld accessed after the thread terminated\n");
+    abort();
+    mi_tld = NULL;
+  }
   if (mi_tld==NULL) {
     mi_tld = mi_tld_alloc();
-  }
+  }  
   return mi_tld;
 }
 
@@ -409,9 +423,6 @@ static bool _mi_thread_heap_done(mi_heap_t* heap) {
     #endif
   }
 
-  // free the tld
-  mi_tld_t* tld = _mi_tld();
-  _mi_meta_free(_mi_tld(), sizeof(mi_tld_t), tld->memid);
   return false;
 }
 
@@ -497,10 +508,7 @@ void _mi_thread_done(mi_heap_t* heap)
   _mi_thread_heap_done(heap);  // returns true if already ran
 
   // free thread local data
-  if (mi_tld != NULL) {
-    _mi_meta_free(mi_tld, sizeof(mi_tld_t), mi_tld->memid);
-    mi_tld = NULL;
-  }
+  mi_tld_free();
 }
 
 void _mi_heap_set_default_direct(mi_heap_t* heap)  {