diff --git a/tcmalloc/internal/percpu_tcmalloc.h b/tcmalloc/internal/percpu_tcmalloc.h index b947a1800..83dd6de18 100644 --- a/tcmalloc/internal/percpu_tcmalloc.h +++ b/tcmalloc/internal/percpu_tcmalloc.h @@ -274,6 +274,14 @@ class TcmallocSlab { void StopCpu(int cpu); void StartCpu(int cpu); + // Stop/start all cpus at once. StopAllCpus also executes the Fence for all + // cpus, so that no local operation is in flight when it returns. + void StopAllCpus(); + void StartAllCpus(); + + // Asserts that is currently stopped, as required by remote operations. + void AssertCpuStopped(int cpu) const; + // Grows the cpu/size_class slab's capacity to no greater than // min(capacity+len, max_capacity()) and returns the increment // applied. @@ -1192,7 +1200,7 @@ std::pair TcmallocSlab::CacheCpuSlabSlow() { template void TcmallocSlab::DrainCpu(void* slabs, Shift shift, int cpu, DrainHandler drain_handler) { - TC_ASSERT(state_[cpu].stopped.load(std::memory_order_relaxed)); + AssertCpuStopped(cpu); for (size_t size_class = 1; size_class < NumClasses; ++size_class) { uint16_t begin = begins_[size_class].load(std::memory_order_relaxed); auto* hdrp = GetHeader(slabs, shift, cpu, size_class); @@ -1277,11 +1285,7 @@ ResizeSlabsInfo TcmallocSlab::UpdateMaxCapacities( begins_[size_class].load(std::memory_order_relaxed); } - for (auto& state : state_) { - TC_CHECK(!state.stopped.load(std::memory_order_relaxed)); - state.stopped.store(true, std::memory_order_relaxed); - } - FenceAllCpus(); + StopAllCpus(); #ifdef TCMALLOC_INTERNAL_LATENCY_INJECTION // TODO(b/29448043): Remove latency injection. @@ -1299,9 +1303,7 @@ ResizeSlabsInfo TcmallocSlab::UpdateMaxCapacities( InitSlabs(new_slabs, shift, capacity); // Phase 4: Re-start all CPUs. - for (auto& state : state_) { - state.stopped.store(false, std::memory_order_release); - } + StartAllCpus(); // Phase 5: Return pointers from the old slab to the TransferCache. for (size_t cpu = 0; cpu < n_cpus; ++cpu) { @@ -1336,11 +1338,7 @@ auto TcmallocSlab::ResizeSlabs( } } - for (auto& state : state_) { - TC_CHECK(!state.stopped.load(std::memory_order_relaxed)); - state.stopped.store(true, std::memory_order_relaxed); - } - FenceAllCpus(); + StopAllCpus(); #ifdef TCMALLOC_INTERNAL_LATENCY_INJECTION // TODO(b/29448043): Remove latency injection. @@ -1351,9 +1349,7 @@ auto TcmallocSlab::ResizeSlabs( InitSlabs(new_slabs, new_shift, capacity); // Phase 3: Re-start all CPUs. - for (auto& state : state_) { - state.stopped.store(false, std::memory_order_release); - } + StartAllCpus(); // Phase 4: Return pointers from the old slab to the TransferCache. for (size_t cpu = 0; cpu < n_cpus; ++cpu) { @@ -1388,7 +1384,7 @@ template size_t TcmallocSlab::GrowOtherCache( int cpu, size_t size_class, size_t len, absl::FunctionRef max_capacity) { - TC_ASSERT(state_[cpu].stopped.load(std::memory_order_relaxed)); + AssertCpuStopped(cpu); const auto [slabs, shift] = GetSlabsAndShift(std::memory_order_relaxed); const size_t max_cap = max_capacity(ToUint8(shift)); auto* hdrp = GetHeader(slabs, shift, cpu, size_class); @@ -1403,7 +1399,7 @@ size_t TcmallocSlab::GrowOtherCache( template size_t TcmallocSlab::ShrinkOtherCache( int cpu, size_t size_class, size_t len, ShrinkHandler shrink_handler) { - TC_ASSERT(state_[cpu].stopped.load(std::memory_order_relaxed)); + AssertCpuStopped(cpu); const auto [slabs, shift] = GetSlabsAndShift(std::memory_order_relaxed); auto* hdrp = GetHeader(slabs, shift, cpu, size_class); @@ -1462,11 +1458,7 @@ void TcmallocSlab::ReleaseSlabMetadataForDrainedCpus( // Stop all CPUs. They must also be locked, since we are touching the // populated bit later. - for (auto& state : state_) { - TC_CHECK(!state.stopped.load(std::memory_order_relaxed)); - state.stopped.store(true, std::memory_order_relaxed); - } - FenceAllCpus(); + StopAllCpus(); // See which ones are actually drained, and which hugepages we can free. const auto [slabs, shift] = GetSlabsAndShift(std::memory_order_relaxed); @@ -1549,9 +1541,7 @@ void TcmallocSlab::ReleaseSlabMetadataForDrainedCpus( } // Restart the CPUs again. - for (auto& state : state_) { - state.stopped.store(false, std::memory_order_release); - } + StartAllCpus(); } template @@ -1565,10 +1555,31 @@ void TcmallocSlab::StopCpu(int cpu) { template void TcmallocSlab::StartCpu(int cpu) { TC_ASSERT(cpu >= 0 && cpu < num_cpus(), "cpu=%d", cpu); - TC_ASSERT(state_[cpu].stopped.load(std::memory_order_relaxed)); + AssertCpuStopped(cpu); state_[cpu].stopped.store(false, std::memory_order_release); } +template +void TcmallocSlab::StopAllCpus() { + for (auto& state : state_) { + TC_CHECK(!state.stopped.load(std::memory_order_relaxed)); + state.stopped.store(true, std::memory_order_relaxed); + } + FenceAllCpus(); +} + +template +void TcmallocSlab::StartAllCpus() { + for (auto& state : state_) { + state.stopped.store(false, std::memory_order_release); + } +} + +template +void TcmallocSlab::AssertCpuStopped(int cpu) const { + TC_ASSERT(state_[cpu].stopped.load(std::memory_order_relaxed)); +} + template PerCPUMetadataState TcmallocSlab::MetadataMemoryUsage() const { PerCPUMetadataState result;