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); },