Add TLS-based reentrancy check in an off state.

tcmalloc/internal/logging.cc

@@ -17,6 +17,8 @@
 #include <fcntl.h>
 
 #include <iterator>
+
+#include "absl/base/thread_annotations.h"
 #ifdef _WINDOWS
 #include <io.h>
 #endif
@@ -55,7 +57,7 @@ namespace tcmalloc_internal {
 // may not be able to heap-allocate while crashing.
 ABSL_CONST_INIT static absl::base_internal::SpinLock crash_lock(
     absl::base_internal::SCHEDULE_KERNEL_ONLY);
-static bool crashed = false;
+static bool crashed ABSL_GUARDED_BY(crash_lock) = false;
 
 static const size_t kStatsBufferSize = 16 << 10;
 #ifndef __APPLE__
@@ -167,6 +169,14 @@ void RecordCrash(absl::string_view detector, absl::string_view error) {
     if (!crashed) {
       crashed = true;
       first_crash = true;
+#if !defined(NDEBUG)
+      // Reset reentrancy count to allow calls to malloc/free under any signal
+      // handlers that we may trigger from abort().
+      //
+      // We only reset this, rather than permit an unlimited amount of
+      // reentrancy while crashing, in order to maintain modest stack depths.
+      tcmalloc_reentrancy_count = 0;
+#endif
     }
   }
 
@@ -287,6 +297,10 @@ PbtxtRegion PbtxtRegion::CreateSubRegion(absl::string_view key) {
   return sub;
 }
 
+#if !defined(NDEBUG)
+ABSL_CONST_INIT thread_local int tcmalloc_reentrancy_count = 0;
+#endif
+
 }  // namespace tcmalloc_internal
 }  // namespace tcmalloc
 GOOGLE_MALLOC_SECTION_END

tcmalloc/internal/logging.h

@@ -358,7 +358,26 @@ class PbtxtRegion {
   PbtxtRegionType type_;
 };
 
+#if !defined(NDEBUG)
+ABSL_CONST_INIT extern thread_local int tcmalloc_reentrancy_count;
+#endif
 
+#if 0 && !defined(NDEBUG)
+class [[maybe_unused]] ReentrancyGuard {
+ public:
+  ReentrancyGuard() {
+    TC_CHECK_EQ(tcmalloc_reentrancy_count, 0);
+    ++tcmalloc_reentrancy_count;
+  }
+  ~ReentrancyGuard() { --tcmalloc_reentrancy_count; }
+};
+#else
+class [[maybe_unused]] ReentrancyGuard {
+ public:
+  constexpr ReentrancyGuard() = default;
+  ~ReentrancyGuard() = default;
+};
+#endif
 
 }  // namespace tcmalloc_internal
 }  // namespace tcmalloc

tcmalloc/tcmalloc.cc

@@ -860,6 +860,8 @@ ABSL_ATTRIBUTE_NOINLINE static void do_unsized_free_irregular(
 
 template <typename Policy>
 inline ABSL_ATTRIBUTE_ALWAYS_INLINE void do_free(void* ptr, Policy policy) {
+  ReentrancyGuard guard;
+
   // TODO(b/404341539):  Improve the bound.
   TC_ASSERT(CorrectAlignment(ptr, static_cast<std::align_val_t>(1)));
 
@@ -972,6 +974,8 @@ template <typename Policy>
 inline ABSL_ATTRIBUTE_ALWAYS_INLINE void do_free_with_size(void* ptr,
                                                            size_t size,
                                                            Policy policy) {
+  ReentrancyGuard guard;
+
   TC_ASSERT(CorrectAlignment(ptr, policy.align()));
 
   // This is an optimized path that may be taken if the binary is compiled
@@ -1364,6 +1368,8 @@ ABSL_ATTRIBUTE_NOINLINE static typename Policy::pointer_type slow_alloc_large(
 template <typename Policy, typename Pointer = typename Policy::pointer_type>
 static inline Pointer ABSL_ATTRIBUTE_ALWAYS_INLINE fast_alloc(size_t size,
                                                               Policy policy) {
+  ReentrancyGuard guard;
+
   // If size is larger than kMaxSize, it's not fast-path anymore. In
   // such case, GetSizeClass will return false, and we'll delegate to the slow
   // path. If malloc is not yet initialized, we may end up with size_class == 0
@@ -1373,7 +1379,10 @@ static inline Pointer ABSL_ATTRIBUTE_ALWAYS_INLINE fast_alloc(size_t size,
       tc_globals.sizemap().GetSizeClass(policy, size);
   if (ABSL_PREDICT_FALSE(!is_small)) {
     SLOW_PATH_BARRIER();
-    TCMALLOC_MUSTTAIL return slow_alloc_large(size, policy);
+#if defined(NDEBUG) && !defined(TCMALLOC_INTERNAL_WITH_ASSERTIONS)
+    TCMALLOC_MUSTTAIL
+#endif
+    return slow_alloc_large(size, policy);
   }
 
   // TryRecordAllocationFast() returns true if no extra logic is required, e.g.:
@@ -1716,6 +1725,12 @@ static inline ABSL_ATTRIBUTE_ALWAYS_INLINE void* do_realloc(void* old_ptr,
     do_free(old_ptr, MallocPolicy());
     return new_ptr;
   } else {
+    // We guard here since callers to
+    // `TCMallocInternalRealloc`/`TCMallocInternalReallocArray` and the above
+    // conditional will check for reentrancy inside of `fast_alloc` and
+    // `do_free`.
+    tcmalloc::tcmalloc_internal::ReentrancyGuard guard;
+
     // We still need to call hooks to report the updated size:
     tcmalloc::MallocHook::InvokeDeleteHook(
         {const_cast<void*>(old_ptr), std::nullopt, old_size,

tcmalloc/testing/BUILD

@@ -1392,3 +1392,16 @@ create_tcmalloc_testsuite(
         "@com_google_googletest//:gtest_main",
     ],
 )
+
+create_tcmalloc_testsuite(
+    name = "reentrancy_test",
+    srcs = ["reentrancy_test.cc"],
+    deps = [
+        "//tcmalloc:malloc_extension",
+        "//tcmalloc:malloc_hook",
+        "//tcmalloc/internal:config",
+        "//tcmalloc/internal:logging",
+        "@com_google_absl//absl/base:core_headers",
+        "@com_google_googletest//:gtest_main",
+    ],
+)

tcmalloc/testing/CMakeLists.txt

@@ -1311,3 +1311,19 @@ tcmalloc_cc_test_variants(
     "tcmalloc::testing_testutil"
 )
 
+tcmalloc_cc_test_variants(
+  NAME
+    tcmalloc_testing_reentrancy_test
+  SRCS
+    "reentrancy_test.cc"
+  DEPS
+    "GTest::gtest_main"
+    "GTest::gmock_main"
+    "GTest::gmock"
+    "absl::core_headers"
+    "tcmalloc::internal_config"
+    "tcmalloc::internal_logging"
+    "tcmalloc::malloc_extension"
+    "tcmalloc::malloc_hook"
+)
+