From 53974e2fd50b5fb97b4a5836361b76e2854b64c8 Mon Sep 17 00:00:00 2001
From: Vaibhav Gogte <vgogte@google.com>
Date: Wed, 29 Jan 2025 14:36:14 -0800
Subject: [PATCH] Enable TCMalloc to use reuse size classes by default.

PiperOrigin-RevId: 721111704
Change-Id: I21f01afd223d68ec1a34c073fa7204f05ed0d028
---
 tcmalloc/BUILD             |  1 +
 tcmalloc/cpu_cache.h       | 16 ++++++++++++++--
 tcmalloc/cpu_cache_test.cc | 11 +++++++----
 tcmalloc/static_vars.cc    |  9 +++------
 4 files changed, 25 insertions(+), 12 deletions(-)

diff --git a/tcmalloc/BUILD b/tcmalloc/BUILD
index 456d90689..feaf22b3c 100644
--- a/tcmalloc/BUILD
+++ b/tcmalloc/BUILD
@@ -1323,6 +1323,7 @@ cc_test(
     deps = [
         ":common_8k_pages",
         ":mock_transfer_cache",
+        ":size_class_info",
         "//tcmalloc/internal:affinity",
         "//tcmalloc/internal:logging",
         "//tcmalloc/internal:optimization",
diff --git a/tcmalloc/cpu_cache.h b/tcmalloc/cpu_cache.h
index 0de0bde25..164e06f0a 100644
--- a/tcmalloc/cpu_cache.h
+++ b/tcmalloc/cpu_cache.h
@@ -141,6 +141,11 @@ class StaticForwarder {
     return Parameters::per_cpu_caches_dynamic_slab_shrink_threshold();
   }
 
+  static bool reuse_size_classes() {
+    return tc_globals.size_class_configuration() ==
+           SizeClassConfiguration::kReuse;
+  }
+
   static size_t class_to_size(int size_class) {
     return tc_globals.sizemap().class_to_size(size_class);
   }
@@ -851,8 +856,15 @@ inline size_t CpuCache<Forwarder>::MaxCapacity(size_t size_class) const {
   if (ColdFeatureActive()) {
     // We reduce the number of cached objects for some sizes to fit into the
     // slab.
-    const uint16_t kLargeUninterestingObjectDepth = 133 * kWiderSlabMultiplier;
-    const uint16_t kLargeInterestingObjectDepth = 28 * kWiderSlabMultiplier;
+    //
+    // We use fewer number of size classes when using reuse size classes. So,
+    // we may use larger capacity for some sizes.
+    const uint16_t kLargeUninterestingObjectDepth =
+        forwarder_.reuse_size_classes() ? 246 * kWiderSlabMultiplier
+                                        : 133 * kWiderSlabMultiplier;
+    const uint16_t kLargeInterestingObjectDepth =
+        forwarder_.reuse_size_classes() ? 46 * kWiderSlabMultiplier
+                                        : 28 * kWiderSlabMultiplier;
 
     absl::Span<const size_t> cold = forwarder_.cold_size_classes();
     if (absl::c_binary_search(cold, size_class)) {
diff --git a/tcmalloc/cpu_cache_test.cc b/tcmalloc/cpu_cache_test.cc
index 3203f533d..68eb45b67 100644
--- a/tcmalloc/cpu_cache_test.cc
+++ b/tcmalloc/cpu_cache_test.cc
@@ -46,6 +46,7 @@
 #include "tcmalloc/internal/sysinfo.h"
 #include "tcmalloc/mock_transfer_cache.h"
 #include "tcmalloc/parameters.h"
+#include "tcmalloc/size_class_info.h"
 #include "tcmalloc/sizemap.h"
 #include "tcmalloc/static_vars.h"
 #include "tcmalloc/tcmalloc_policy.h"
@@ -161,6 +162,8 @@ class TestStaticForwarder {
                : -1.0;
   }
 
+  bool reuse_size_classes() const { return true; }
+
   size_t class_to_size(int size_class) const {
     if (size_map_.has_value()) {
       return size_map_->class_to_size(size_class);
@@ -1167,7 +1170,7 @@ TEST(CpuCacheTest, DynamicSlabThreshold) {
   forwarder.dynamic_slab_enabled_ = true;
   forwarder.dynamic_slab_grow_threshold_ = kDynamicSlabGrowThreshold;
   SizeMap size_map;
-  size_map.Init(kSizeClasses.classes);
+  size_map.Init(size_map.CurrentClasses().classes);
   forwarder.size_map_ = size_map;
 
   cache.Activate();
@@ -1221,7 +1224,7 @@ TEST(CpuCacheTest, DynamicSlabParamsChange) {
 #endif
 
   SizeMap size_map;
-  size_map.Init(kSizeClasses.classes);
+  size_map.Init(size_map.CurrentClasses().classes);
   for (bool initially_enabled : {false, true}) {
     for (DynamicSlab initial_dynamic_slab :
          {DynamicSlab::kGrow, DynamicSlab::kShrink, DynamicSlab::kNoop}) {
@@ -1281,7 +1284,7 @@ TEST(CpuCacheTest, MaxCapacityResizeFailedBytesMlocked) {
   cache.Activate();
 
   SizeMap size_map;
-  size_map.Init(kSizeClasses.classes);
+  size_map.Init(size_map.CurrentClasses().classes);
   forwarder.size_map_ = size_map;
 
   std::vector<std::thread> threads;
@@ -1330,7 +1333,7 @@ TEST(CpuCacheTest, SlabResizeFailedBytesMlocked) {
   cache.Activate();
 
   SizeMap size_map;
-  size_map.Init(kSizeClasses.classes);
+  size_map.Init(size_map.CurrentClasses().classes);
   forwarder.size_map_ = size_map;
 
   std::vector<std::thread> threads;
diff --git a/tcmalloc/static_vars.cc b/tcmalloc/static_vars.cc
index 2c55a8764..aaacbbb84 100644
--- a/tcmalloc/static_vars.cc
+++ b/tcmalloc/static_vars.cc
@@ -188,18 +188,15 @@ SizeClassConfiguration Static::size_class_configuration() {
 
   const char* e = thread_safe_getenv("TCMALLOC_LEGACY_SIZE_CLASSES");
   if (e == nullptr) {
-    // TODO(b/358126781): Change this to use reuse size classes.
-    return SizeClassConfiguration::kPow2Below64;
+    return SizeClassConfiguration::kReuse;
   } else if (!strcmp(e, "pow2below64")) {
     return SizeClassConfiguration::kPow2Below64;
   } else if (!strcmp(e, "0")) {
-    // TODO(b/358126781): Change this to use reuse size classes.
-    return SizeClassConfiguration::kPow2Below64;
+    return SizeClassConfiguration::kReuse;
   } else {
     TC_BUG("bad TCMALLOC_LEGACY_SIZE_CLASSES env var '%s'", e);
   }
-  // TODO(b/358126781): Change this to use reuse size classes.
-  return SizeClassConfiguration::kPow2Below64;
+  return SizeClassConfiguration::kReuse;
 }
 
 ABSL_ATTRIBUTE_COLD ABSL_ATTRIBUTE_NOINLINE void Static::SlowInitIfNecessary() {