From b92cd4e87cd65ee5c935d596b4f66c2315e7f6c9 Mon Sep 17 00:00:00 2001 From: Ryan Karpuszka Date: Mon, 31 Aug 2026 07:17:07 -0700 Subject: [PATCH] Refactor UpdateMaxCapacities and ResizeSlabs to initialize CPU slabs before stopping all CPUs. This avoids page faulting in the newly allocated slabs while CPUs are stopped. Since updating capacities still needs to happen while CPUs are stopped, the capacity callback needs to return the new capacities rather than the current ones. This way initializing the CPU slabs can happen before update_capacity is called. PiperOrigin-RevId: 973865650 --- tcmalloc/cpu_cache.h | 21 +++++++++++-- tcmalloc/internal/percpu_tcmalloc.h | 37 ++++++++++++++++------- tcmalloc/internal/percpu_tcmalloc_fuzz.cc | 8 ++++- tcmalloc/internal/percpu_tcmalloc_test.cc | 22 ++++++++++++-- 4 files changed, 72 insertions(+), 16 deletions(-) diff --git a/tcmalloc/cpu_cache.h b/tcmalloc/cpu_cache.h index 7e98c1b5b..923ce5b19 100644 --- a/tcmalloc/cpu_cache.h +++ b/tcmalloc/cpu_cache.h @@ -1697,13 +1697,30 @@ void CpuCache::ResizeSizeClassMaxCapacities() ShiftOffset(per_cpu_shift, shift_bounds_.initial_shift), new_resize_slab_offset); + std::atomic 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{*this, nullptr}, new_max_capacities, to_update); } for (int cpu = 0; cpu < num_cpus; ++cpu) resize_[cpu].lock.unlock(); diff --git a/tcmalloc/internal/percpu_tcmalloc.h b/tcmalloc/internal/percpu_tcmalloc.h index 56fd3cc63..31502b42a 100644 --- a/tcmalloc/internal/percpu_tcmalloc.h +++ b/tcmalloc/internal/percpu_tcmalloc.h @@ -162,8 +162,12 @@ class TcmallocSlab { // Update maximum capacities allocated to each size class. // Build and initialize so as to use new maximum capacities // provided by callback for the . + // should return the new maximum capacity for the given size + // class, regardless of whether has been called or not. // updates capacities for the with the new // once the slabs are initialized. + // returns whether the given cpu's slab is populated. The + // return value should remain constant for the duration of the call. // provides an array of new maximum capacities to be // updated for size classes. // provides the number of size classes for which the @@ -1083,7 +1087,8 @@ template void TcmallocSlab::InitCpuImpl( void* slabs, Shift shift, int cpu, absl::FunctionRef 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::UpdateMaxCapacities( absl::FunctionRef update_capacity, absl::FunctionRef 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 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::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::ResizeSlabs( absl::FunctionRef capacity, absl::FunctionRef 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 old_begins; @@ -1318,12 +1330,15 @@ auto TcmallocSlab::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 diff --git a/tcmalloc/internal/percpu_tcmalloc_fuzz.cc b/tcmalloc/internal/percpu_tcmalloc_fuzz.cc index fce33e84d..ad36ccd33 100644 --- a/tcmalloc/internal/percpu_tcmalloc_fuzz.cc +++ b/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; }, diff --git a/tcmalloc/internal/percpu_tcmalloc_test.cc b/tcmalloc/internal/percpu_tcmalloc_test.cc index e8796536f..ad0ceaf20 100644 --- a/tcmalloc/internal/percpu_tcmalloc_test.cc +++ b/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 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 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); },