allow initial main thread to terminate before the process terminates (see issue #1287)

Author: daanx

include/mimalloc/prim.h

@@ -248,10 +248,9 @@ extern mi_decl_hidden bool _mi_process_is_initialized;                // has mi_
 // fast path of `_mi_free` and we specialize for various platforms as
 // inlined definitions. Regular code should call `init.c:_mi_thread_id()`.
 // We only require _mi_prim_thread_id() to return a unique id
-// for each thread (unequal to zero).
+// for each thread (unequal to zero) with the bottom 2 bits clear.
 //-------------------------------------------------------------------
 
-
 // Do we have __builtin_thread_pointer? This would be the preferred way to get a unique thread id
 // but unfortunately, it seems we cannot test for this reliably at this time (see issue #883)
 // Nevertheless, it seems needed on older graviton platforms (see issue #851).

src/alloc.c

@@ -133,7 +133,7 @@ static mi_decl_forceinline mi_decl_restrict void* mi_theap_malloc_small_zero_non
   mi_assert(size <= MI_SMALL_SIZE_MAX);
   #if MI_DEBUG
   const uintptr_t tid = _mi_thread_id();
-  mi_assert(theap->tld->thread_id == 0 || theap->tld->thread_id == tid); // theaps are thread local
+  mi_assert(theap->tld->thread_id == ~0UL || theap->tld->thread_id == tid); // theaps are thread local
   #endif
   #if (MI_PADDING || MI_GUARDED)
   if mi_unlikely(size == 0) { size = sizeof(void*); }

src/init.c

@@ -163,11 +163,13 @@ mi_decl_cache_align const mi_theap_t _mi_theap_empty_wrong = {
 
 // Heap for the main thread
 
+#define MI_THREADID_UNINIT ((mi_threadid_t)(~0))
+
 extern mi_decl_hidden mi_decl_cache_align mi_theap_t theap_main;
 extern mi_decl_hidden mi_decl_cache_align mi_heap_t  heap_main;
 
 static mi_decl_cache_align mi_tld_t tld_main = {
-  0,                      // thread_id
+  MI_THREADID_UNINIT,     // thread_id
   0,                      // thread_seq
   0,                      // numa node
   &subproc_main,          // subproc
@@ -214,7 +216,7 @@ mi_decl_cache_align mi_heap_t heap_main
 mi_decl_hidden mi_decl_thread mi_theap_t* __mi_theap_main = NULL;
 
 mi_threadid_t _mi_thread_id(void) mi_attr_noexcept {
-  mi_threadid_t tid = _mi_prim_thread_id();
+  const mi_threadid_t tid = _mi_prim_thread_id();
   mi_assert_internal( (tid & 0x03) == 0 ); // mimalloc reserves the bottom 2 bits
   return tid;
 }
@@ -298,7 +300,7 @@ static void mi_subproc_main_init(void) {
 
 // Initialize main tld
 static void mi_tld_main_init(void) {
-  if (tld_main.thread_id == 0) {
+  if (tld_main.thread_id == MI_THREADID_UNINIT) {
     tld_main.thread_id = _mi_prim_thread_id();
     mi_lock_init(&tld_main.theaps_lock);
   }
@@ -379,17 +381,11 @@ static mi_tld_t* mi_tld_alloc(void) {
 #define MI_TLD_INVALID  ((mi_tld_t*)1)
 
 mi_decl_noinline static void mi_tld_free(mi_tld_t* tld) {
-  mi_lock_done(&tld->theaps_lock);
-  if (tld != NULL && tld != MI_TLD_INVALID) {
-    mi_atomic_decrement_relaxed(&tld->subproc->thread_count);
-    _mi_meta_free(tld, sizeof(mi_tld_t), tld->memid);
-  }
-  #if 0
-  // do not read/write to `thread_tld` on older macOS <= 14 as that will re-initialize the thread local storage
-  // (since we are calling this during pthread shutdown)
-  // (and this could happen on other systems as well, so let's never do it)
-  thread_tld = MI_TLD_INVALID;
-  #endif
+  if (tld==NULL || tld==MI_TLD_INVALID) return; 
+  mi_atomic_decrement_relaxed(&tld->subproc->thread_count);
+  tld->thread_id = 0;
+  mi_lock_done(&tld->theaps_lock);  
+  _mi_meta_free(tld, sizeof(mi_tld_t), tld->memid);  // note: safe for static tld_main
 }
 
 // return the thread local heap ensuring it is initialized (and not `NULL` or `&_mi_theap_empty`);
@@ -704,7 +700,7 @@ static void mi_process_setup_auto_thread_done(void) {
 
 
 bool _mi_is_main_thread(void) {
-  return (tld_main.thread_id==0 || tld_main.thread_id == _mi_thread_id());
+  return (tld_main.thread_id==MI_THREADID_UNINIT || tld_main.thread_id == _mi_thread_id());
 }
 
 

src/prim/prim.c

@@ -27,7 +27,7 @@ terms of the MIT license. A copy of the license can be found in the file
 #endif
 
 // Generic process initialization
-#ifndef MI_PRIM_HAS_PROCESS_ATTACH
+#if !defined(MI_PRIM_HAS_PROCESS_ATTACH)
 #if defined(__GNUC__) || defined(__clang__)
   // gcc,clang: use the constructor/destructor attribute
   // which for both seem to run before regular constructors/destructors

src/theap.c

@@ -289,7 +289,7 @@ void _mi_theap_incref(mi_theap_t* theap) {
 }
 
 void _mi_theap_decref(mi_theap_t* theap) {
-  if (theap!=NULL && theap->memid.memkind > MI_MEM_STATIC) {
+  if (theap!=NULL && !mi_memid_needs_no_free(theap->memid)) {
     if (mi_atomic_decrement_acq_rel(&theap->refcount) == 1) {
       mi_theap_free_mem(theap);
     }
@@ -324,8 +324,9 @@ bool _mi_theap_free(mi_theap_t* theap, bool acquire_heap_theaps_lock, bool acqui
       if (theap->tnext != NULL) { theap->tnext->tprev = theap->tprev;  }
       if (theap->tprev != NULL) { theap->tprev->tnext = theap->tnext;  }
                           else { mi_assert_internal(theap->tld->theaps == theap); theap->tld->theaps = theap->tnext; }
-      theap->tnext = theap->tprev = NULL;                        
+      theap->tnext = theap->tprev = NULL;           
     }
+
     theap->tld = NULL;
     _mi_theap_decref(theap);
     return true;

test/main-static-dep.cpp

@@ -24,12 +24,12 @@ static Static s = Static();
 
 void Test(void) {
   char* s = mi_mallocn_tp(char, 128);
-  strcpy_s(s, 128, "hello world!");
+  strlcpy(s, "hello world!", 128);
   printf("message from static dll: %s\n", s);
   mi_free(s);
 }
 
-
+#ifdef WIN32
 #include <windows.h>
 
 BOOL WINAPI DllMain(HINSTANCE module, DWORD reason, LPVOID reserved) {
@@ -44,3 +44,4 @@ BOOL WINAPI DllMain(HINSTANCE module, DWORD reason, LPVOID reserved) {
   }
   return TRUE;
 }
+#endif

test/main-static-dep.h

@@ -1,11 +1,17 @@
 #pragma once
 
-typedef void (__stdcall * TestFun)(void);
-
 #if __cplusplus
-extern "C" 
+extern "C" {
 #endif
-__declspec(dllexport) void __cdecl Test(void);
-
 
+#ifdef WIN32
+typedef void (__cdecl *TestFun)(void);
+__declspec(dllexport) void __cdecl Test(void);
+#else
+typedef void (*TestFun)(void);
+void Test(void);
+#endif
 
+#if __cplusplus
+}
+#endif

test/main-static.cpp

@@ -0,0 +1,44 @@
+#include <stdlib.h>
+#include <stdio.h>
+#include <assert.h>
+#include <string.h>
+#include <stdint.h>
+
+#include <mimalloc.h>
+#include <new>
+#include <vector>
+#include <future>
+#include <iostream>
+#include <thread>
+#include <random>
+#include <chrono>
+#include <assert.h>
+
+#include <dlfcn.h>
+
+#include "main-static-dep.h"
+
+TestFun fun;
+void*   so;
+
+void testso() {
+  fun();
+}
+
+void loadso() {
+  so = dlopen("./libstatic.so", RTLD_LAZY);
+  fun = (TestFun)dlsym(so,"Test");
+  testso();
+}
+
+static void test_static(void) {
+  auto t1 = std::thread(&loadso);
+  t1.join();
+  auto t2 = std::thread(&testso);
+  t2.join();
+}
+
+int main(int argc, char** argv) {
+  test_static();
+  return 0;
+}