Improving handling first allocation being for TLS (#767)

* handle reentrancy during initialization

* use finialization list if possible

* Add test for reentrancy of C++ destructors and allocation

* Add test for reentrancy of setspecific

* Add new mode for directly calling __cxa_thread_atexit_impl

---------

Co-authored-by: Schrodinger ZHU Yifan <yifanzhu@rochester.edu>

src/test/func/multi_atexit/multi_atexit.cc

@@ -0,0 +1,44 @@
+#ifndef __has_feature
+#  define __has_feature(x) 0
+#endif
+
+// These test partially override the libc malloc/free functions to test
+// interesting corner cases.  This breaks the sanitizers as they will be
+// partially overridden.  So we disable the tests if any of the sanitizers are
+// enabled.
+#if defined(__linux__) && !__has_feature(address_sanitizer) && \
+  !defined(__SANITIZE_ADDRESS__) && !defined(__SANITIZE_THREAD__) && \
+  !defined(SNMALLOC_THREAD_SANITIZER_ENABLED)
+#  define RUN_TEST
+#endif
+
+#ifdef RUN_TEST
+#  include <snmalloc/snmalloc.h>
+#  include <stdlib.h>
+
+void do_nothing() {}
+
+// We only selectively override these functions. Otherwise, malloc may be called
+// before atexit triggers the first initialization attempt.
+
+extern "C" void* calloc(size_t num, size_t size)
+{
+  return snmalloc::libc::calloc(num, size);
+}
+
+extern "C" void free(void* p)
+{
+  if (snmalloc::is_owned(p))
+    return snmalloc::libc::free(p);
+  // otherwise, just leak the memory
+}
+
+#endif
+
+int main()
+{
+#ifdef RUN_TEST
+  for (int i = 0; i < 8192; ++i)
+    atexit(do_nothing);
+#endif
+}

src/test/func/multi_setspecific/multi_setspecific.cc

@@ -0,0 +1,87 @@
+#ifndef __has_feature
+#  define __has_feature(x) 0
+#endif
+
+// These test partially override the libc malloc/free functions to test
+// interesting corner cases.  This breaks the sanitizers as they will be
+// partially overridden.  So we disable the tests if any of the sanitizers are
+// enabled.
+#if defined(__linux__) && !__has_feature(address_sanitizer) && \
+  !defined(__SANITIZE_ADDRESS__) && !defined(__SANITIZE_THREAD__) && \
+  !defined(SNMALLOC_THREAD_SANITIZER_ENABLED)
+#  define RUN_TEST
+#endif
+
+#ifdef RUN_TEST
+#  include <array>
+#  include <snmalloc/snmalloc.h>
+#  include <stdlib.h>
+#  include <thread>
+
+// A key in the second "second level" block of the pthread key table.
+// First second level block is statically allocated.
+// This is be 33.
+pthread_key_t key;
+
+void thread_setspecific()
+{
+  // If the following line is uncommented then the test will pass.
+  // free(calloc(1, 1));
+  pthread_setspecific(key, (void*)1);
+}
+
+// We only selectively override these functions. Otherwise, malloc may be called
+// before atexit triggers the first initialization attempt.
+
+extern "C" void* calloc(size_t num, size_t size)
+{
+  snmalloc::message("calloc({}, {})", num, size);
+  return snmalloc::libc::calloc(num, size);
+}
+
+extern "C" void free(void* p)
+{
+  snmalloc::message("free({})", p);
+  // Just leak it
+  if (snmalloc::is_owned(p))
+    return snmalloc::libc::free(p);
+}
+
+void callback(void*)
+{
+  snmalloc::message("callback");
+}
+
+int main()
+{
+  // The first 32 keys are statically allocated, so we need to create 33 keys
+  // to create a key for which pthread_setspecific will call the calloc.
+  for (size_t i = 0; i < 33; i++)
+  {
+    pthread_key_create(&key, callback);
+  }
+
+  // The first calloc occurs here, after the first [0, 32] keys have been
+  // created thus snmalloc will choose the key 33, `key` contains the key `32`
+  // and snmalloc `33`. Both of these keys are not in the statically allocated
+  // part of the pthread key space.
+  std::thread(thread_setspecific).join();
+
+  // There should be a single allocator that can be extracted.
+  if (snmalloc::AllocPool<snmalloc::Config>::extract() == nullptr)
+  {
+    // The thread has not torn down its allocator.
+    snmalloc::report_fatal_error(
+      "Teardown of thread allocator has not occurred.");
+    return 1;
+  }
+
+  return 0;
+}
+#else
+int main()
+{
+  // This test is not run, but it is used to check that the code compiles.
+  return 0;
+}
+#endif

src/test/func/multi_threadatexit/multi_threadatexit.cc

@@ -0,0 +1,63 @@
+#include <thread>
+
+#ifndef __has_feature
+#  define __has_feature(x) 0
+#endif
+
+// These test partially override the libc malloc/free functions to test
+// interesting corner cases.  This breaks the sanitizers as they will be
+// partially overridden.  So we disable the tests if any of the sanitizers are
+// enabled.
+#if defined(__linux__) && !__has_feature(address_sanitizer) && \
+  !defined(__SANITIZE_ADDRESS__) && !defined(__SANITIZE_THREAD__) && \
+  !defined(SNMALLOC_THREAD_SANITIZER_ENABLED)
+#  define RUN_TEST
+#endif
+
+#ifdef RUN_TEST
+#  include <snmalloc/snmalloc.h>
+#  include <stdlib.h>
+
+template<size_t N, size_t M = 0>
+void thread_destruct()
+{
+  snmalloc::message("thread_destruct<{}, {}> start", N, M);
+  static thread_local snmalloc::OnDestruct destruct{
+    []() { snmalloc::message("thread_destruct<{}, {}> destructor", N, M); }};
+  snmalloc::message("thread_destruct<{}, {}> end", N, M);
+
+  if constexpr (N > M + 1)
+  {
+    // destructor
+    thread_destruct<N, (M + N) / 2>();
+    thread_destruct<(M + N) / 2, M>();
+  }
+}
+
+// We only selectively override these functions. Otherwise, malloc may be called
+// before atexit triggers the first initialization attempt.
+
+extern "C" void* calloc(size_t num, size_t size)
+{
+  snmalloc::message("calloc({}, {})", num, size);
+  return snmalloc::libc::calloc(num, size);
+}
+
+extern "C" void free(void* p)
+{
+  snmalloc::message("free({})", p);
+  if (snmalloc::is_owned(p))
+    return snmalloc::libc::free(p);
+  // otherwise, just leak the memory
+}
+
+int main()
+{
+  std::thread(thread_destruct<1000>).join();
+}
+#else
+int main()
+{
+  return 0;
+}
+#endif

src/test/func/thread_alloc_external/thread_alloc_external.cc

@@ -1,6 +1,12 @@
 #ifdef SNMALLOC_USE_PTHREAD_DESTRUCTORS
 #  undef SNMALLOC_USE_PTHREAD_DESTRUCTORS
 #endif
+#ifdef SNMALLOC_USE_CXX11_THREAD_ATEXIT_DIRECT
+#  undef SNMALLOC_USE_CXX11_THREAD_ATEXIT_DIRECT
+#endif
+#ifdef SNMALLOC_USE_CXX11_DESTRUCTORS
+#  undef SNMALLOC_USE_CXX11_DESTRUCTORS
+#endif
 
 #include <new>
 #include <snmalloc/snmalloc_core.h>