Transfer pointers and capacity during slab resizing, too.

tcmalloc/internal/percpu_tcmalloc.h

@@ -463,6 +463,11 @@ class TcmallocSlab {
   void DrainOldSlabs(void* slabs, Shift shift, int cpu,
                      const std::array<uint16_t, NumClasses>& old_begins,
                      DrainHandler drain_handler);
+  void TransferFreelistsToNewSlabs(
+      void* old_slabs, void* new_slabs,
+      const std::array<uint16_t, NumClasses>& old_begins, Shift old_shift,
+      Shift new_shift, absl::FunctionRef<size_t(size_t)> max_capacity,
+      absl::FunctionRef<bool(size_t)> populated);
 
   // Implementation of InitCpu() allowing for reuse in ResizeSlabs().
   void InitCpuImpl(void* slabs, Shift shift, int cpu,
@@ -1348,79 +1353,40 @@ void TcmallocSlab<NumClasses>::DrainOldSlabs(
   }
 }
 
+// Note that since we are always moving and never copying or dropping,
+// none of these will change the amount of RAM that we are using,
+// so we will not need to call back into CpuCache to update the available
+// counter, nor will we need to worry about going over that limit.
 template <size_t NumClasses>
-ResizeSlabsInfo TcmallocSlab<NumClasses>::UpdateMaxCapacities(
-    void* new_slabs, absl::FunctionRef<size_t(size_t)> max_capacity,
-    absl::FunctionRef<void(int, uint16_t)> update_max_capacity,
-    absl::FunctionRef<bool(size_t)> populated, DrainHandler drain_handler,
-    PerSizeClassMaxCapacity* new_max_capacity, int classes_to_resize) {
+void TcmallocSlab<NumClasses>::TransferFreelistsToNewSlabs(
+    void* old_slabs, void* new_slabs,
+    const std::array<uint16_t, NumClasses>& old_begins, Shift old_shift,
+    Shift new_shift, absl::FunctionRef<size_t(size_t)> max_capacity,
+    absl::FunctionRef<bool(size_t)> populated) {
   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, max_capacity);
-  }
-
-  // Phase 1: Stop all CPUs and initialize any CPUs in the new slab that have
-  // already been populated in the old slab.
-  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);
-  }
-
-  StopAllCpus();
-
-#ifdef TCMALLOC_INTERNAL_LATENCY_INJECTION
-  // TODO(b/29448043): Remove latency injection.
-  ScopedDelay delay(ScopedDelay::update_max_capacities_delay);
-#endif
-
-  // Phase 2: Update max capacity of the size classes.
-  for (int i = 0; i < classes_to_resize; ++i) {
-    size_t size_class = new_max_capacity[i].size_class;
-    size_t cap = new_max_capacity[i].max_capacity;
-    update_max_capacity(size_class, cap);
-  }
-
-  // Phase 3: Initialize slabs.
-  InitSlabs(new_slabs, shift, max_capacity);
-
-  // Phase 4: Copy the old capacities and pointers to the new slabs,
-  // subject to possibly reduced new limits.
-  //
-  // Note that since we are always moving and never copying or dropping,
-  // none of these will change the amount of RAM that we are using,
-  // so we will not need to call back into CpuCache to update the available
-  // counter, nor will we need to worry about going over that limit.
   for (size_t cpu = 0; cpu < n_cpus; ++cpu) {
     if (!populated(cpu)) continue;
 
     void** old_start =
-        reinterpret_cast<void**>(CpuMemoryStart(old_slabs, shift, cpu));
+        reinterpret_cast<void**>(CpuMemoryStart(old_slabs, old_shift, cpu));
     void** new_start =
-        reinterpret_cast<void**>(CpuMemoryStart(new_slabs, shift, cpu));
+        reinterpret_cast<void**>(CpuMemoryStart(new_slabs, new_shift, cpu));
     for (int size_class = 1; size_class < NumClasses; ++size_class) {
       uint16_t new_max_capacity = max_capacity(size_class);
       if (new_max_capacity == 0) {
         // Nothing to transfer.
         continue;
       }
 
-      auto* ohdrp = GetHeader(old_slabs, shift, cpu, size_class);
+      auto* ohdrp = GetHeader(old_slabs, old_shift, cpu, size_class);
       Header ohdr = LoadHeader(ohdrp);
       uint16_t pointers_in_list = ohdr.length(old_begins[size_class]);
       if (pointers_in_list == 0 && ohdr.remaining_capacity == 0) {
         // Nothing to transfer.
         continue;
       }
 
-      auto* nhdrp = GetHeader(new_slabs, shift, cpu, size_class);
+      auto* nhdrp = GetHeader(new_slabs, new_shift, cpu, size_class);
       Header nhdr = LoadHeader(nhdrp);
 
       // First move any elements in the freelist that we can. Note that we also
@@ -1433,6 +1399,7 @@ ResizeSlabsInfo TcmallocSlab<NumClasses>::UpdateMaxCapacities(
                elems_to_copy * sizeof(*new_start));
         ohdr.current -= elems_to_copy;
         nhdr.current += elems_to_copy;
+        TC_ASSERT_NE(nhdr.current, Header::kReservedIndexValue);
       }
 
       // Now transfer any unused capacity that we can. Most of the time, this
@@ -1450,6 +1417,55 @@ ResizeSlabsInfo TcmallocSlab<NumClasses>::UpdateMaxCapacities(
       StoreHeader(nhdrp, nhdr);
     }
   }
