diff --git a/tcmalloc/page_allocator.cc b/tcmalloc/page_allocator.cc index a608ddc5f..5816cedfa 100644 --- a/tcmalloc/page_allocator.cc +++ b/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,6 +227,7 @@ 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(span_alloc_info.density), tag); @@ -232,6 +235,7 @@ ABSL_ATTRIBUTE_NOINLINE void PageAllocator::InvokeDeleteHookSlow( 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); } diff --git a/tcmalloc/testing/BUILD b/tcmalloc/testing/BUILD index ca40314af..5b971d7ef 100644 --- a/tcmalloc/testing/BUILD +++ b/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", diff --git a/tcmalloc/testing/CMakeLists.txt b/tcmalloc/testing/CMakeLists.txt index 1026d234b..ceb27fcd0 100644 --- a/tcmalloc/testing/CMakeLists.txt +++ b/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" diff --git a/tcmalloc/testing/hooks_test.cc b/tcmalloc/testing/hooks_test.cc index b899c0049..277b3a96a 100644 --- a/tcmalloc/testing/hooks_test.cc +++ b/tcmalloc/testing/hooks_test.cc @@ -21,6 +21,8 @@ #include #include +#include +#include #include #include #include @@ -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