Refactor UpdateMaxCapacities and ResizeSlabs to initialize CPU slabs before stopping all CPUs.

tcmalloc/cpu_cache.h

@@ -1697,13 +1697,30 @@ void CpuCache<Forwarder>::ResizeSizeClassMaxCapacities()
         ShiftOffset(per_cpu_shift, shift_bounds_.initial_shift),
         new_resize_slab_offset);
 
+    std::atomic<uint16_t> updated_max_capacities[kNumClasses];
+    for (size_t i = 0; i < kNumClasses; ++i) {
+      updated_max_capacities[i].store(
+          max_capacity_[i].load(std::memory_order_relaxed),
+          std::memory_order_relaxed);
+    }
+
+    for (int i = 0; i < to_update; ++i) {
+      updated_max_capacities[new_max_capacities[i].size_class].store(
+          new_max_capacities[i].max_capacity, std::memory_order_relaxed);
+    }
+
     info = freelist_.UpdateMaxCapacities(
         new_slabs,
-        GetShiftMaxCapacity{max_capacity_, per_cpu_shift, shift_bounds_},
+        GetShiftMaxCapacity{updated_max_capacities, per_cpu_shift,
+                            shift_bounds_},
         [this](int size_class, uint16_t cap) {
           UpdateMaxCapacity(size_class, cap);
         },
-        [this](int cpu) { return HasPopulated(cpu); },
+        [this](int cpu) {
+          // Since this is called while holding all resize_ locks,
+          // HasPopulated should not change.
+          return HasPopulated(cpu);
+        },
         DrainHandler<CpuCache>{*this, nullptr}, new_max_capacities, to_update);
   }
   for (int cpu = 0; cpu < num_cpus; ++cpu) resize_[cpu].lock.unlock();

tcmalloc/internal/percpu_tcmalloc.h

