Fix EstimateSlabBytes().

tcmalloc/cpu_cache.h

@@ -988,23 +988,6 @@ inline void CpuCache<Forwarder>::CalculateMaxCapacityForAllClasses(
   }
 }
 
-// Returns estimated bytes required.
-inline size_t EstimateSlabBytes(MaxCapacityFunctor get_capacity,
-                                size_t header_bytes) {
-  size_t bytes_required = header_bytes;
-
-  for (int size_class = 0; size_class < kNumClasses; ++size_class) {
-    // Each non-empty size class region in the slab is preceded by one padding
-    // pointer that points to itself. (We do this because prefetches of invalid
-    // pointers are slow.)
-    size_t num_pointers = get_capacity(size_class);
-    if (num_pointers > 0) ++num_pointers;
-    bytes_required += sizeof(void*) * num_pointers;
-  }
-
-  return bytes_required;
-}
-
 template <class Forwarder>
 inline uint16_t CpuCache<Forwarder>::GetMaxCapacity(int size_class) const {
   return max_capacity_[size_class].load(std::memory_order_relaxed);
@@ -1077,8 +1060,9 @@ inline void CpuCache<Forwarder>::Activate() {
        shift <= shift_bounds_.max_shift; ++shift) {
     CalculateMaxCapacityForAllClasses(shift, max_capacity_);
     const size_t bytes_available = 1 << shift;
-    const size_t bytes_required = EstimateSlabBytes(
-        GetMaxCapacityFunctor(), Freelist::GetTotalClassHeaderSize());
+    const size_t bytes_required =
+        Freelist::EstimateSlabBytes(GetMaxCapacityFunctor());
+
     // We may make certain size classes no-ops by selecting "0" at runtime, so
     // using a compile-time calculation overestimates worst-case memory usage.
     if (ABSL_PREDICT_FALSE(bytes_required > bytes_available)) {

tcmalloc/cpu_cache_test.cc

@@ -110,8 +110,8 @@ class CpuCachePeer {
 
       cpu_cache.CalculateMaxCapacityForAllClasses(shift, max_capacity);
       const size_t bytes_available = 1 << shift;
-      const size_t bytes_required = EstimateSlabBytes(
-          {max_capacity}, CpuCache::Freelist::GetTotalClassHeaderSize());
+      const size_t bytes_required = CpuCache::Freelist::EstimateSlabBytes(
+          MaxCapacityFunctor{max_capacity});
       EXPECT_GT(bytes_required * 20, bytes_available * 17);
       EXPECT_LE(bytes_required, bytes_available);
     }

tcmalloc/cpu_slab_allocation_golden.txt

@@ -149,7 +149,7 @@ Initial max capacity each size class for 8 kB pages, 16 kB slabs:
  - 147:     0 element(s)
  - 148:     0 element(s)
  - 149:     0 element(s)
-Total: 14496 bytes used, 1888 bytes unused.
+Total: 14504 bytes used, 1880 bytes unused.
 
 Initial max capacity each size class for 8 kB pages, 32 kB slabs:
  -   0:     0 element(s)

tcmalloc/dump_cpu_slab_allocation.cc

@@ -94,9 +94,9 @@ void DumpMaxCapacityForTest() {
     // Note that bytes_used includes Header metadata and guard pointers,
     // so it will be larger than just adding together all max_capacity[]
     // and multiplying by sizeof(void*).
-    int bytes_used = EstimateSlabBytes(
-        get_capacity,
-        subtle::percpu::TcmallocSlab<kNumClasses>::GetTotalClassHeaderSize());
+    int bytes_used =
+        subtle::percpu::TcmallocSlab<kNumClasses>::EstimateSlabBytes(
+            get_capacity);
     printf("Total: %d bytes used, %d bytes unused.\n", bytes_used,
            (1 << shift) - bytes_used);
     printf("\n");

tcmalloc/internal/percpu_tcmalloc.h

@@ -334,6 +334,11 @@ class TcmallocSlab {
 
   [[nodiscard]] PerCPUMetadataState MetadataMemoryUsage() const;
 
+  // Estimate how many bytes would be needed for a per-CPU slab with the given
+  // class capacities. Used in asserts and tests to verify we don't overflow.
+  [[nodiscard]] static size_t EstimateSlabBytes(
+      absl::FunctionRef<size_t(size_t)> capacity);
+
   // Gets the current shift of the slabs. Intended for use by the thread that
   // calls ResizeSlabs().
   [[nodiscard]] uint8_t GetShift() const {
@@ -1110,6 +1115,7 @@ void TcmallocSlab<NumClasses>::InitCpu(
   InitCpuImpl(slabs, shift, cpu, capacity);
 }
 
+// Keep EstimateSlabBytes() below in sync with this.
 template <size_t NumClasses>
 void TcmallocSlab<NumClasses>::InitCpuImpl(
     void* slabs, Shift shift, int cpu,
@@ -1156,6 +1162,37 @@ void TcmallocSlab<NumClasses>::InitCpuImpl(
   }
 }
 
+// A dry-run variant of InitCpuImpl() that tracks its logic closely but does not
+// write anything (nor form invalid pointers).
+template <size_t NumClasses>
+size_t TcmallocSlab<NumClasses>::EstimateSlabBytes(
+    absl::FunctionRef<size_t(size_t)> capacity) {
+  // Follows the same increment logic as elems, above. (We could compute the
+  // initial value here using GetHeader() instead of GetTotalClassHeaderSize(),
+  // but we'd then need to invent fake shift and cpu values, which doesn't make
+  // a lot of sense.)
+  size_t bytes_used_on_curr_slab = GetTotalClassHeaderSize();
+
+  bool prev_empty = false;
+  for (size_t size_class = 1; size_class < NumClasses; ++size_class) {
+    size_t cap = capacity(size_class);
+    TC_CHECK_EQ(static_cast<uint16_t>(cap), cap);
+
+    // This item serves both as the marker of slab begin (Pop checks for low bit
+    // set to understand that it reached begin), and as prefetching stub
+    // (Pop prefetches the previous element and prefetching an invalid pointer
+    // is slow, this is a valid pointer for prefetching).
+    if (!prev_empty) {
+      bytes_used_on_curr_slab += sizeof(void*);
+    }
+    prev_empty = cap == 0;
+
+    bytes_used_on_curr_slab += cap * sizeof(void*);
+  }
+
+  return bytes_used_on_curr_slab;
+}
+
 #if TCMALLOC_INTERNAL_PERCPU_USE_RSEQ
 template <size_t NumClasses>
 std::pair<int, bool> TcmallocSlab<NumClasses>::CacheCpuSlabSlow() {