Merge cold and sampled partitions

tcmalloc/central_freelist.cc

@@ -164,7 +164,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

tcmalloc/page_allocator.cc

@@ -44,8 +44,6 @@ 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 +54,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_));
 }
@@ -195,16 +179,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<HugePageAwareAllocator*>(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);

tcmalloc/page_allocator.h

@@ -204,7 +204,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) {}
@@ -213,11 +214,9 @@ class PageAllocator {
     HugePageAwareAllocator hpaa;
   } choices_[kNumHeaps];
   std::array<Interface*, kNormalPartitions> normal_impl_;
-  std::array<Interface*, kSecurityPartitions> sampled_impl_;
   Interface* cold_impl_;
   Algorithm alg_;
   bool has_cold_impl_;
-  bool sampled_partition_active_;
 
   // Max size of backed spans we will attempt to maintain.
   // Crash if we can't maintain below limits_[kHard], which is guaranteed to be
@@ -250,9 +249,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:
@@ -302,30 +299,20 @@ 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();
   }
   return ret;
 }
 
 inline void PageAllocator::GetSmallSpanStats(SmallSpanStats* result) {
-  SmallSpanStats normal, sampled;
+  SmallSpanStats normal;
   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;
+  *result = normal;
   if (has_cold_impl_) {
     SmallSpanStats cold;
     cold_impl_->GetSmallSpanStats(&cold);
@@ -334,19 +321,13 @@ inline void PageAllocator::GetSmallSpanStats(SmallSpanStats* result) {
 }
 
 inline void PageAllocator::GetLargeSpanStats(LargeSpanStats* result) {
-  LargeSpanStats normal, sampled;
+  LargeSpanStats normal;
   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;
+  *result = normal;
   if (has_cold_impl_) {
     LargeSpanStats cold;
     cold_impl_->GetLargeSpanStats(&cold);
@@ -377,13 +358,6 @@ inline Length PageAllocator::ReleaseAtLeastNPages(Length num_pages,
         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;
 }
@@ -398,11 +372,6 @@ inline PageReleaseStats PageAllocator::GetReleaseStats() const {
     stats += normal_impl_[partition]->GetReleaseStats();
   }
 
-  stats += sampled_impl_[0]->GetReleaseStats();
-  if (sampled_partition_active_) {
-    stats += sampled_impl_[1]->GetReleaseStats();
-  }
-
   return stats;
 }
 
@@ -411,6 +380,9 @@ inline void PageAllocator::Print(Printer& out, MemoryTag tag,
   if (tag == MemoryTag::kCold && !has_cold_impl_) {
     return;
   }
+  if (tag == MemoryTag::kSampled || tag == MemoryTag::kSampledP1) {
+    return;
+  }
 
   const absl::string_view label = MemoryTagToLabel(tag);
   if (tag != MemoryTag::kNormal) {
@@ -427,6 +399,9 @@ inline void PageAllocator::PrintInPbtxt(PbtxtRegion& region, MemoryTag tag,
   if (tag == MemoryTag::kCold && !has_cold_impl_) {
     return;
   }
+  if (tag == MemoryTag::kSampled || tag == MemoryTag::kSampledP1) {
+    return;
+  }
 
   PbtxtRegion pa = region.CreateSubRegion("page_allocator");
   pa.PrintRaw("tag", MemoryTagToLabel(tag));

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<Length>(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<uintptr_t>(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(

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]),

tcmalloc/testing/partitioning_fuzz_test.cc

@@ -447,6 +447,7 @@ void RandomizedAllocateAndDeallocateFuzzTest(
   std::vector<AllocationRecord> 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<AllocationOp>(action)) {

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<tcmalloc::hot_cold_t>(label) >= threshold ||
-         (MallocExtension::GetNumericProperty(
-              "tcmalloc.security_partitioning_active")
-                  .value_or(0) == 1 &&
+         (Parameters::heap_partitioning_mode() == HeapPartitioningMode::kFull &&
           std::is_same_v<T, char*>);
 }
 
 TYPED_TEST(HotColdTest, HotColdNew) {
   const bool expectColdTags = tcmalloc_internal::ColdFeatureActive();
+  ScopedNeverSample never_sample;
 
   absl::flat_hash_set<uintptr_t> hot;
   absl::flat_hash_set<uintptr_t> 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";

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.