tls: the pthread key of the thread local slots is created once

CMakeLists.txt

@@ -1147,7 +1147,7 @@ if (MI_BUILD_TESTS)
   enable_testing()
 
   # static link tests
-  set(mi_static_tests api api-fill stress-heaps stress-subprocs stress profile heap-mt heap-teardown heap-delete-race heap-churn heap-release-mt heap-aba heap-burst-destroy abandoned-lazy fork-user-heap snapshot purge-zero park-handoff forced-purge purge-behind-pass free-before-init)
+  set(mi_static_tests api api-fill stress-heaps stress-subprocs stress profile heap-mt heap-teardown heap-delete-race heap-churn heap-release-mt heap-aba heap-burst-destroy abandoned-lazy fork-user-heap snapshot purge-zero park-handoff forced-purge purge-behind-pass free-before-init tls-key-race)
   if(NOT (MI_DEBUG_TSAN OR (MI_TRACK STREQUAL "ASAN") OR MI_DEBUG_UBSAN))
     list(APPEND mi_static_tests thp-optout)  # counts madvise calls by interposing it, which a sanitizer runtime does first
   endif()

include/mimalloc/internal.h

@@ -358,6 +358,7 @@ extern "C" {
 #endif
 extern mi_decl_export _Atomic(uintptr_t) mi_debug_stall_in_thread_theaps_done;
 extern mi_decl_export _Atomic(uintptr_t) mi_debug_stall_in_heap_delete_claim;
+extern mi_decl_export _Atomic(uintptr_t) mi_debug_stall_in_pthread_key_create;
 extern mi_decl_export _Atomic(uintptr_t) mi_debug_abandoned_maps_allocated;
 extern mi_decl_export volatile long      mi_debug_fail_os_commit_after;
 #ifdef __cplusplus
@@ -521,6 +522,7 @@ static inline void __mi_stat_counter_decrease(mi_stat_counter_t* stat, uint64_t
 #endif
 
 mi_decl_noinline bool _mi_pthread_key_create(pthread_key_t* pkey, void (*destruct)(void*), void* init);
+mi_decl_noinline bool _mi_pthread_key_create_once(_Atomic(pthread_key_t)* pkey, void* init);
 
 static inline void* mi_pthread_key_get(pthread_key_t key) {
   #if !MI_PTHREADS_GET_INVALID_KEY_IS_NULL
@@ -529,16 +531,17 @@ static inline void* mi_pthread_key_get(pthread_key_t key) {
   return pthread_getspecific(key);
 }
 
-static inline bool mi_pthread_key_set(pthread_key_t* pkey, void* val) {
-  if mi_likely(*pkey!=MI_PTHREAD_KEY_INVALID) { pthread_setspecific(*pkey,val); return true; }
-  else if (val!=NULL) { return _mi_pthread_key_create(pkey,NULL,val); }
+// A key that is created when a value is first set (by any thread: `pkey` is shared).
+static inline bool mi_pthread_key_set(_Atomic(pthread_key_t)* pkey, void* val) {
+  const pthread_key_t key = mi_atomic_load_acquire(pkey);   // (after the thread that created it; a thread sets a value about once)
+  if mi_likely(key!=MI_PTHREAD_KEY_INVALID) { pthread_setspecific(key,val); return true; }
+  else if (val!=NULL) { return _mi_pthread_key_create_once(pkey,val); }
   else return true;
 }
 
-static inline void mi_pthread_key_delete(pthread_key_t* pkey) {
-  const pthread_key_t key = *pkey;
+static inline void mi_pthread_key_delete(_Atomic(pthread_key_t)* pkey) {
+  const pthread_key_t key = mi_atomic_exchange_relaxed(pkey, MI_PTHREAD_KEY_INVALID);
   if (key!=MI_PTHREAD_KEY_INVALID) {
-    *pkey = MI_PTHREAD_KEY_INVALID;
     pthread_key_delete(key);
   }
 }

src/libc.c

@@ -155,6 +155,44 @@ mi_decl_noinline bool _mi_pthread_key_create(pthread_key_t* pkey, void (*destruc
   mi_assert_internal(pthread_getspecific(*pkey)==init);
   return true;
 }
+
+// For a key that is created when it is first set (`internal.h:mi_pthread_key_set`). Several threads can get here for
+// one key at once: the key of the thread local slots (`threadlocal.c`) is first set when a heap besides the main one
+// is first used, which a program can do on many threads at the same time. One key wins; a thread that stored its value
+// under a key of its own, which the next one replaced, would not find it again.
+#if MI_DEBUG > 0
+mi_decl_export _Atomic(uintptr_t) mi_debug_stall_in_pthread_key_create;   // test hook (test-tls-key-race): 1 has the first thread here wait for a second one
+#endif
+
+mi_decl_noinline bool _mi_pthread_key_create_once(_Atomic(pthread_key_t)* pkey, void* init) {
+  #if MI_DEBUG > 0
+  if (mi_atomic_load_acquire(&mi_debug_stall_in_pthread_key_create) != 0) {
+    if (mi_atomic_increment_acq_rel(&mi_debug_stall_in_pthread_key_create) == 1) {   // the first: both have found no key then
+      for (int i = 0; i < 100000 && mi_atomic_load_acquire(&mi_debug_stall_in_pthread_key_create) == 2; i++) { _mi_prim_thread_yield(); }
+    }
+  }
+  #endif
+  pthread_key_t key;
+  const int err = pthread_key_create(&key,NULL);
+  if mi_unlikely(err!=0) {
+    key = mi_atomic_load_acquire(pkey);   // (another thread may have made it meanwhile)
+    if (key == MI_PTHREAD_KEY_INVALID) {
+      _mi_error_message(ENOMEM,"unable to allocate a thread local variable (error %d)\n", err);
+      return false;
+    }
+  }
+  else {
+    mi_assert_internal(key != MI_PTHREAD_KEY_INVALID);
+    pthread_key_t expected = MI_PTHREAD_KEY_INVALID;
+    if (!mi_atomic_cas_strong_acq_rel(pkey, &expected, key)) {
+      pthread_key_delete(key);   // another thread was first
+      key = expected;
+    }
+  }
+  pthread_setspecific(key,init);
+  mi_assert_internal(pthread_getspecific(key)==init);
+  return true;
+}
 #endif
 
 // --------------------------------------------------------

src/threadlocal.c

@@ -44,8 +44,8 @@ static mi_thread_locals_t mi_thread_locals_empty = mi_init_struct_zero;
 #if MI_TLS_MODEL_PTHREADS || defined(__APPLE__)   // macOS has fast pthreads
 // Use pthreads
 #define mi_define_thread_local(tp,name,initval) \
-  static pthread_key_t __##name##_key = MI_PTHREAD_KEY_INVALID; \
-  static inline tp   name##_peek(void)    { return (tp)mi_pthread_key_get(__##name##_key); } \
+  static _Atomic(pthread_key_t) __##name##_key = MI_ATOMIC_VAR_INIT(MI_PTHREAD_KEY_INVALID); \
+  static inline tp   name##_peek(void)    { return (tp)mi_pthread_key_get(mi_atomic_load_relaxed(&__##name##_key)); } \
   static inline tp   name##_get(void)     { tp result = name##_peek(); return (result!=NULL ? result : initval); } \
   static inline bool name##_set(tp val)   { return mi_pthread_key_set(&__##name##_key,val); } \
   static inline void name##_delete(void)  { mi_pthread_key_delete(&__##name##_key); }

test/test-tls-key-race.c

@@ -0,0 +1,103 @@
+/* ----------------------------------------------------------------------------
+Copyright (c) 2018-2026 Microsoft Research, Daan Leijen
+This is free software; you can redistribute it and/or modify it under the
+terms of the MIT license.
+-----------------------------------------------------------------------------*/
+
+/* The thread local slots of a thread (`threadlocal.c`), which hold the theap of every heap besides the main
+   one, hang off a thread local of their own. Where that is a pthread key (macOS, MI_TLS_MODEL_PTHREADS) the
+   key is created when it is first set, which is when a heap besides the main one is first used: and a program
+   can do that on many threads at once. Every thread has to end up with the same key, or the one whose key was
+   replaced does not find its slots again (and makes a second theap for its heap; a debug build asserts in
+   `heap.c:mi_heap_init_theap`).
+
+   There is one first time in a process. A debug build has a hook that holds the first thread that creates the
+   key until a second one is there too (`libc.c:mi_debug_stall_in_pthread_key_create`), which makes that one
+   time the race; without it the threads make their heaps, wait for each other, and allocate from them at once.
+   Where thread locals are not pthread keys this only checks that it works.
+
+   > mimalloc-test-tls-key-race
+*/
+
+#include <stdio.h>
+#include <stdlib.h>
+#include <stdint.h>
+#include <stdbool.h>
+
+#include "mimalloc.h"
+#include "mimalloc-stats.h"
+
+#if defined(_WIN32) || defined(__EMSCRIPTEN__)
+int main(void) { printf("skipped (no pthreads)\n"); return 0; }
+#else
+
+#include <pthread.h>
+#include <sched.h>
+#include <stdatomic.h>
+
+#if MI_DEBUG > 0
+extern _Atomic(uintptr_t) mi_debug_stall_in_pthread_key_create;
+#endif
+
+#define THREADS 8
+
+static atomic_int ready, go, bad;
+
+// A thread that lost its slots does not notice: it makes new ones, and a second theap for its heap, and goes on
+// with those. The count of the theaps that were made tells.
+static size_t theaps_made(void) {
+  static mi_stats_t stats;
+  stats.size = sizeof(stats);
+  stats.version = MI_STAT_VERSION;
+  if (!mi_subproc_stats_get_exclusive(mi_subproc_main(), &stats)) return 0;
+  return (size_t)stats.theaps.total;
+}
+
+static void* worker(void* arg) {
+  (void)arg;
+  mi_heap_t* const heap = mi_heap_new();   // (takes a lock: the threads do not get here together)
+  atomic_fetch_add(&ready, 1);
+  while (!atomic_load(&go)) { sched_yield(); }
+  void* const p = mi_heap_malloc(heap, 64);   // the first use of the thread local slots
+  mi_theap_t* const theap = mi_heap_theap(heap);
+  void* const q = mi_heap_malloc(heap, 64);
+  if (p == NULL || q == NULL || theap == NULL || mi_heap_theap(heap) != theap || !mi_heap_contains(heap, p) || !mi_heap_contains(heap, q)) {
+    atomic_fetch_add(&bad, 1);
+  }
+  mi_free(p);
+  mi_free(q);
+  mi_heap_delete(heap);
+  return NULL;
+}
+
+int main(void) {
+  mi_free(mi_malloc(8));   // the main heap only
+  #if MI_DEBUG > 0
+  atomic_store(&mi_debug_stall_in_pthread_key_create, (uintptr_t)1);
+  #endif
+  const size_t theaps_before = theaps_made();
+  pthread_t threads[THREADS];
+  int started = 0;
+  for (; started < THREADS; started++) {
+    if (pthread_create(&threads[started], NULL, &worker, NULL) != 0) break;
+  }
+  while (atomic_load(&ready) < started) { sched_yield(); }
+  atomic_store(&go, 1);
+  for (int i = 0; i < started; i++) { pthread_join(threads[i], NULL); }
+  // every thread makes a theap for the main heap (its heap is allocated there) and one for its own
+  const size_t made = theaps_made() - theaps_before;
+  const bool counted = (theaps_made() != 0);   // (not in a build without statistics)
+  bool ok = (started >= 2 && atomic_load(&bad) == 0 && (!counted || made == 2*(size_t)started));
+  printf("%d threads made %zu theaps, %d of them saw its heap change", started, made, atomic_load(&bad));
+  #if MI_DEBUG > 0
+  const uintptr_t arrived = atomic_exchange(&mi_debug_stall_in_pthread_key_create, (uintptr_t)0) - 1;
+  printf("; %d created the key at the same time", (int)arrived);
+  #if MI_TLS_MODEL_PTHREADS || defined(__APPLE__)
+  ok = ok && (arrived >= 2);   // or this was not the race
+  #endif
+  #endif
+  printf("\n");
+  return (ok ? 0 : 1);
+}
+
+#endif