From d0dde7ce39dcefbb8d64375955364ac20bb05884 Mon Sep 17 00:00:00 2001 From: Ben Reisner Date: Thu, 10 Sep 2026 09:22:29 -0700 Subject: [PATCH] Consolidate SAMPLED and COLD page heap partitions. Consolidates the SAMPLED and COLD page heap partitions into a single COLD partition, reducing kNumHeaps from 3 to 2. Both partitions use MADV_NOHUGEPAGE. - PageAllocator: Route sampled memory tags to cold_impl_, allocate cold_impl_ unconditionally, and remove sampled_impl_ and has_cold_impl_. - tcmalloc.cc: Update do_free_with_size and handle_sampled_or_illformed_ptrs to properly route and handle sampled allocations in cold_impl_. - central_freelist.cc: Assert !free_span->sampled() on deallocation. - Update tests to account for sampled allocations carrying MemoryTag::kCold. PiperOrigin-RevId: 979226826 --- tcmalloc/central_freelist.cc | 2 +- tcmalloc/page_allocator.cc | 43 ++---------- tcmalloc/page_allocator.h | 76 ++++------------------ tcmalloc/tcmalloc.cc | 16 +++-- tcmalloc/testing/heap_profiling_test.cc | 4 +- tcmalloc/testing/partitioning_fuzz_test.cc | 1 + tcmalloc/testing/tcmalloc_test.cc | 11 +++- tcmalloc/testing/want_hpaa_test_helper.cc | 3 +- 8 files changed, 46 insertions(+), 110 deletions(-) diff --git a/tcmalloc/central_freelist.cc b/tcmalloc/central_freelist.cc index c47c2bb8b..94cdf4762 100644 --- a/tcmalloc/central_freelist.cc +++ b/tcmalloc/central_freelist.cc @@ -160,7 +160,7 @@ void StaticForwarder::DeallocateSpans(size_t objects_per_span, // Unregister size class doesn't require holding any locks. for (Span* const free_span : free_spans) { TC_ASSERT_EQ(GetMemoryTag(free_span->start_address()), tag); - TC_ASSERT(!IsSampledMemory(free_span->start_address())); + TC_ASSERT(!free_span->sampled()); tc_globals.pagemap().UnregisterSizeClass(free_span); // Before taking pageheap_lock, prefetch the PageTrackers these spans are diff --git a/tcmalloc/page_allocator.cc b/tcmalloc/page_allocator.cc index d5d61b5e3..71aebc670 100644 --- a/tcmalloc/page_allocator.cc +++ b/tcmalloc/page_allocator.cc @@ -43,9 +43,6 @@ namespace tcmalloc_internal { using huge_page_allocator_internal::HugePageAwareAllocatorOptions; PageAllocator::PageAllocator() { - has_cold_impl_ = ColdFeatureActive(); - sampled_partition_active_ = - Parameters::heap_partitioning_mode() == HeapPartitioningMode::kFull; size_t part = 0; normal_impl_[0] = new (&choices_[part++].hpaa) @@ -56,22 +53,8 @@ PageAllocator::PageAllocator() { HugePageAwareAllocator( HugePageAwareAllocatorOptions{MemoryTag::kNormalP1}); } - sampled_impl_[0] = new (&choices_[part++].hpaa) HugePageAwareAllocator( - HugePageAwareAllocatorOptions{MemoryTag::kSampled}); - if (sampled_partition_active_) { - // this is not the case for NUMA partitions, hence, we can't use the - // active_partitions() check. - sampled_impl_[1] = - new (tc_globals.arena().Alloc(sizeof(HugePageAwareAllocator))) - HugePageAwareAllocator( - HugePageAwareAllocatorOptions{MemoryTag::kSampledP1}); - } - if (has_cold_impl_) { - cold_impl_ = new (&choices_[part++].hpaa) - HugePageAwareAllocator(HugePageAwareAllocatorOptions{MemoryTag::kCold}); - } else { - cold_impl_ = normal_impl_[0]; - } + cold_impl_ = new (&choices_[part++].hpaa) + HugePageAwareAllocator(HugePageAwareAllocatorOptions{MemoryTag::kCold}); alg_ = HPAA; TC_CHECK_LE(part, std::size(choices_)); } @@ -183,13 +166,11 @@ bool PageAllocator::ShrinkHardBy(Length pages, LimitKind limit_kind) { limit); warned_hugepages = true; } - if (has_cold_impl_) { - ret += static_cast(cold_impl_) - ->ReleaseAtLeastNPagesBreakingHugepages(pages - ret, - release_reason); - if (ret >= pages) { - return true; - } + ret += static_cast(cold_impl_) + ->ReleaseAtLeastNPagesBreakingHugepages(pages - ret, + release_reason); + if (ret >= pages) { + return true; } for (int partition = 0; partition < active_partitions(); partition++) { ret += static_cast(normal_impl_[partition]) @@ -199,16 +180,6 @@ bool PageAllocator::ShrinkHardBy(Length pages, LimitKind limit_kind) { return true; } } - for (int partition = 0; - partition < (sampled_partition_active_ ? kSecurityPartitions : 1); - partition++) { - ret += static_cast(sampled_impl_[partition]) - ->ReleaseAtLeastNPagesBreakingHugepages(pages - ret, - release_reason); - if (ret >= pages) { - return true; - } - } } // Return "true", if we got back under the limit. return (pages <= ret); diff --git a/tcmalloc/page_allocator.h b/tcmalloc/page_allocator.h index 11e74f969..f376dda8a 100644 --- a/tcmalloc/page_allocator.h +++ b/tcmalloc/page_allocator.h @@ -218,7 +218,8 @@ class PageAllocator { size_t active_partitions() const; - static constexpr size_t kNumHeaps = 3; // 3 heaps: normal, sampled, cold. + static constexpr size_t kNumHeaps = + 2; // 2 heaps: normal, cold (including sampled). union Choices { Choices() : dummy(0) {} @@ -227,11 +228,8 @@ class PageAllocator { HugePageAwareAllocator hpaa; } choices_[kNumHeaps]; std::array normal_impl_; - std::array sampled_impl_; Interface* cold_impl_; Algorithm alg_; - bool has_cold_impl_; - bool sampled_partition_active_; bool over_limit_ ABSL_GUARDED_BY(pageheap_lock) = false; // Max size of backed spans we will attempt to maintain. @@ -265,9 +263,7 @@ inline PageAllocator::Interface* PageAllocator::impl(MemoryTag tag) const { case MemoryTag::kNormalP1: return normal_impl_[1]; case MemoryTag::kSampled: - return sampled_impl_[0]; case MemoryTag::kSampledP1: - return sampled_impl_[1]; case MemoryTag::kCold: return cold_impl_; default: @@ -317,63 +313,35 @@ inline BackingStats PageAllocator::stats() const { for (int partition = 1; partition < active_partitions(); partition++) { ret += normal_impl_[partition]->stats(); } - ret += sampled_impl_[0]->stats(); - if (sampled_partition_active_) { - ret += sampled_impl_[1]->stats(); - } - if (has_cold_impl_) { - ret += cold_impl_->stats(); - } + ret += cold_impl_->stats(); return ret; } inline void PageAllocator::GetSmallSpanStats(SmallSpanStats* result) { - SmallSpanStats normal, sampled; + SmallSpanStats normal, cold; for (int partition = 0; partition < active_partitions(); partition++) { SmallSpanStats part_stats; normal_impl_[partition]->GetSmallSpanStats(&part_stats); normal += part_stats; } - sampled_impl_[0]->GetSmallSpanStats(&sampled); - if (sampled_partition_active_) { - SmallSpanStats part_stats; - sampled_impl_[1]->GetSmallSpanStats(&part_stats); - sampled += part_stats; - } - *result = normal + sampled; - if (has_cold_impl_) { - SmallSpanStats cold; - cold_impl_->GetSmallSpanStats(&cold); - *result += cold; - } + cold_impl_->GetSmallSpanStats(&cold); + *result = normal + cold; } inline void PageAllocator::GetLargeSpanStats(LargeSpanStats* result) { - LargeSpanStats normal, sampled; + LargeSpanStats normal, cold; for (int partition = 0; partition < active_partitions(); partition++) { LargeSpanStats part_stats; normal_impl_[partition]->GetLargeSpanStats(&part_stats); normal += part_stats; } - sampled_impl_[0]->GetLargeSpanStats(&sampled); - if (sampled_partition_active_) { - LargeSpanStats part_stats; - sampled_impl_[1]->GetLargeSpanStats(&part_stats); - sampled += part_stats; - } - *result = normal + sampled; - if (has_cold_impl_) { - LargeSpanStats cold; - cold_impl_->GetLargeSpanStats(&cold); - *result = *result + cold; - } + cold_impl_->GetLargeSpanStats(&cold); + *result = normal + cold; } inline void PageAllocator::TreatHugepageTrackers( EnableCollapse enable_collapse) { - if (has_cold_impl_) { - cold_impl_->TreatHugepageTrackers(EnableCollapse::kDisabled); - } + cold_impl_->TreatHugepageTrackers(EnableCollapse::kDisabled); for (int partition = 0; partition < active_partitions(); partition++) { normal_impl_[partition]->TreatHugepageTrackers(enable_collapse); } @@ -384,21 +352,12 @@ inline Length PageAllocator::ReleaseAtLeastNPages(Length num_pages, Length released; // TODO(ckennelly): Refine this policy. Cold data should be the most // resilient to not being on huge pages. - if (has_cold_impl_) { - released = cold_impl_->ReleaseAtLeastNPages(num_pages, reason); - } + released = cold_impl_->ReleaseAtLeastNPages(num_pages, reason); for (int partition = 0; partition < active_partitions(); partition++) { released += normal_impl_[partition]->ReleaseAtLeastNPages( num_pages > released ? num_pages - released : Length(0), reason); } - released += sampled_impl_[0]->ReleaseAtLeastNPages( - num_pages > released ? num_pages - released : Length(0), reason); - if (sampled_partition_active_) { - released += sampled_impl_[1]->ReleaseAtLeastNPages( - num_pages > released ? num_pages - released : Length(0), reason); - } - InvokeReleaseHook(num_pages, released, reason); return released; } @@ -406,24 +365,17 @@ inline Length PageAllocator::ReleaseAtLeastNPages(Length num_pages, inline PageReleaseStats PageAllocator::GetReleaseStats() const { PageReleaseStats stats; - if (has_cold_impl_) { - stats += cold_impl_->GetReleaseStats(); - } + stats += cold_impl_->GetReleaseStats(); for (int partition = 0; partition < active_partitions(); partition++) { stats += normal_impl_[partition]->GetReleaseStats(); } - stats += sampled_impl_[0]->GetReleaseStats(); - if (sampled_partition_active_) { - stats += sampled_impl_[1]->GetReleaseStats(); - } - return stats; } inline void PageAllocator::Print(Printer& out, MemoryTag tag, PageFlagsBase& pageflags) { - if (tag == MemoryTag::kCold && !has_cold_impl_) { + if (tag == MemoryTag::kSampled || tag == MemoryTag::kSampledP1) { return; } @@ -439,7 +391,7 @@ inline void PageAllocator::Print(Printer& out, MemoryTag tag, inline void PageAllocator::PrintInPbtxt(PbtxtRegion& region, MemoryTag tag, PageFlagsBase& pageflags) { - if (tag == MemoryTag::kCold && !has_cold_impl_) { + if (tag == MemoryTag::kSampled || tag == MemoryTag::kSampledP1) { return; } diff --git a/tcmalloc/tcmalloc.cc b/tcmalloc/tcmalloc.cc index e4a38a796..b846bc352 100644 --- a/tcmalloc/tcmalloc.cc +++ b/tcmalloc/tcmalloc.cc @@ -634,7 +634,7 @@ inline sized_ptr_t do_malloc_pages(size_t size, size_t weight, Policy policy) { Length num_pages = std::max(BytesToLengthCeil(size), Length(1)); MemoryTag tag = MemoryTag::kNormal; - if (policy.is_cold() && + if (ColdFeatureActive() && policy.is_cold() && (Parameters::heap_partitioning_mode() != HeapPartitioningMode::kFull || policy.security_partition() == 0)) { tag = MemoryTag::kCold; @@ -849,10 +849,11 @@ ABSL_ATTRIBUTE_NOINLINE static void handle_sampled_or_illformed_ptrs( auto tag = GetMemoryTag(ptr); const uintptr_t uptr = absl::bit_cast(ptr); TC_ASSERT((uptr & (kBadAlignmentMask | kBadDeallocationHighMask)) != 0 || - (tag != MemoryTag::kNormal && tag != MemoryTag::kNormalP1 && - tag != MemoryTag::kCold)); + (tag != MemoryTag::kNormal && tag != MemoryTag::kNormalP1)); - if (ABSL_PREDICT_TRUE(IsSampledMemory(ptr))) { + if (ABSL_PREDICT_TRUE(IsSampledMemory(ptr) || + tc_globals.pagemap().sizeclass(PageIdContaining(ptr)) == + 0)) { // we don't know true class size of the ptr return InvokeHooksAndFreePages(ptr, size, policy); } @@ -911,7 +912,8 @@ inline ABSL_ATTRIBUTE_ALWAYS_INLINE void do_free_with_size(void* ptr, if (ABSL_PREDICT_FALSE(ptr == nullptr)) { return; } - bool is_cold = ((uptr & kTagOrBadDeallocationMask) == kColdMask); + bool is_cold = ((uptr & kTagOrBadDeallocationMask) == kColdMask) && + (tc_globals.pagemap().sizeclass(PageIdContaining(ptr)) != 0); if (ABSL_PREDICT_FALSE(!is_cold)) { // Outline cold path to avoid putting cold size lookup on the fast path. SLOW_PATH_BARRIER(); @@ -1039,7 +1041,9 @@ bool CorrectSize(const void* ptr, const size_t provided_size, Policy policy) { // Recompute the provided size and how it maps onto a size class. const hot_cold_t access_hint = - ABSL_PREDICT_FALSE(GetMemoryTag(ptr) == MemoryTag::kCold) + ABSL_PREDICT_FALSE( + (GetMemoryTag(ptr) == MemoryTag::kCold || policy.is_cold()) && + policy.allocation_type() == AllocationType::New) ? hot_cold_t{0} : hot_cold_t{255}; auto [is_small, provided_size_class] = tc_globals.sizemap().GetSizeClass( diff --git a/tcmalloc/testing/heap_profiling_test.cc b/tcmalloc/testing/heap_profiling_test.cc index 719a240b4..b04e3ad50 100644 --- a/tcmalloc/testing/heap_profiling_test.cc +++ b/tcmalloc/testing/heap_profiling_test.cc @@ -351,7 +351,9 @@ TEST(HeapProfilingTest, MadviseSampledAllocations) { allocs[i] = allocate(); switch (test_case.heap) { case AllocationHeap::kSampled: - EXPECT_TRUE(tcmalloc_internal::IsSampledMemory(allocs[i])); + EXPECT_EQ(tcmalloc_internal::GetMemoryTag(allocs[i]), + test_case.guarded ? tcmalloc_internal::MemoryTag::kSampled + : tcmalloc_internal::MemoryTag::kCold); break; case AllocationHeap::kCold: EXPECT_EQ(tcmalloc_internal::GetMemoryTag(allocs[i]), diff --git a/tcmalloc/testing/partitioning_fuzz_test.cc b/tcmalloc/testing/partitioning_fuzz_test.cc index f58f55508..23856ac1d 100644 --- a/tcmalloc/testing/partitioning_fuzz_test.cc +++ b/tcmalloc/testing/partitioning_fuzz_test.cc @@ -447,6 +447,7 @@ void RandomizedAllocateAndDeallocateFuzzTest( std::vector live_allocs; live_allocs.reserve(actions.size()); tcmalloc::ScopedGuardedSamplingInterval no_guarded_sampling(-1); + tcmalloc::ScopedNeverSample never_sample; for (const auto& action : actions) { if (std::holds_alternative(action)) { diff --git a/tcmalloc/testing/tcmalloc_test.cc b/tcmalloc/testing/tcmalloc_test.cc index cbffd4053..38deb9bb4 100644 --- a/tcmalloc/testing/tcmalloc_test.cc +++ b/tcmalloc/testing/tcmalloc_test.cc @@ -1298,14 +1298,13 @@ static bool IsHot(uint8_t label, // allocations as hot to avoid mixing pointer-containing and pointerless // allocations in the same cold partition. return static_cast(label) >= threshold || - (MallocExtension::GetNumericProperty( - "tcmalloc.security_partitioning_active") - .value_or(0) == 1 && + (Parameters::heap_partitioning_mode() == HeapPartitioningMode::kFull && std::is_same_v); } TYPED_TEST(HotColdTest, HotColdNew) { const bool expectColdTags = tcmalloc_internal::ColdFeatureActive(); + ScopedNeverSample never_sample; absl::flat_hash_set hot; absl::flat_hash_set cold; @@ -1390,6 +1389,7 @@ hot_cold_t MinHotAccessHint() { } TYPED_TEST(HotColdTest, NothrowHotColdNew) { + ScopedNeverSample never_sample; const bool expectColdTags = tcmalloc_internal::ColdFeatureActive(); if (!expectColdTags) { GTEST_SKIP() << "Cold allocations not enabled"; @@ -1436,6 +1436,7 @@ TYPED_TEST(HotColdTest, NothrowHotColdNew) { } TYPED_TEST(HotColdTest, AlignedNothrowHotColdNew) { + ScopedNeverSample never_sample; const bool expectColdTags = tcmalloc_internal::ColdFeatureActive(); if (!expectColdTags) { GTEST_SKIP() << "Cold allocations not enabled"; @@ -1486,6 +1487,7 @@ TYPED_TEST(HotColdTest, AlignedNothrowHotColdNew) { } TYPED_TEST(HotColdTest, ArrayNothrowHotColdNew) { + ScopedNeverSample never_sample; const bool expectColdTags = tcmalloc_internal::ColdFeatureActive(); if (!expectColdTags) { GTEST_SKIP() << "Cold allocations not enabled"; @@ -1532,6 +1534,7 @@ TYPED_TEST(HotColdTest, ArrayNothrowHotColdNew) { } TYPED_TEST(HotColdTest, ArrayAlignedNothrowHotColdNew) { + ScopedNeverSample never_sample; const bool expectColdTags = tcmalloc_internal::ColdFeatureActive(); if (!expectColdTags) { GTEST_SKIP() << "Cold allocations not enabled"; @@ -1582,6 +1585,7 @@ TYPED_TEST(HotColdTest, ArrayAlignedNothrowHotColdNew) { } TYPED_TEST(HotColdTest, SizeReturningHotColdNew) { + ScopedNeverSample never_sample; const bool expectColdTags = tcmalloc_internal::ColdFeatureActive(); if (!expectColdTags) { GTEST_SKIP() << "Cold allocations not enabled"; @@ -1645,6 +1649,7 @@ TYPED_TEST(HotColdTest, SizeReturningHotColdNew) { // Test that setting the min_hot_access_hint parameter has the expected effect // on treatment of the allocated data as cold. TYPED_TEST(HotColdTest, HotColdNewMinHotFlag) { + ScopedNeverSample never_sample; const bool expectColdTags = tcmalloc_internal::ColdFeatureActive(); if (!expectColdTags) { GTEST_SKIP() << "Cold allocations not enabled"; diff --git a/tcmalloc/testing/want_hpaa_test_helper.cc b/tcmalloc/testing/want_hpaa_test_helper.cc index 245e4469c..ef6bc479f 100644 --- a/tcmalloc/testing/want_hpaa_test_helper.cc +++ b/tcmalloc/testing/want_hpaa_test_helper.cc @@ -30,7 +30,8 @@ int main(int argc, char** argv) { bool hpaa = false; int subrelease = -1; for (absl::string_view line : absl::StrSplit(input, '\n')) { - if (absl::StrContains(line, "Begin SAMPLED page allocator")) { + if (absl::StrContains(line, "Begin COLD page allocator") || + absl::StrContains(line, "Begin SAMPLED page allocator")) { // Stop when we reach the end of the main page allocator. We don't // want to look at the sampled or cold allocator parameters for this // test.