From 405a774bc987d5b873215e607d3df7e71ec5bb27 Mon Sep 17 00:00:00 2001 From: Chris Kennelly Date: Mon, 9 Dec 2024 08:49:23 -0800 Subject: [PATCH] Avoid taking pageheap_lock for span allocations. span_allocator is internally thread-safe, so we do not need to access the page heap lock to guard it. PiperOrigin-RevId: 704302332 Change-Id: I11f91ad80aa0349587bede7c7c6cfc291f7359f2 --- tcmalloc/huge_page_aware_allocator.h | 113 +++++++++++++++------ tcmalloc/mock_huge_page_static_forwarder.h | 7 +- tcmalloc/span.h | 6 +- tcmalloc/span_test.cc | 2 + 4 files changed, 96 insertions(+), 32 deletions(-) diff --git a/tcmalloc/huge_page_aware_allocator.h b/tcmalloc/huge_page_aware_allocator.h index 17ce30b61..4ae016188 100644 --- a/tcmalloc/huge_page_aware_allocator.h +++ b/tcmalloc/huge_page_aware_allocator.h @@ -106,8 +106,13 @@ class StaticForwarder { static void SetHugepage(HugePage p, void* pt); // SpanAllocator state. - static Span* NewSpan(Range r) ABSL_EXCLUSIVE_LOCKS_REQUIRED(pageheap_lock) - ABSL_ATTRIBUTE_RETURNS_NONNULL; + static Span* NewSpan(Range r) +#ifdef TCMALLOC_INTERNAL_LEGACY_LOCKING + ABSL_EXCLUSIVE_LOCKS_REQUIRED(pageheap_lock) +#else + ABSL_LOCKS_EXCLUDED(pageheap_lock) +#endif // TCMALLOC_INTERNAL_LEGACY_LOCKING + ABSL_ATTRIBUTE_RETURNS_NONNULL; static void DeleteSpan(Span* span) ABSL_EXCLUSIVE_LOCKS_REQUIRED(pageheap_lock) ABSL_ATTRIBUTE_NONNULL(); @@ -361,19 +366,30 @@ class HugePageAwareAllocator final : public PageAllocatorInterface { ABSL_EXCLUSIVE_LOCKS_REQUIRED(pageheap_lock); // Helpers for New(). - Span* LockAndAlloc(Length n, SpanAllocInfo span_alloc_info, - bool* from_released); +#ifdef TCMALLOC_INTERNAL_LEGACY_LOCKING + using FinalizeType = Span*; +#else // !TCMALLOC_INTERNAL_LEGACY_LOCKING + struct FinalizeType { + Range r; + bool donated = false; + }; +#endif // !TCMALLOC_INTERNAL_LEGACY_LOCKING + + FinalizeType LockAndAlloc(Length n, SpanAllocInfo span_alloc_info, + bool* from_released); - Span* AllocSmall(Length n, SpanAllocInfo span_alloc_info, bool* from_released) + FinalizeType AllocSmall(Length n, SpanAllocInfo span_alloc_info, + bool* from_released) ABSL_EXCLUSIVE_LOCKS_REQUIRED(pageheap_lock); - Span* AllocLarge(Length n, SpanAllocInfo span_alloc_info, bool* from_released) + FinalizeType AllocLarge(Length n, SpanAllocInfo span_alloc_info, + bool* from_released) ABSL_EXCLUSIVE_LOCKS_REQUIRED(pageheap_lock); - Span* AllocEnormous(Length n, SpanAllocInfo span_alloc_info, - bool* from_released) + FinalizeType AllocEnormous(Length n, SpanAllocInfo span_alloc_info, + bool* from_released) ABSL_EXCLUSIVE_LOCKS_REQUIRED(pageheap_lock); - Span* AllocRawHugepages(Length n, SpanAllocInfo span_alloc_info, - bool* from_released) + FinalizeType AllocRawHugepages(Length n, SpanAllocInfo span_alloc_info, + bool* from_released) ABSL_EXCLUSIVE_LOCKS_REQUIRED(pageheap_lock); bool AddRegion() ABSL_EXCLUSIVE_LOCKS_REQUIRED(pageheap_lock); @@ -385,7 +401,9 @@ class HugePageAwareAllocator final : public PageAllocatorInterface { ABSL_EXCLUSIVE_LOCKS_REQUIRED(pageheap_lock); // Finish an allocation request - give it a span and mark it in the pagemap. - Span* Finalize(Range r); + FinalizeType Finalize(Range r); + + Span* Spanify(FinalizeType f); // Whether this HPAA should use subrelease. This delegates to the appropriate // parameter depending whether this is for the cold heap or another heap. @@ -466,23 +484,30 @@ inline PageId HugePageAwareAllocator::RefillFiller( } template -inline Span* HugePageAwareAllocator::Finalize(Range r) +inline typename HugePageAwareAllocator::FinalizeType +HugePageAwareAllocator::Finalize(Range r) ABSL_EXCLUSIVE_LOCKS_REQUIRED(pageheap_lock) { TC_ASSERT_NE(r.p, PageId{0}); + info_.RecordAlloc(r); + forwarder_.ShrinkToUsageLimit(r.n); +#ifdef TCMALLOC_INTERNAL_LEGACY_LOCKING // TODO(b/175334169): Lift Span creation out of LockAndAlloc. Span* ret = forwarder_.NewSpan(r); forwarder_.Set(r.p, ret); TC_ASSERT(!ret->sampled()); - info_.RecordAlloc(r); - forwarder_.ShrinkToUsageLimit(r.n); return ret; +#else + return {r, false}; +#endif } // For anything <= half a huge page, we will unconditionally use the filler // to pack it into a single page. If we need another page, that's fine. template -inline Span* HugePageAwareAllocator::AllocSmall( - Length n, SpanAllocInfo span_alloc_info, bool* from_released) { +inline typename HugePageAwareAllocator::FinalizeType +HugePageAwareAllocator::AllocSmall(Length n, + SpanAllocInfo span_alloc_info, + bool* from_released) { auto [pt, page, released] = filler_.TryGet(n, span_alloc_info); *from_released = released; if (ABSL_PREDICT_TRUE(pt != nullptr)) { @@ -491,14 +516,16 @@ inline Span* HugePageAwareAllocator::AllocSmall( page = RefillFiller(n, span_alloc_info, from_released); if (ABSL_PREDICT_FALSE(page == PageId{0})) { - return nullptr; + return {}; } return Finalize(Range(page, n)); } template -inline Span* HugePageAwareAllocator::AllocLarge( - Length n, SpanAllocInfo span_alloc_info, bool* from_released) { +inline typename HugePageAwareAllocator::FinalizeType +HugePageAwareAllocator::AllocLarge(Length n, + SpanAllocInfo span_alloc_info, + bool* from_released) { // If it's an exact page multiple, just pull it from pages directly. HugeLength hl = HLFromPages(n); if (hl.in_pages() == n) { @@ -561,18 +588,21 @@ inline Span* HugePageAwareAllocator::AllocLarge( } template -inline Span* HugePageAwareAllocator::AllocEnormous( - Length n, SpanAllocInfo span_alloc_info, bool* from_released) { +inline typename HugePageAwareAllocator::FinalizeType +HugePageAwareAllocator::AllocEnormous(Length n, + SpanAllocInfo span_alloc_info, + bool* from_released) { return AllocRawHugepages(n, span_alloc_info, from_released); } template -inline Span* HugePageAwareAllocator::AllocRawHugepages( +inline typename HugePageAwareAllocator::FinalizeType +HugePageAwareAllocator::AllocRawHugepages( Length n, SpanAllocInfo span_alloc_info, bool* from_released) { HugeLength hl = HLFromPages(n); HugeRange r = cache_.Get(hl, from_released); - if (!r.valid()) return nullptr; + if (!r.valid()) return {}; // We now have a huge page range that covers our request. There // might be some slack in it if n isn't a multiple of @@ -593,9 +623,14 @@ inline Span* HugePageAwareAllocator::AllocRawHugepages( Length here = kPagesPerHugePage - slack; TC_ASSERT_GT(here, Length(0)); AllocAndContribute(last, here, span_alloc_info, /*donated=*/true); - Span* span = Finalize(Range(r.start().first_page(), n)); + auto span = Finalize(Range(r.start().first_page(), n)); +#ifdef TCMALLOC_INTERNAL_LEGACY_LOCKING span->set_donated(/*value=*/true); return span; +#else + span.donated = true; + return span; +#endif } inline static void BackSpan(Span* span) { @@ -608,7 +643,7 @@ inline Span* HugePageAwareAllocator::New( Length n, SpanAllocInfo span_alloc_info) { TC_CHECK_GT(n, Length(0)); bool from_released; - Span* s = LockAndAlloc(n, span_alloc_info, &from_released); + Span* s = Spanify(LockAndAlloc(n, span_alloc_info, &from_released)); if (s) { // Prefetch for writing, as we anticipate using the memory soon. PrefetchW(s->start_address()); @@ -619,8 +654,10 @@ inline Span* HugePageAwareAllocator::New( } template -inline Span* HugePageAwareAllocator::LockAndAlloc( - Length n, SpanAllocInfo span_alloc_info, bool* from_released) { +inline typename HugePageAwareAllocator::FinalizeType +HugePageAwareAllocator::LockAndAlloc(Length n, + SpanAllocInfo span_alloc_info, + bool* from_released) { PageHeapSpinLockHolder l; // Our policy depends on size. For small things, we will pack them // into single hugepages. @@ -651,16 +688,34 @@ inline Span* HugePageAwareAllocator::NewAligned( // TODO(b/134690769): support higher align. TC_CHECK_LE(align, kPagesPerHugePage); bool from_released; - Span* s; + FinalizeType f; { PageHeapSpinLockHolder l; - s = AllocRawHugepages(n, span_alloc_info, &from_released); + f = AllocRawHugepages(n, span_alloc_info, &from_released); } + Span* s = Spanify(f); if (s && from_released) BackSpan(s); TC_ASSERT(!s || GetMemoryTag(s->start_address()) == tag_); return s; } +template +inline Span* HugePageAwareAllocator::Spanify(FinalizeType f) { +#ifdef TCMALLOC_INTERNAL_LEGACY_LOCKING + return f; +#else + if (ABSL_PREDICT_FALSE(f.r.p == PageId{0})) { + return nullptr; + } + + Span* s = forwarder_.NewSpan(f.r); + forwarder_.Set(f.r.p, s); + TC_ASSERT(!s->sampled()); + s->set_donated(f.donated); + return s; +#endif +} + template inline void HugePageAwareAllocator::DeleteFromHugepage( FillerType::Tracker* pt, Range r, bool might_abandon) { diff --git a/tcmalloc/mock_huge_page_static_forwarder.h b/tcmalloc/mock_huge_page_static_forwarder.h index e3d4d77ac..d50011bb2 100644 --- a/tcmalloc/mock_huge_page_static_forwarder.h +++ b/tcmalloc/mock_huge_page_static_forwarder.h @@ -126,8 +126,11 @@ class FakeStaticForwarder { void SetHugepage(HugePage p, void* pt) { trackers_[p] = pt; } // SpanAllocator state. - [[nodiscard]] Span* NewSpan(Range r) ABSL_EXCLUSIVE_LOCKS_REQUIRED( - pageheap_lock) ABSL_ATTRIBUTE_RETURNS_NONNULL { + [[nodiscard]] Span* NewSpan(Range r) +#ifdef TCMALLOC_INTERNAL_LEGACY_LOCKING + ABSL_EXCLUSIVE_LOCKS_REQUIRED(pageheap_lock) +#endif // TCMALLOC_INTERNAL_LEGACY_LOCKING + ABSL_ATTRIBUTE_RETURNS_NONNULL { Span* span; void* result = absl::base_internal::LowLevelAlloc::AllocWithArena( sizeof(*span) + alignof(Span) + sizeof(void*), ll_arena()); diff --git a/tcmalloc/span.h b/tcmalloc/span.h index da913218c..7d3375a51 100644 --- a/tcmalloc/span.h +++ b/tcmalloc/span.h @@ -102,7 +102,11 @@ class Span final : public SpanList::Elem { // Allocator/deallocator for spans. Note that these functions are defined // in static_vars.h, which is weird: see there for why. - static Span* New(Range r); + static Span* New(Range r) +#ifndef TCMALLOC_INTERNAL_LEGACY_LOCKING + ABSL_LOCKS_EXCLUDED(pageheap_lock) +#endif + ; static void Delete(Span* span); static void operator delete(void*) = delete; diff --git a/tcmalloc/span_test.cc b/tcmalloc/span_test.cc index 1c02c1993..57e21b9cf 100644 --- a/tcmalloc/span_test.cc +++ b/tcmalloc/span_test.cc @@ -256,7 +256,9 @@ TEST(SpanAllocatorTest, Alignment) { spans.reserve(kNumSpans); { +#ifdef TCMALLOC_INTERNAL_LEGACY_LOCKING PageHeapSpinLockHolder l; +#endif // TCMALLOC_INTERNAL_LEGACY_LOCKING for (int i = 0; i < kNumSpans; ++i) { spans.push_back(Span::New(r)); }