+}
+
+template <size_t NumClasses>
+ResizeSlabsInfo TcmallocSlab<NumClasses>::UpdateMaxCapacities(
+    void* new_slabs, absl::FunctionRef<size_t(size_t)> max_capacity,
+    absl::FunctionRef<void(int, uint16_t)> update_max_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, max_capacity);
+  }
+
+  // Phase 1: Stop all CPUs and initialize any CPUs in the new slab that have
+  // already been populated in the old slab.
+  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);
+  }
+
+  StopAllCpus();
+
+#ifdef TCMALLOC_INTERNAL_LATENCY_INJECTION
+  // TODO(b/29448043): Remove latency injection.
+  ScopedDelay delay(ScopedDelay::update_max_capacities_delay);
+#endif
+
+  // Phase 2: Update max capacity of the size classes.
+  for (int i = 0; i < classes_to_resize; ++i) {
+    size_t size_class = new_max_capacity[i].size_class;
+    size_t cap = new_max_capacity[i].max_capacity;
+    update_max_capacity(size_class, cap);
+  }
+
+  // Phase 3: Initialize slabs.
+  InitSlabs(new_slabs, shift, max_capacity);
+
+  // Phase 4: Copy the old capacities and pointers to the new slabs,
+  // subject to possibly reduced new limits.
+  TransferFreelistsToNewSlabs(old_slabs, new_slabs, old_begins, shift, shift,
+                              max_capacity, populated);
 
   // Phase 5: Re-start all CPUs.
   StartAllCpus();
