use Tls expansion slots on windows if direct slots are exhausted (see also PR #1202)

Author: daanx

CMakeLists.txt

@@ -18,7 +18,7 @@ option(MI_SEE_ASM           "Generate assembly files" OFF)
 option(MI_OSX_INTERPOSE     "Use interpose to override standard malloc on macOS" ON)
 option(MI_OSX_ZONE          "Use malloc zone to override standard malloc on macOS" ON)
 option(MI_WIN_REDIRECT      "Use redirection module ('mimalloc-redirect') on Windows if compiling mimalloc as a DLL" ON)
-option(MI_WIN_USE_FIXED_TLS "Use a fixed TLS slot on Windows to avoid extra tests in the malloc fast path" OFF)
+option(MI_WIN_DIRECT_TLS    "Use only direct TLS slots on Windows to avoid extra tests in the malloc fast path (only works if the program uses less than 64 TlsAlloc'd slots in total)" OFF)
 option(MI_LOCAL_DYNAMIC_TLS "Use local-dynamic-tls, a slightly slower but dlopen-compatible thread local storage mechanism (Unix)" OFF)
 option(MI_LIBC_MUSL         "Enable this when linking with musl libc" OFF)
 
@@ -394,9 +394,9 @@ if(MI_WIN_USE_FLS)
   list(APPEND mi_defines MI_WIN_USE_FLS=1)
 endif()
 
-if(MI_WIN_USE_FIXED_TLS)
-  message(STATUS "Use fixed TLS slot on Windows to avoid extra tests in the malloc fast path (MI_WIN_USE_FIXED_TLS=ON)")
-  list(APPEND mi_defines MI_WIN_USE_FIXED_TLS=1)
+if(MI_WIN_DIRECT_TLS)
+  message(STATUS "Use only direct TLS slots on Windows to avoid extra tests in the malloc fast path -- this only works if the program uses less than 64 TlsAlloc'd slots in total! (MI_WIN_USE_ONLY_DIRECT_TLS=ON)")
+  list(APPEND mi_defines MI_WIN_DIRECT_TLS=1)
 endif()
 
 # Check /proc/cpuinfo for an SV39 MMU and limit the virtual address bits.

include/mimalloc/prim.h

@@ -408,15 +408,38 @@ static inline mi_theap_t* _mi_theap_cached(void) {
 // Dynamic TLS slot (windows)
 #define MI_THEAP_INITASNULL  1
 
+// We try to use direct slots, but can also use the expansion slots (upto 1024 available)
+#if MI_SIZE_SIZE==4
+#define MI_TLS_EXPANSION_SLOT    (0x0F94 / MI_SIZE_SIZE)
+#else
+#define MI_TLS_EXPANSION_SLOT    (0x1780 / MI_SIZE_SIZE)
+#endif
+
 extern mi_decl_hidden size_t _mi_theap_default_slot;
 extern mi_decl_hidden size_t _mi_theap_cached_slot;
+extern mi_decl_hidden size_t _mi_theap_default_expansion_slot;
+extern mi_decl_hidden size_t _mi_theap_cached_expansion_slot;
 
 static inline mi_theap_t* _mi_theap_default(void) {
-  return (mi_theap_t*)mi_prim_tls_slot(_mi_theap_default_slot); // valid initial "last user slot" so it returns NULL at first leading to slot initialization
+  const size_t slot = _mi_theap_default_slot;
+  mi_theap_t* theap  = (mi_theap_t*)mi_prim_tls_slot(slot);
+  #if !MI_WIN_DIRECT_TLS
+  if (slot==MI_TLS_EXPANSION_SLOT && theap!=NULL) {  // in initialized TlsExpansionSlots ?
+    theap = ((mi_theap_t**)theap)[_mi_theap_default_expansion_slot];
+  }
+  #endif
+  return theap;
 }
 
 static inline mi_theap_t* _mi_theap_cached(void) {
-  return (mi_theap_t*)mi_prim_tls_slot(_mi_theap_cached_slot);
+  const size_t slot = _mi_theap_cached_slot;
+  mi_theap_t* theap = (mi_theap_t*)mi_prim_tls_slot(slot);
+  #if !MI_WIN_DIRECT_TLS
+  if (slot==MI_TLS_EXPANSION_SLOT && theap!=NULL) {  // in initialized TlsExpansionSlots ?
+    theap = ((mi_theap_t**)theap)[_mi_theap_cached_expansion_slot];
+  }
+  #endif
+  return theap;
 }
 
 #elif MI_TLS_MODEL_DYNAMIC_PTHREADS
@@ -479,7 +502,7 @@ static inline mi_theap_t* _mi_heap_theap_peek(const mi_heap_t* heap) {
 static inline mi_theap_t* _mi_page_associated_theap_peek(mi_page_t* page) {
   mi_heap_t* const heap = page->heap;
   mi_theap_t* theap;
-  if mi_likely(heap==NULL) { theap = __mi_theap_main; }  // note: on macOS accessing the thread_local can cause allocation during thread shutdown (and reinitialize the thread)! 
+  if mi_likely(heap==NULL) { theap = __mi_theap_main; }  // note: on macOS accessing the thread_local can cause allocation during thread shutdown (and reinitialize the thread)!
                       else { theap = _mi_heap_theap_peek(heap); }
   mi_assert_internal(theap==NULL || _mi_thread_id()==theap->tld->thread_id);
   return theap;

src/alloc.c

@@ -252,7 +252,7 @@ mi_decl_nodiscard extern inline mi_decl_restrict void* mi_theap_malloc(mi_theap_
 }
 
 mi_decl_nodiscard mi_decl_restrict void* mi_malloc(size_t size) mi_attr_noexcept {
-  return mi_theap_malloc(_mi_theap_default(), size);
+   return mi_theap_malloc(_mi_theap_default(), size);
 }
 
 mi_decl_nodiscard mi_decl_restrict void* mi_heap_malloc(mi_heap_t* heap, size_t size) mi_attr_noexcept {

src/init.c

@@ -682,32 +682,83 @@ mi_decl_cold mi_decl_noinline mi_theap_t* _mi_theap_empty_get(void) {
 }
 
 #if MI_TLS_MODEL_DYNAMIC_WIN32
+#include <winternl.h>
 
-// only for win32 for now
+// If we can, we use one of the 64 direct TLS slots (but fall back to expansion slots if needed)
 #if MI_SIZE_SIZE==4
-#define MI_TLS_USER_BASE  (0x0E10 / MI_SIZE_SIZE)
+#define MI_TLS_DIRECT_FIRST             (0x0E10 / MI_SIZE_SIZE)
 #else
-#define MI_TLS_USER_BASE  (0x1480 / MI_SIZE_SIZE)
+#define MI_TLS_DIRECT_FIRST             (0x1480 / MI_SIZE_SIZE)
 #endif
-#define MI_TLS_USER_LAST_SLOT  (MI_TLS_USER_BASE + 63)
+#define MI_TLS_DIRECT_SLOTS             (64)
+#define MI_TLS_EXPANSION_SLOTS          (1024)
 
-// we initially use the last user slot so NULL is returned
-// when allocating a slot, we check we get a slot before the last one (so it wasn't used yet)
-mi_decl_hidden size_t _mi_theap_default_slot = MI_TLS_USER_LAST_SLOT;
-mi_decl_hidden size_t _mi_theap_cached_slot  = MI_TLS_USER_LAST_SLOT;
+#if !MI_WIN_DIRECT_TLS
+#define MI_TLS_INITIAL_SLOT             MI_TLS_EXPANSION_SLOT
+#define MI_TLS_INITIAL_EXPANSION_SLOT   (MI_TLS_EXPANSION_SLOTS-1)
+#else
+// with only direct entries, use the "arbitrary user data" field 
+// and assume it is NULL (see also <http://www.nynaeve.net/?p=98>)
+#define MI_TLS_INITIAL_EXPANSION_SLOT   (0)
+#if MI_SIZE_SIZE==4
+#define MI_TLS_INITIAL_SLOT             (0x14 / MI_SIZE_SIZE)
+#else
+#define MI_TLS_INITIAL_SLOT             (0x28 / MI_SIZE_SIZE)
+#endif
+#endif
+
+// we initially use the last of the expansion slots as the default NULL.
+// note: this will fail if the program allocates exactly 1024+64 slots with TlsAlloc (which is quite unlikely)
+mi_decl_hidden mi_decl_cache_align size_t _mi_theap_default_slot = MI_TLS_INITIAL_SLOT;
+mi_decl_hidden size_t _mi_theap_default_expansion_slot = MI_TLS_INITIAL_EXPANSION_SLOT;
+mi_decl_hidden size_t _mi_theap_cached_slot            = MI_TLS_INITIAL_SLOT;
+mi_decl_hidden size_t _mi_theap_cached_expansion_slot  = MI_TLS_INITIAL_EXPANSION_SLOT;
+
+static size_t mi_win_tls_slot_alloc(size_t* extended) {
+  const DWORD slot = TlsAlloc();
+  if (slot==TLS_OUT_OF_INDEXES || slot >= MI_TLS_DIRECT_SLOTS + MI_TLS_EXPANSION_SLOTS - 1) {
+    // note: we also fail if the program already allocated the maximum number of expansion slots (as we use the last one as the default)
+    *extended = 0;
+    return 0;
+  }
+  else if (slot<MI_TLS_DIRECT_SLOTS) {
+    *extended = 0;
+    return (slot + MI_TLS_DIRECT_FIRST);
+  }
+  else {
+    #if MI_WIN_DIRECT_TLS
+    *extended = 0;
+    return 0;
+    #else
+    *extended = (slot - MI_TLS_DIRECT_SLOTS);
+    return MI_TLS_EXPANSION_SLOT;
+    #endif
+  }
+}
 
-mi_decl_cold mi_theap_t* _mi_tls_slots_init(void) {
+mi_decl_cold mi_theap_t* _mi_win_tls_slots_init(void) {
   static mi_atomic_once_t tls_slots_init;
   if (mi_atomic_once(&tls_slots_init)) {
-    _mi_theap_default_slot = TlsAlloc() + MI_TLS_USER_BASE;
-    _mi_theap_cached_slot  = TlsAlloc() + MI_TLS_USER_BASE;
-    if (_mi_theap_cached_slot >= MI_TLS_USER_LAST_SLOT) {
+    _mi_theap_default_slot = mi_win_tls_slot_alloc(&_mi_theap_default_expansion_slot);
+    _mi_theap_cached_slot = mi_win_tls_slot_alloc(&_mi_theap_cached_expansion_slot);
+    if (_mi_theap_cached_slot==0) {
       _mi_error_message(EFAULT, "unable to allocate fast TLS user slot (0x%zx)\n", _mi_theap_cached_slot);
     }
   }
   return (mi_theap_t*)&_mi_theap_empty;
 }
 
+static void mi_win_tls_slot_set(size_t slot, size_t extended_slot, void* value) {
+  mi_assert_internal((slot >= MI_TLS_DIRECT_FIRST && slot < MI_TLS_DIRECT_FIRST + MI_TLS_DIRECT_SLOTS) || slot == MI_TLS_EXPANSION_SLOT);
+  if (slot < MI_TLS_DIRECT_FIRST + MI_TLS_DIRECT_SLOTS) {
+    mi_prim_tls_slot_set(slot, value);
+  }
+  else {
+    mi_assert_internal(extended_slot < MI_TLS_EXPANSION_SLOTS);
+    TlsSetValue((DWORD)(extended_slot + MI_TLS_DIRECT_SLOTS), value);  // use TlsSetValue to initialize the TlsExpansion array if needed
+  }
+}
+
 #elif MI_TLS_MODEL_DYNAMIC_PTHREADS
 
 // only for pthreads for now
@@ -731,8 +782,8 @@ void _mi_theap_cached_set(mi_theap_t* theap) {
   #elif MI_TLS_MODEL_FIXED_SLOT
     mi_prim_tls_slot_set(MI_TLS_MODEL_FIXED_SLOT_CACHED, theap);
   #elif MI_TLS_MODEL_DYNAMIC_WIN32
-    _mi_tls_slots_init();
-    mi_prim_tls_slot_set(_mi_theap_cached_slot, theap);
+    _mi_win_tls_slots_init();
+    mi_win_tls_slot_set(_mi_theap_cached_slot, _mi_theap_cached_expansion_slot, theap);
   #elif MI_TLS_MODEL_DYNAMIC_PTHREADS
     _mi_tls_keys_init();
     if (_mi_theap_cached_key!=0) pthread_setspecific(_mi_theap_cached_key, theap);
@@ -747,8 +798,8 @@ void _mi_theap_default_set(mi_theap_t* theap)  {
   #elif MI_TLS_MODEL_FIXED_SLOT
     mi_prim_tls_slot_set(MI_TLS_MODEL_FIXED_SLOT_DEFAULT, theap);
   #elif MI_TLS_MODEL_DYNAMIC_WIN32
-    _mi_tls_slots_init();
-    mi_prim_tls_slot_set(_mi_theap_default_slot, theap);
+    _mi_win_tls_slots_init();
+    mi_win_tls_slot_set(_mi_theap_default_slot, _mi_theap_default_expansion_slot, theap);
   #elif MI_TLS_MODEL_DYNAMIC_PTHREADS
     _mi_tls_keys_init();
     if (_mi_theap_default_key!=0) pthread_setspecific(_mi_theap_default_key, theap);