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/snmalloc/global/threadalloc.h

@@ -21,8 +21,25 @@
 #  endif
 #endif
 
+#if defined(SNMALLOC_USE_CXX11_THREAD_ATEXIT_DIRECT)
+#  if defined(SNMALLOC_THREAD_TEARDOWN_DEFINED)
+#    error At most one out of method of thread teardown can be specified.
+#  endif
+#  define SNMALLOC_THREAD_TEARDOWN_DEFINED
+extern "C" int __cxa_thread_atexit_impl(void(func)(void*), void*, void*);
+extern "C" void* __dso_handle;
+#endif
+
+#if defined(SNMALLOC_USE_CXX11_DESTRUCTORS)
+#  if defined(SNMALLOC_THREAD_TEARDOWN_DEFINED)
+#    error At most one out of method of thread teardown can be specified.
+#  endif
+#  define SNMALLOC_THREAD_TEARDOWN_DEFINED
+#endif
+
 #if !defined(SNMALLOC_THREAD_TEARDOWN_DEFINED)
-#  define SNMALLOC_USE_CXX_THREAD_DESTRUCTORS
+// Default to C++11 destructors if nothing has been specified.
+#  define SNMALLOC_USE_CXX11_DESTRUCTORS
 #endif
 extern "C" void _malloc_thread_cleanup();
 
@@ -56,6 +73,7 @@ namespace snmalloc
   class CheckInitPthread;
   class CheckInitCXX;
   class CheckInitThreadCleanup;
+  class CheckInitCXXAtExitDirect;
 
   /**
    * Holds the thread local state for the allocator.  The state is constant
@@ -159,8 +177,10 @@ namespace snmalloc
     };
 #  ifdef SNMALLOC_USE_PTHREAD_DESTRUCTORS
     using CheckInit = CheckInitPthread;
-#  elif defined(SNMALLOC_USE_CXX_THREAD_DESTRUCTORS)
+#  elif defined(SNMALLOC_USE_CXX11_DESTRUCTORS)
     using CheckInit = CheckInitCXX;
+#  elif defined(SNMALLOC_USE_CXX11_THREAD_ATEXIT_DIRECT)
+    using CheckInit = CheckInitCXXAtExitDirect;
 #  else
     using CheckInit = CheckInitThreadCleanup;
 #  endif
@@ -180,8 +200,15 @@ namespace snmalloc
 
     /**
      * Used to give correct signature to teardown required by atexit.
+     * If [[gnu::destructor]] is available, we use the attribute to register
+     * the finalisation statically. In VERY rare cases, dynamic registration
+     * can trigger deadlocks.
      */
-    static void pthread_cleanup_main_thread()
+#    if __has_attribute(destructor)
+    [[gnu::destructor]]
+#    endif
+    static void
+    pthread_cleanup_main_thread()
     {
       ThreadAlloc::teardown();
     }
@@ -198,7 +225,9 @@ namespace snmalloc
       // run at least once.  If the main thread exits with `pthread_exit` then
       // it will be called twice but this case is already handled because other
       // destructors can cause the per-thread allocator to be recreated.
+#    if !__has_attribute(destructor)
       atexit(&pthread_cleanup_main_thread);
+#    endif
     }
 
   public:
@@ -220,7 +249,31 @@ namespace snmalloc
 #    endif
     }
   };
-#  elif defined(SNMALLOC_USE_CXX_THREAD_DESTRUCTORS)
+#  elif defined(SNMALLOC_USE_CXX11_THREAD_ATEXIT_DIRECT)
+  class CheckInitCXXAtExitDirect
+  : public ThreadAlloc::CheckInitBase<CheckInitCXXAtExitDirect>
+  {
+    static void cleanup(void*)
+    {
+      ThreadAlloc::teardown();
+    }
+
+  public:
+    /**
+     * This function is called by each thread once it starts using the
+     * thread local allocator.
+     *
+     * This implementation depends on the libstdc++ requirement on libc
+     * that provides the functionality to register a destructor for the
+     * thread locals.  This allows to not require either pthread or
+     * C++ std lib.
+     */
+    static void register_clean_up()
+    {
+      __cxa_thread_atexit_impl(cleanup, nullptr, &__dso_handle);
+    }
+  };
+#  elif defined(SNMALLOC_USE_CXX11_DESTRUCTORS)
   class CheckInitCXX : public ThreadAlloc::CheckInitBase<CheckInitCXX>
   {
   public:

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>