Add AllocationGuard to page allocator hooks.

tcmalloc/page_allocator.cc

@@ -24,6 +24,7 @@
 #include "absl/base/optimization.h"
 #include "tcmalloc/common.h"
 #include "tcmalloc/huge_page_aware_allocator.h"
+#include "tcmalloc/internal/allocation_guard.h"
 #include "tcmalloc/internal/config.h"
 #include "tcmalloc/internal/environment.h"
 #include "tcmalloc/internal/logging.h"
@@ -216,6 +217,7 @@ size_t PageAllocator::active_partitions() const {
 ABSL_ATTRIBUTE_NOINLINE void PageAllocator::InvokeNewHookSlow(
     Span* span, Length n, Length align, SpanAllocInfo span_alloc_info,
     MemoryTag tag) {
+  AllocationGuard g;
   size_t start_page_index = span ? span->first_page().index() : 0;
   page_allocator_new_hooks.Invoke(
       start_page_index, n.raw_num(), align.raw_num(),
@@ -225,13 +227,15 @@ ABSL_ATTRIBUTE_NOINLINE void PageAllocator::InvokeNewHookSlow(
 
 ABSL_ATTRIBUTE_NOINLINE void PageAllocator::InvokeDeleteHookSlow(
     PageId start_page, Length n, SpanAllocInfo span_alloc_info, MemoryTag tag) {
+  AllocationGuard g;
   page_allocator_delete_hooks.Invoke(
       start_page.index(), n.raw_num(), span_alloc_info.objects_per_span,
       static_cast<uint8_t>(span_alloc_info.density), tag);
 }
 
 ABSL_ATTRIBUTE_NOINLINE void PageAllocator::InvokeReleaseHookSlow(
     Length num_pages, Length released, PageReleaseReason reason) {
+  AllocationGuard g;
   page_allocator_release_hooks.Invoke(num_pages.raw_num(), released.raw_num(),
                                       reason);
 }

tcmalloc/testing/BUILD

@@ -184,6 +184,7 @@ create_tcmalloc_testsuite(
         "//tcmalloc/internal:allocation_guard",
         "//tcmalloc/internal:config",
         "//tcmalloc/internal:logging",
+        "//tcmalloc/internal:page_allocator_hooks",
         "@com_github_google_benchmark//:benchmark",
         "@com_google_absl//absl/base:malloc_internal",
         "@com_google_absl//absl/debugging:leak_check",

tcmalloc/testing/CMakeLists.txt

@@ -154,6 +154,7 @@ tcmalloc_cc_test_variants(
     "tcmalloc::internal_allocation_guard"
     "tcmalloc::internal_config"
     "tcmalloc::internal_logging"
+    "tcmalloc::internal_page_allocator_hooks"
     "tcmalloc::malloc_extension"
     "tcmalloc::malloc_hook"
     "tcmalloc::testing_malloc_hook_recorder"

tcmalloc/testing/hooks_test.cc

@@ -21,6 +21,8 @@
 #include <sys/mman.h>
 #include <unistd.h>
 
+#include <cstddef>
+#include <cstdint>
 #include <new>
 #include <optional>
 #include <string>
@@ -33,6 +35,7 @@
 #include "tcmalloc/internal/allocation_guard.h"
 #include "tcmalloc/internal/config.h"
 #include "tcmalloc/internal/logging.h"
+#include "tcmalloc/internal/page_allocator_hooks.h"
 #include "tcmalloc/malloc_extension.h"
 #include "tcmalloc/malloc_hook.h"
 #include "tcmalloc/malloc_hook_invoke.h"
@@ -43,6 +46,8 @@ namespace tcmalloc {
 namespace {
 
 using tcmalloc_internal::kSanitizerPresent;
+using tcmalloc_internal::page_allocator_delete_hooks;
+using tcmalloc_internal::page_allocator_new_hooks;
 using tcmalloc_testing::MallocHookRecorder;
 using testing::Contains;
 using testing::ElementsAre;
@@ -598,5 +603,37 @@ TEST(TCMallocTest, MarkThreadBusy) {
   log_hooks = false;
 }
 
+TEST(HooksTest, AllocationInHookFails) {
+  if (kSanitizerPresent) {
+    GTEST_SKIP() << "Sanitizers intercept malloc/new";
+  }
+
+  const auto new_hook = [](size_t, size_t, size_t, size_t, uint8_t,
+                           tcmalloc_internal::MemoryTag) {
+    ::operator delete(::operator new(1));
+  };
+  EXPECT_DEBUG_DEATH(
+      {
+        EXPECT_TRUE(page_allocator_new_hooks.Add(new_hook));
+        void* ptr = ::operator new(1024 * 1024 * 16);
+        ::operator delete(ptr);
+        EXPECT_TRUE(page_allocator_new_hooks.Remove(new_hook));
+      },
+      "");
+
+  const auto delete_hook = [](size_t, size_t, size_t, uint8_t,
+                              tcmalloc_internal::MemoryTag) {
+    ::operator delete(::operator new(1));
+  };
+  EXPECT_DEBUG_DEATH(
+      {
+        void* ptr = ::operator new(1024 * 1024 * 16);
+        EXPECT_TRUE(page_allocator_delete_hooks.Add(delete_hook));
+        ::operator delete(ptr);
+        EXPECT_TRUE(page_allocator_delete_hooks.Remove(delete_hook));
+      },
+      "");
+}
+
 }  // namespace
 }  // namespace tcmalloc