[libc] add _malloc_thread_cleanup option

Author: SchrodingerZhu

libc/config/config.json

@@ -89,6 +89,10 @@
     "LIBC_CONF_RWLOCK_DEFAULT_SPIN_COUNT": {
       "value": 100,
       "doc": "Default number of spins before blocking if a rwlock is in contention (default to 100)."
+    },
+    "LIBC_CONF_ENABLE_MALLOC_THREAD_CLEANUP": {
+      "value": false,
+      "doc": "Enable the `_malloc_thread_cleanup` weak symbol. When defined, this is function is called after `__cxa` and pthread-specific dtors. On main thread, this will be called after `atexit` functions and `.fini` dtors, right before TLS tearing down. This function can be overridden by allocators to perform cleanup. Allocators can use this symbol to avoid registering thread dtors using potentially reentrant routines."
     }
   },
   "math": {

libc/docs/configure.rst

@@ -47,6 +47,7 @@ to learn about the defaults for your platform and target.
     - ``LIBC_CONF_PRINTF_FLOAT_TO_STR_USE_MEGA_LONG_DOUBLE_TABLE``: Use large table for better printf long double performance.
     - ``LIBC_CONF_PRINTF_RUNTIME_DISPATCH``: Use dynamic dispatch for the output mechanism to reduce code size.
 * **"pthread" options**
+    - ``LIBC_CONF_ENABLE_MALLOC_THREAD_CLEANUP``: Enable the `_malloc_thread_cleanup` weak symbol. When defined, this is function is called after `__cxa` and pthread-specific dtors. On main thread, this will be called after `atexit` functions and `.fini` dtors, right before TLS tearing down. This function can be overridden by allocators to perform cleanup. Allocators can use this symbol to avoid registering thread dtors using potentially reentrant routines.
     - ``LIBC_CONF_RAW_MUTEX_DEFAULT_SPIN_COUNT``: Default number of spins before blocking if a mutex is in contention (default to 100).
     - ``LIBC_CONF_RWLOCK_DEFAULT_SPIN_COUNT``: Default number of spins before blocking if a rwlock is in contention (default to 100).
     - ``LIBC_CONF_TIMEOUT_ENSURE_MONOTONICITY``: Automatically adjust timeout to CLOCK_MONOTONIC (default to true). POSIX API may require CLOCK_REALTIME, which can be unstable and leading to unexpected behavior. This option will convert the real-time timestamp to monotonic timestamp relative to the time of call.

libc/src/__support/threads/linux/CMakeLists.txt

@@ -71,6 +71,12 @@ add_header_library(
     libc.src.__support.threads.mutex_common
 )
 
+if (LIBC_CONF_ENABLE_MALLOC_THREAD_CLEANUP)
+  set(malloc_cleanup_flags -DLIBC_COPT_ENABLE_MALLOC_THREAD_CLEANUP)
+else()
+  set(malloc_cleanup_flags)
+endif()
+
 add_object_library(
   thread
   SRCS
@@ -89,6 +95,7 @@ add_object_library(
     libc.src.__support.threads.thread_common
   COMPILE_OPTIONS
     ${libc_opt_high_flag}
+    ${malloc_cleanup_flags}
     -fno-omit-frame-pointer # This allows us to sniff out the thread args from
                             # the new thread's stack reliably.
     -Wno-frame-address      # Yes, calling __builtin_return_address with a

libc/src/__support/threads/linux/thread.cpp

@@ -482,6 +482,10 @@ int Thread::get_name(cpp::StringStream &name) const {
   return 0;
 }
 
+#ifdef LIBC_COPT_ENABLE_MALLOC_THREAD_CLEANUP
+extern "C" [[gnu::weak]] void _malloc_thread_cleanup();
+#endif // LIBC_COPT_ENABLE_MALLOC_THREAD_CLEANUP
+
 void thread_exit(ThreadReturnValue retval, ThreadStyle style) {
   auto attrib = self.attrib;
 
@@ -494,6 +498,11 @@ void thread_exit(ThreadReturnValue retval, ThreadStyle style) {
   // different thread. The destructors of thread local and TSS objects should
   // be called by the thread which owns them.
   internal::call_atexit_callbacks(attrib);
+#ifdef LIBC_COPT_ENABLE_MALLOC_THREAD_CLEANUP
+  // call _malloc_thread_cleanup after the atexit callbacks
+  if (_malloc_thread_cleanup)
+    _malloc_thread_cleanup();
+#endif // LIBC_COPT_ENABLE_MALLOC_THREAD_CLEANUP
 
   uint32_t joinable_state = uint32_t(DetachState::JOINABLE);
   if (!attrib->detach_state.compare_exchange_strong(

libc/src/stdlib/CMakeLists.txt

@@ -589,6 +589,12 @@ add_header_library(
 )
 endif()
 
+if (LIBC_CONF_ENABLE_MALLOC_THREAD_CLEANUP)
+  set(malloc_cleanup_flags -DLIBC_COPT_ENABLE_MALLOC_THREAD_CLEANUP)
+else()
+  set(malloc_cleanup_flags)
+endif()
+
 add_entrypoint_object(
   atexit
   SRCS
@@ -599,6 +605,8 @@ add_entrypoint_object(
     20 # For constinit
   DEPENDS
     .exit_handler
+  COMPILE_OPTIONS
+    ${malloc_cleanup_flags}
 )
 
 add_entrypoint_object(

libc/src/stdlib/atexit.cpp

@@ -24,9 +24,17 @@ int __cxa_atexit(AtExitCallback *callback, void *payload, void *) {
   return add_atexit_unit(atexit_callbacks, {callback, payload});
 }
 
+#ifdef LIBC_COPT_ENABLE_MALLOC_THREAD_CLEANUP
+extern "C" [[gnu::weak]] void _malloc_thread_cleanup();
+#endif // LIBC_COPT_ENABLE_MALLOC_THREAD_CLEANUP
+
 void __cxa_finalize(void *dso) {
   if (!dso) {
     call_exit_callbacks(atexit_callbacks);
+#ifdef LIBC_COPT_ENABLE_MALLOC_THREAD_CLEANUP
+    if (_malloc_thread_cleanup)
+      _malloc_thread_cleanup();
+#endif // LIBC_COPT_ENABLE_MALLOC_THREAD_CLEANUP
     if (teardown_main_tls)
       teardown_main_tls();
   }

libc/test/integration/src/__support/threads/CMakeLists.txt

@@ -36,6 +36,21 @@ add_integration_test(
     libc.src.__support.threads.thread
 )
 
+if (LIBC_CONF_ENABLE_MALLOC_THREAD_CLEANUP)
+  add_integration_test(
+    malloc_cleanup_test
+    SUITE
+      libc-support-threads-integration-tests
+    SRCS
+      malloc_cleanup_test.cpp
+    DEPENDS
+      libc.src.__support.threads.thread
+      libc.src.stdlib.atexit
+      libc.src.stdlib.exit
+      libc.src.stdlib._Exit
+  )
+endif()
+
 add_integration_test(
   double_exit_test
   SUITE

libc/test/integration/src/__support/threads/malloc_cleanup_test.cpp

@@ -0,0 +1,47 @@
+//===-- Test handling of malloc TLS data ----------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "src/__support/threads/thread.h"
+#include "src/stdlib/atexit.h"
+#include "src/stdlib/exit.h"
+#include "src/stdlib/_Exit.h"
+#include "test/IntegrationTest/test.h"
+
+bool cxa_dtor_called = false;
+bool fini_dtor_called = false;
+bool atexit_called = false;
+volatile thread_local int tls_accessible = 0;
+
+extern "C" {
+[[gnu::weak]]
+void *__dso_handle = nullptr;
+int __cxa_thread_atexit_impl(void (*func)(void *), void *arg, void *dso);
+void _malloc_thread_cleanup() {
+  // make sure that TLS is still alive
+  tls_accessible = 1;
+  if (!cxa_dtor_called)
+    __builtin_trap();
+  if (!fini_dtor_called)
+    __builtin_trap();
+  if (!atexit_called)
+    __builtin_trap();
+  LIBC_NAMESPACE::_Exit(0);
+}
+}
+
+[[gnu::destructor]]
+void destructor() {
+  fini_dtor_called = true;
+}
+
+TEST_MAIN() {
+  __cxa_thread_atexit_impl([](void *) { cxa_dtor_called = true; }, nullptr,
+                           __dso_handle);
+  LIBC_NAMESPACE::atexit([]() { atexit_called = true; });
+  LIBC_NAMESPACE::exit(1);
+}