@@ -1467,7 +1483,7 @@ ResizeSlabsInfo TcmallocSlab<NumClasses>::UpdateMaxCapacities(
 template <size_t NumClasses>
 auto TcmallocSlab<NumClasses>::ResizeSlabs(
     Shift new_shift, void* new_slabs,
-    absl::FunctionRef<size_t(size_t)> capacity,
+    absl::FunctionRef<size_t(size_t)> max_capacity,
     absl::FunctionRef<void()> commit_new_capacities,
     absl::FunctionRef<bool(size_t)> populated, DrainHandler drain_handler)
     -> ResizeSlabsInfo {
@@ -1486,7 +1502,7 @@ auto TcmallocSlab<NumClasses>::ResizeSlabs(
   const int n_cpus = num_cpus();
   for (size_t cpu = 0; cpu < n_cpus; ++cpu) {
     if (populated(cpu)) {
-      InitCpuImpl(new_slabs, new_shift, cpu, capacity);
+      InitCpuImpl(new_slabs, new_shift, cpu, max_capacity);
     }
   }
 
@@ -1498,16 +1514,21 @@ auto TcmallocSlab<NumClasses>::ResizeSlabs(
 #endif
 
   // Phase 2: Atomically update slabs and shift.
-  InitSlabs(new_slabs, new_shift, capacity);
+  InitSlabs(new_slabs, new_shift, max_capacity);
 
   // Phase 3: Update the global max_capacity_ array, now that the CPUs
   // are stopped.
   commit_new_capacities();
 
-  // Phase 4: Re-start all CPUs.
+  // Phase 4: Copy the old capacities and pointers to the new slabs,
+  // subject to possibly reduced new limits.
+  TransferFreelistsToNewSlabs(old_slabs, new_slabs, old_begins, old_shift,
+                              new_shift, max_capacity, populated);
+
+  // Phase 5: Re-start all CPUs.
   StartAllCpus();
 
-  // Phase 5: Return pointers from the old slab to the TransferCache.
+  // Phase 6: Return pointers from the old slab to the TransferCache.
   for (size_t cpu = 0; cpu < n_cpus; ++cpu) {
     if (!populated(cpu)) continue;
     DrainOldSlabs(old_slabs, old_shift, cpu, old_begins, drain_handler);

tcmalloc/internal/percpu_tcmalloc_fuzz.cc

@@ -572,17 +572,20 @@ struct ResizeSlabs {
         [&state](size_t cpu) { return state.cpu_initialized[cpu]; },
         [&state](int cpu, size_t size_class, void** batch, size_t size,
                  size_t cap) {
-          state.HandleDrain(cpu, size_class, batch, size, cap,
-                            state.expected_capacity[cpu][size_class]);
+          const size_t old_cap = state.expected_capacity[cpu][size_class];
+          const size_t kept =
+              std::min<size_t>(old_cap, state.max_capacity[size_class]);
+          state.HandleDrain(cpu, size_class, batch, size, cap, old_cap - kept);
         });
     TC_CHECK_EQ(state.slab.GetShift(), target_shift);
     const Shift old_shift = ToShiftType(current_shift);
     sized_aligned_delete(old_slabs, old_slabs_size, SlabAlignment(old_shift));
-    // ResizeSlabs drains every CPU's old slab completely and starts at
-    // capacity 0 in the new slabs.
+    // ResizeSlabs() transfers existing pointers and capacities up to the
+    // new max capacities into the new slabs, draining only any excess via
+    // HandleDrain.
     for (const auto& caps : state.expected_capacity) {
-      for (size_t sc = 1; sc < kNumClasses; ++sc) {
-        TC_CHECK_EQ(caps[sc], 0);
+      for (size_t size_class = 1; size_class < kNumClasses; ++size_class) {
+        TC_CHECK_LE(caps[size_class], state.max_capacity[size_class]);
       }
     }
     auto [got_cpu, cached] = state.slab.CacheCpuSlab();

tcmalloc/internal/percpu_tcmalloc_test.cc

@@ -680,6 +680,54 @@ TEST_F(TcmallocSlabTest, ShrinkEmptyCache) {
   slab_.StartCpu(kCpu);
 }
 
+TEST_F(TcmallocSlabTest, ResizeSlabKeepsCapacity) {
+  if (!IsFast()) {
+    GTEST_SKIP() << "Need fast percpu. Skipping.";
+    return;
+  }
+
+  constexpr unsigned kCpu = 0;
+  ScopedFakeCpuId fake_cpu_id(kCpu);
+  constexpr unsigned kSizeClass = 1;
+
+  // Initialize a core.
+  slab_.InitCpu(kCpu, [](size_t size_class) { return kCapacity; });
+  auto [got_cpu, cached] = slab_.CacheCpuSlab();
+  ASSERT_TRUE(cached);
+  ASSERT_EQ(got_cpu, kCpu);
+
+  // Push some things onto a size class, growing it as we go.
+  EXPECT_FALSE(slab_.Push(kSizeClass, (void*)0x40));  // No capacity.
+  EXPECT_EQ(3, slab_.Grow(kCpu, kSizeClass, 3, [&]() {
+    return 200;  // Any max capacity, really.
+  }));
+  EXPECT_TRUE(slab_.Push(kSizeClass, (void*)0x40));
+  EXPECT_TRUE(slab_.Push(kSizeClass, (void*)0x80));
+
+  // Resize the blocks.
+  auto alloc = [&](size_t size, std::align_val_t alignment) {
+    return ByteCountingMalloc(size, alignment);
+  };
+  void* slabs = AllocSlabs(alloc, kShift + 1);
+  (void)slab_.ResizeSlabs(
+      subtle::percpu::ToShiftType(kShift + 1), slabs,
+      [](size_t) { return kCapacity * 2; }, [] { /* Nothing to commit. */ },
+      [](int cpu) { return cpu == kCpu; },
+      [&](int cpu, size_t size_class, void** batch, size_t size, size_t cap) {
+        EXPECT_EQ(size, 0);
+        EXPECT_EQ(cap, 0);
+      });
+
+  auto [got_cpu2, cached2] = slab_.CacheCpuSlab();
+  ASSERT_TRUE(cached2);
+  ASSERT_EQ(got_cpu2, kCpu);
+
+  // Both the two elements and the remaining capacity should still be there.
+  EXPECT_EQ(3, slab_.Capacity(kCpu, kSizeClass));
+  EXPECT_EQ((void*)0x80, slab_.Pop(kSizeClass));
+  EXPECT_EQ((void*)0x40, slab_.Pop(kSizeClass));
+}
+
 TEST_F(TcmallocSlabTest, SimulatedMadviseFailure) {
   if (!IsFast()) {
     GTEST_SKIP() << "Need fast percpu. Skipping.";