@@ -162,8 +162,12 @@ class TcmallocSlab {
   // Update maximum capacities allocated to each size class.
   // Build and initialize <new_slabs> so as to use new maximum capacities
   // provided by <capacity> callback for the <size_class>.
+  // <capacity> should return the new maximum capacity for the given size
+  // class, regardless of whether <update_capacity> has been called or not.
   // <update_capacity> updates capacities for the <size_class> with the new
   // <cap> once the slabs are initialized.
+  // <populated> returns whether the given cpu's slab is populated. The
+  // return value should remain constant for the duration of the call.
   // <new_max_capacity> provides an array of new maximum capacities to be
   // updated for size classes.
   // <classes_to_resize> provides the number of size classes for which the
@@ -1083,7 +1087,8 @@ template <size_t NumClasses>
 void TcmallocSlab<NumClasses>::InitCpuImpl(
     void* slabs, Shift shift, int cpu,
     absl::FunctionRef<size_t(size_t)> capacity) {
-  TC_CHECK(state_[cpu].stopped.load(std::memory_order_relaxed));
+  TC_CHECK(slabs != GetSlabsAndShift(std::memory_order_relaxed).first ||
+           state_[cpu].stopped.load(std::memory_order_relaxed));
   TC_CHECK_LE((1 << ToUint8(shift)), (1 << 16) * sizeof(void*));
 
   // Initialize prefetch target and compute the offsets for the
@@ -1251,16 +1256,26 @@ ResizeSlabsInfo TcmallocSlab<NumClasses>::UpdateMaxCapacities(
     absl::FunctionRef<void(int, uint16_t)> update_capacity,
     absl::FunctionRef<bool(size_t)> populated, DrainHandler drain_handler,
     PerSizeClassMaxCapacity* new_max_capacity, int classes_to_resize) {
+  const int n_cpus = num_cpus();
+  const auto [old_slabs, shift] = GetSlabsAndShift(std::memory_order_relaxed);
+
+  // Phase 0: Initialize slabs for populated CPUs BEFORE stopping CPUs.
+  // This prefaults pages while all CPUs continue running undisturbed.
+  for (size_t cpu = 0; cpu < n_cpus; ++cpu) {
+    // populated should not change while this function runs because the caller
+    // should be holding all the CPU resize locks.
+    if (!populated(cpu)) continue;
+    InitCpuImpl(new_slabs, shift, cpu, capacity);
+  }
+
   // Phase 1: Stop all CPUs and initialize any CPUs in the new slab that have
   // already been populated in the old slab.
-  const auto [old_slabs, shift] = GetSlabsAndShift(std::memory_order_relaxed);
   std::array<uint16_t, NumClasses> old_begins;
   for (int size_class = 1; size_class < NumClasses; ++size_class) {
     old_begins[size_class] =
         begins_[size_class].load(std::memory_order_relaxed);
   }
 
-  const int n_cpus = num_cpus();
   for (auto& state : state_) {
     TC_CHECK(!state.stopped.load(std::memory_order_relaxed));
     state.stopped.store(true, std::memory_order_relaxed);
@@ -1280,10 +1295,6 @@ ResizeSlabsInfo TcmallocSlab<NumClasses>::UpdateMaxCapacities(
   }
 
   // Phase 3: Initialize slabs.
-  for (size_t cpu = 0; cpu < n_cpus; ++cpu) {
-    if (!populated(cpu)) continue;
-    InitCpuImpl(new_slabs, shift, cpu, capacity);
-  }
   InitSlabs(new_slabs, shift, capacity);
 
   // Phase 4: Re-start all CPUs.
@@ -1305,8 +1316,9 @@ auto TcmallocSlab<NumClasses>::ResizeSlabs(
     absl::FunctionRef<size_t(size_t)> capacity,
     absl::FunctionRef<bool(size_t)> populated, DrainHandler drain_handler)
     -> ResizeSlabsInfo {
-  // Phase 1: Collect begins, stop all CPUs and initialize any CPUs in the new
-  // slab that have already been populated in the old slab.
+  // Phase 1: Collect begins, initialize any CPUs in the new
+  // slab that have already been populated in the old slab,
+  // then stop all CPUs.
   const auto [old_slabs, old_shift] =
       GetSlabsAndShift(std::memory_order_relaxed);
   std::array<uint16_t, NumClasses> old_begins;
@@ -1318,12 +1330,15 @@ auto TcmallocSlab<NumClasses>::ResizeSlabs(
   TC_ASSERT_NE(new_shift, old_shift);
   const int n_cpus = num_cpus();
   for (size_t cpu = 0; cpu < n_cpus; ++cpu) {
-    TC_CHECK(!state_[cpu].stopped.load(std::memory_order_relaxed));
-    state_[cpu].stopped.store(true, std::memory_order_relaxed);
     if (populated(cpu)) {
       InitCpuImpl(new_slabs, new_shift, cpu, capacity);
     }
   }
+
+  for (auto& state : state_) {
+    TC_CHECK(!state.stopped.load(std::memory_order_relaxed));
+    state.stopped.store(true, std::memory_order_relaxed);
+  }
   FenceAllCpus();
 
 #ifdef TCMALLOC_INTERNAL_LATENCY_INJECTION

tcmalloc/internal/percpu_tcmalloc_fuzz.cc

@@ -565,7 +565,13 @@ struct UpdateMaxCapacities {
     void* new_slabs = Malloc(slabs_size, SlabAlignment(shift));
     const auto [old_slabs, old_slabs_size] = state.slab.UpdateMaxCapacities(
         new_slabs,
-        [&state](size_t size_class) { return state.MaxCapacity(size_class); },
+        [&state, sc, target_cap](size_t size_class) {
+          // Return the new max capacity for the size class we want to grow.
+          if (size_class == sc) return target_cap;
+          // Since other size classes' max capacities are not changed, we can
+          // just return their current max capacities.
+          return state.MaxCapacity(size_class);
+        },
         [&state](int size_class, uint16_t cap) {
           state.max_capacity[size_class] = cap;
         },

tcmalloc/internal/percpu_tcmalloc_test.cc

@@ -480,9 +480,16 @@ TEST_F(TcmallocSlabTest, ResizeMaxCapacities) {
     new_max_capacity[1] = PerSizeClassMaxCapacity{
         .size_class = kSizeClassToShrink,
         .max_capacity = max_capacity[kSizeClassToShrink] - 1};
+    std::array<size_t, kStressSlabs> updated_max_capacity;
+    for (int sc = 0; sc < kStressSlabs; ++sc) {
+      updated_max_capacity[sc] = max_capacity[sc];
+    }
+    updated_max_capacity[kSizeClassToGrow] += 1;
+    updated_max_capacity[kSizeClassToShrink] -= 1;
     void* slabs = AllocSlabs(allocator, kShift);
     const auto [old_slabs, old_slabs_size] = slab_.UpdateMaxCapacities(
-        slabs, [&](size_t size_class) { return max_capacity[size_class]; },
+        slabs,
+        [&](size_t size_class) { return updated_max_capacity[size_class]; },
         [&](int size, uint16_t cap) { max_capacity[size] = cap; },
         [](int cpu) { return cpu == kCpu; },
         [&](int cpu, size_t size_class, void** batch, size_t size, size_t cap) {
@@ -889,10 +896,21 @@ void ResizeMaxCapacitiesThread(
     int to_resize = GetResizedMaxCapacities(ctx, new_max_capacity);
     size_t old_slabs_idx = 0;
 
+    std::atomic<size_t> updated_max_capacity[kStressSlabs];
+    for (size_t sc = 0; sc < kStressSlabs; ++sc) {
+      updated_max_capacity[sc].store(
+          ctx.max_capacity[sc].load(std::memory_order_relaxed),
+          std::memory_order_relaxed);
+    }
+    for (int i = 0; i < to_resize; ++i) {
+      updated_max_capacity[new_max_capacity[i].size_class].store(
+          new_max_capacity[i].max_capacity, std::memory_order_relaxed);
+    }
+
     uint8_t shift = ctx.slab->GetShift();
     void* slabs = AllocSlabs(allocator, shift);
     const auto [old_slabs, old_slabs_size] = ctx.slab->UpdateMaxCapacities(
-        slabs, ctx.GetMaxCapacityFunctor(),
+        slabs, GetMaxCapacity{updated_max_capacity},
         [&](int size, uint16_t cap) {
           ctx.max_capacity[size].store(cap, std::memory_order_relaxed);
         },