Preserve broken hugepage state in PageMap across HugeCache/HugeAllocator reuse.

tcmalloc/BUILD

@@ -98,6 +98,7 @@ tcmalloc_deps = [
     "//tcmalloc/internal:logging",
     "//tcmalloc/internal:memory_tag",
     "//tcmalloc/internal:optimization",
+    "//tcmalloc/internal:packed_ptr",
     "//tcmalloc/internal:page_allocation_status",
     "//tcmalloc/internal:percpu",
     "//tcmalloc/internal:range_tracker",
@@ -331,6 +332,7 @@ create_tcmalloc_libraries(
         "//tcmalloc/internal:mincore",
         "//tcmalloc/internal:numa",
         "//tcmalloc/internal:optimization",
+        "//tcmalloc/internal:packed_ptr",
         "//tcmalloc/internal:page_allocation_status",
         "//tcmalloc/internal:page_allocator_hooks",
         "//tcmalloc/internal:page_size",
@@ -712,6 +714,7 @@ cc_library(
         "//tcmalloc/internal:config",
         "//tcmalloc/internal:logging",
         "//tcmalloc/internal:memory_tag",
+        "//tcmalloc/internal:packed_ptr",
         "//tcmalloc/internal:system_allocator",
         "@com_google_absl//absl/base",
         "@com_google_absl//absl/base:core_headers",

tcmalloc/CMakeLists.txt

@@ -76,6 +76,7 @@ tcmalloc_cc_library(
     "tcmalloc::internal_memory_tag"
     "tcmalloc::internal_optimization"
     "tcmalloc::internal_overflow"
+    "tcmalloc::internal_packed_ptr"
     "tcmalloc::internal_page_allocation_status"
     "tcmalloc::internal_page_size"
     "tcmalloc::internal_percpu"
@@ -127,6 +128,7 @@ tcmalloc_cc_library(
     "tcmalloc::internal_memory_tag"
     "tcmalloc::internal_optimization"
     "tcmalloc::internal_overflow"
+    "tcmalloc::internal_packed_ptr"
     "tcmalloc::internal_page_allocation_status"
     "tcmalloc::internal_page_size"
     "tcmalloc::internal_percpu"
@@ -178,6 +180,7 @@ tcmalloc_cc_library(
     "tcmalloc::internal_memory_tag"
     "tcmalloc::internal_optimization"
     "tcmalloc::internal_overflow"
+    "tcmalloc::internal_packed_ptr"
     "tcmalloc::internal_page_allocation_status"
     "tcmalloc::internal_page_size"
     "tcmalloc::internal_percpu"
@@ -351,6 +354,7 @@ tcmalloc_cc_library_variants(
     "tcmalloc::internal_mincore"
     "tcmalloc::internal_numa"
     "tcmalloc::internal_optimization"
+    "tcmalloc::internal_packed_ptr"
     "tcmalloc::internal_page_allocation_status"
     "tcmalloc::internal_page_allocator_hooks"
     "tcmalloc::internal_page_size"
@@ -485,6 +489,7 @@ tcmalloc_cc_library(
     "tcmalloc::internal_memory_tag"
     "tcmalloc::internal_optimization"
     "tcmalloc::internal_overflow"
+    "tcmalloc::internal_packed_ptr"
     "tcmalloc::internal_page_allocation_status"
     "tcmalloc::internal_page_size"
     "tcmalloc::internal_percpu"
@@ -536,6 +541,7 @@ tcmalloc_cc_library(
     "tcmalloc::internal_memory_tag"
     "tcmalloc::internal_optimization"
     "tcmalloc::internal_overflow"
+    "tcmalloc::internal_packed_ptr"
     "tcmalloc::internal_page_allocation_status"
     "tcmalloc::internal_page_size"
     "tcmalloc::internal_percpu"
@@ -588,6 +594,7 @@ tcmalloc_cc_library(
     "tcmalloc::internal_memory_tag"
     "tcmalloc::internal_optimization"
     "tcmalloc::internal_overflow"
+    "tcmalloc::internal_packed_ptr"
     "tcmalloc::internal_page_allocation_status"
     "tcmalloc::internal_page_size"
     "tcmalloc::internal_percpu"
@@ -639,6 +646,7 @@ tcmalloc_cc_library(
     "tcmalloc::internal_memory_tag"
     "tcmalloc::internal_optimization"
     "tcmalloc::internal_overflow"
+    "tcmalloc::internal_packed_ptr"
     "tcmalloc::internal_page_allocation_status"
     "tcmalloc::internal_page_size"
     "tcmalloc::internal_percpu"
@@ -690,6 +698,7 @@ tcmalloc_cc_library(
     "tcmalloc::internal_memory_tag"
     "tcmalloc::internal_optimization"
     "tcmalloc::internal_overflow"
+    "tcmalloc::internal_packed_ptr"
     "tcmalloc::internal_page_allocation_status"
     "tcmalloc::internal_page_size"
     "tcmalloc::internal_percpu"
@@ -743,6 +752,7 @@ tcmalloc_cc_library(
     "tcmalloc::internal_memory_tag"
     "tcmalloc::internal_optimization"
     "tcmalloc::internal_overflow"
+    "tcmalloc::internal_packed_ptr"
     "tcmalloc::internal_page_allocation_status"
     "tcmalloc::internal_page_size"
     "tcmalloc::internal_percpu"
@@ -869,6 +879,7 @@ tcmalloc_cc_library(
     "tcmalloc::internal_config"
     "tcmalloc::internal_logging"
     "tcmalloc::internal_memory_tag"
+    "tcmalloc::internal_packed_ptr"
     "tcmalloc::internal_system_allocator"
 )
 

tcmalloc/huge_page_aware_allocator.h

@@ -41,6 +41,7 @@
 #include "tcmalloc/internal/config.h"
 #include "tcmalloc/internal/logging.h"
 #include "tcmalloc/internal/metadata_allocator.h"
+#include "tcmalloc/internal/packed_ptr.h"
 #include "tcmalloc/internal/pageflags.h"
 #include "tcmalloc/internal/parameter_accessors.h"
 #include "tcmalloc/internal/prefetch.h"
@@ -332,9 +333,16 @@ class HugePageAwareAllocator final : public PageAllocatorInterface {
   MetadataObjectAllocator<HugeRegion, ArenaAlloc::kHugeRegion> region_allocator_
       ABSL_GUARDED_BY(pageheap_lock);
 
+  enum class HugePageSplitState : bool {
+    kUnbroken,
+    kBroken,
+  };
+
   [[nodiscard]] FillerType::Tracker* GetTracker(HugePage p) const;
+  [[nodiscard]] bool IsHugepageBroken(HugePage p) const;
 
   void SetTracker(HugePage p, FillerType::Tracker* pt);
+  void ClearTracker(HugePage p, HugePageSplitState split_state);
 
   [[nodiscard]] AddressRange AllocAndReport(size_t bytes, size_t align)
       ABSL_EXCLUSIVE_LOCKS_REQUIRED(pageheap_lock);
@@ -436,24 +444,45 @@ inline HugePageAwareAllocator<Forwarder>::HugePageAwareAllocator(
 template <class Forwarder>
 inline typename HugePageAwareAllocator<Forwarder>::FillerType::Tracker*
 HugePageAwareAllocator<Forwarder>::GetTracker(HugePage p) const {
-  void* v = forwarder_.GetHugepage(p);
-  FillerType::Tracker* pt = reinterpret_cast<FillerType::Tracker*>(v);
+  FillerType::Tracker* pt = forwarder_.GetHugepage(p).Get();
   TC_ASSERT(pt == nullptr || pt->location() == p);
   return pt;
 }
 
+template <class Forwarder>
+inline bool HugePageAwareAllocator<Forwarder>::IsHugepageBroken(
+    HugePage p) const {
+  return forwarder_.GetHugepage(p).Bits() != 0;
+}
+
 template <class Forwarder>
 inline void HugePageAwareAllocator<Forwarder>::SetTracker(
     HugePage p, HugePageAwareAllocator<Forwarder>::FillerType::Tracker* pt) {
-  forwarder_.SetHugepage(p, pt);
+  forwarder_.SetHugepage(p, PackedPtr<FillerType::Tracker>(pt));
+}
+
+template <class Forwarder>
+inline void HugePageAwareAllocator<Forwarder>::ClearTracker(
+    HugePage p, HugePageSplitState split_state) {
+  forwarder_.SetHugepage(
+      p,
+      PackedPtr<FillerType::Tracker>(
+          /*ptr=*/nullptr, split_state == HugePageSplitState::kBroken ? 1 : 0));
 }
 
 template <class Forwarder>
 inline PageId HugePageAwareAllocator<Forwarder>::AllocAndContribute(
     HugePage p, Length n, SpanAllocInfo span_alloc_info, bool donated) {
   TC_CHECK_NE(p.start_addr(), nullptr);
+  const bool was_broken = IsHugepageBroken(p);
   FillerType::Tracker* pt =
       tracker_allocator_.New(p, donated, forwarder_.clock().now());
+  if (was_broken &&
+      forwarder_.track_broken_hugepages() == TrackBrokenHugepages::kEnabled) {
+    pt->set_unbroken(/*status=*/false);
+  }
+  // TODO(b/572530574): When allocating from an unbacked range (*from_released),
+  // mark the unallocated tail pages [n, kPagesPerHugePage) unbacked in pt.
   TC_ASSERT_GE(pt->longest_free_range(), n);
   TC_ASSERT_EQ(pt->was_donated(), donated);
   // if the page was donated, we track its size so that we can potentially
@@ -839,8 +868,12 @@ inline void HugePageAwareAllocator<Forwarder>::ReleaseHugepage(
   }
 
   HugeRange r = {pt->location(), NHugePages(1)};
-  SetTracker(pt->location(), nullptr);
+  ClearTracker(pt->location(), pt->unbroken() ? HugePageSplitState::kUnbroken
+                                              : HugePageSplitState::kBroken);
 
+  // TODO(b/572530574): Avoid coalescing or preferring broken hugepages ahead of
+  // intact THPs in HugeCache, and evict broken cached pages first in
+  // ShrinkCache.
   if (pt->released()) {
     cache_.ReleaseUnbacked(r);
   } else {

tcmalloc/huge_page_aware_allocator_fuzz.cc

@@ -418,6 +418,17 @@ struct SetSubreleaseUnbackedHugepages {
   }
 };
 
+struct SetTrackBrokenHugepages {
+  bool value;
+
+  void Perform(State& state) const;
+
+  template <typename Sink>
+  friend void AbslStringify(Sink& sink, const SetTrackBrokenHugepages& s) {
+    absl::Format(&sink, "SetTrackBrokenHugepages{.value=%v}", s.value);
+  }
+};
+
 struct SetReleaseSucceeds {
   bool value;
 
@@ -590,10 +601,10 @@ using ParamOp = std::variant<
     SetFillerSkipSubreleaseLongInterval,
     SetFillerSkipSubreleaseColdShortInterval,
     SetFillerSkipSubreleaseColdLongInterval, SetReleasePartialAllocPages,
-    SetHpaaSubrelease, SetSubreleaseUnbackedHugepages, SetReleaseSucceeds,
-    SetCollapseSucceeds, SetHugeRegionAdaptiveRelease, SetAllocateSucceeds,
-    SetBackAllocations, SetBackSizeThresholdBytes, ReentrantSubprogram,
-    SetEnableUnfilteredCollapse, SetReleaseMaxFillerPages,
+    SetHpaaSubrelease, SetSubreleaseUnbackedHugepages, SetTrackBrokenHugepages,
+    SetReleaseSucceeds, SetCollapseSucceeds, SetHugeRegionAdaptiveRelease,
+    SetAllocateSucceeds, SetBackAllocations, SetBackSizeThresholdBytes,
+    ReentrantSubprogram, SetEnableUnfilteredCollapse, SetReleaseMaxFillerPages,
     SetEnableReleaseStalePages, SetMadvNoHugepageHugeRegions, UpdateBitmaps,
     SetUsageLimitPressure>;
 
@@ -1159,6 +1170,11 @@ void SetSubreleaseUnbackedHugepages::Perform(State& state) const {
             : SubreleaseUnbackedMode::kDisabled);
 }
 
+void SetTrackBrokenHugepages::Perform(State& state) const {
+  state.allocator.forwarder().set_track_broken_hugepages(
+      value ? TrackBrokenHugepages::kEnabled : TrackBrokenHugepages::kDisabled);
+}
+
 void SetReleaseSucceeds::Perform(State& state) const {
   state.allocator.forwarder().set_release_succeeds(value);
 }
@@ -1331,6 +1347,8 @@ fuzztest::Domain<ChangeParam> GetChangeParamDomain(int depth) {
       fuzztest::Map(
           [](SetSubreleaseUnbackedHugepages s) { return ChangeParam{s}; },
           fuzztest::Arbitrary<SetSubreleaseUnbackedHugepages>()),
+      fuzztest::Map([](SetTrackBrokenHugepages s) { return ChangeParam{s}; },
+                    fuzztest::Arbitrary<SetTrackBrokenHugepages>()),
       fuzztest::Map([](SetReleaseSucceeds s) { return ChangeParam{s}; },
                     fuzztest::Arbitrary<SetReleaseSucceeds>()),
       fuzztest::Map([](SetCollapseSucceeds s) { return ChangeParam{s}; },
@@ -1976,6 +1994,8 @@ TEST(HugePageAwareAllocatorTest, PrinterTest) {
             "SetCollapseSucceeds{.value=true}");
   EXPECT_EQ(absl::StrCat(SetSubreleaseUnbackedHugepages{.value = false}),
             "SetSubreleaseUnbackedHugepages{.value=false}");
+  EXPECT_EQ(absl::StrCat(SetTrackBrokenHugepages{.value = false}),
+            "SetTrackBrokenHugepages{.value=false}");
   EXPECT_EQ(absl::StrCat(SetUsageLimitPressure{.pages = 3}),
             "SetUsageLimitPressure{.pages=3}");
 }

tcmalloc/huge_page_aware_allocator_test.cc

@@ -311,6 +311,26 @@ class HugePageAwareAllocatorTest
     return allocator_->region().UseHugeRegionMoreOften();
   }
 
+  struct HugepageEntryState {
+    PageTracker* tracker = nullptr;
+    uint64_t bits = 0;
+    bool unbroken = false;
+    Length released_pages;
+  };
+
+  [[nodiscard]] HugepageEntryState GetHugepageState(HugePage hp) {
+    PageHeapSpinLockHolder l;
+    const auto packed = allocator_->forwarder().GetHugepage(hp);
+    HugepageEntryState state;
+    state.tracker = packed.Get();
+    state.bits = packed.Bits();
+    if (state.tracker != nullptr) {
+      state.unbroken = state.tracker->unbroken();
+      state.released_pages = state.tracker->released_pages();
+    }
+    return state;
+  }
+
   std::string Print() {
     return PrintToString(1 << 20, [&](Printer& p) {
       PageFlags pageflags;
@@ -858,7 +878,7 @@ TEST_P(HugePageAwareAllocatorTest, ParkedDonatedHugepageIsNotAbandoned) {
   ASSERT_TRUE(large);
   ASSERT_TRUE(large.donated);
   const HugePage hp = HugePageContaining(large.r.p);
-  auto* pt = static_cast<PageTracker*>(allocator_->forwarder().GetHugepage(hp));
+  PageTracker* pt = allocator_->forwarder().GetHugepage(hp).Get();
   ASSERT_NE(pt, nullptr);
 
   HugeLength donated_huge_pages;
@@ -2571,6 +2591,77 @@ TEST_P(HugePageAwareAllocatorTest, GetPageAllocationStatus) {
   }
 }
 
+TEST_P(HugePageAwareAllocatorTest, BrokenHugepageStatePersistsAcrossReuse) {
+  allocator_->forwarder().set_hpaa_subrelease(true);
+  allocator_->forwarder().set_release_partial_alloc_pages(true);
+
+  const SpanAllocInfo kSpanInfo = {1, AccessDensityPrediction::kSparse};
+  std::vector<AllocationState> keep_alive;
+
+  for (const TrackBrokenHugepages track_broken_hugepages :
+       {TrackBrokenHugepages::kDisabled, TrackBrokenHugepages::kEnabled}) {
+    allocator_->forwarder().set_track_broken_hugepages(track_broken_hugepages);
+
+    // 1. Subrelease free pages on `hp`, refill it so `released_pages() == 0`
+    //    while `!pt->unbroken()`, and drain it into `HugeCache`.
+    AllocationState a0 = New(Length(1), kSpanInfo);
+    ASSERT_TRUE(a0);
+    const HugePage hp = HugePageContaining(a0.r.p);
+    ReleasePages(kPagesPerHugePage - Length(1),
+                 PageReleaseReason::kProcessBackgroundActions);
+    HugepageEntryState state = GetHugepageState(hp);
+    ASSERT_NE(state.tracker, nullptr);
+    EXPECT_FALSE(state.unbroken);
+    EXPECT_EQ(state.released_pages, kPagesPerHugePage - Length(1));
+
+    AllocationState fill_rest = New(kPagesPerHugePage - Length(1), kSpanInfo);
+    ASSERT_TRUE(fill_rest);
+    ASSERT_EQ(HugePageContaining(fill_rest.r.p), hp);
+    state = GetHugepageState(hp);
+    ASSERT_NE(state.tracker, nullptr);
+    EXPECT_FALSE(state.unbroken);
+    EXPECT_EQ(state.released_pages, Length(0));
+
+    // Draining both spans returns `hp` to `HugeCache`; PageMap preserves the
+    // low bit, and reuse restores `!pt->unbroken()` iff
+    // `track_broken_hugepages` is enabled.
+    Delete(a0, 1);
+    Delete(fill_rest, 1);
+    state = GetHugepageState(hp);
+    EXPECT_EQ(state.tracker, nullptr);
+    EXPECT_EQ(state.bits, 1u);
+
+    AllocationState reused_from_cache = New(Length(1), kSpanInfo);
+    ASSERT_TRUE(reused_from_cache);
+    ASSERT_EQ(HugePageContaining(reused_from_cache.r.p), hp);
+    state = GetHugepageState(hp);
+    ASSERT_NE(state.tracker, nullptr);
+    EXPECT_EQ(state.unbroken,
+              track_broken_hugepages == TrackBrokenHugepages::kDisabled);
+
+    // 2. Subreleasing and freeing all allocations (`HandleFullyFreedTracker`)
+    //    returns `hp` to `HugeAllocator` unbacked; reusing `hp` restores
+    //    `!pt->unbroken()` iff `track_broken_hugepages` is enabled.
+    ReleasePages(kPagesPerHugePage - Length(1),
+                 PageReleaseReason::kProcessBackgroundActions);
+    Delete(reused_from_cache, 1);
+    AllocationState reused_from_allocator = New(Length(1), kSpanInfo);
+    ASSERT_TRUE(reused_from_allocator);
+    ASSERT_EQ(HugePageContaining(reused_from_allocator.r.p), hp);
+    state = GetHugepageState(hp);
+    ASSERT_NE(state.tracker, nullptr);
+    EXPECT_EQ(state.unbroken,
+              track_broken_hugepages == TrackBrokenHugepages::kDisabled);
+
+    keep_alive.push_back(reused_from_allocator);
+    keep_alive.push_back(New(kPagesPerHugePage - Length(1), kSpanInfo));
+  }
+
+  for (AllocationState& s : keep_alive) {
+    Delete(s, 1);
+  }
+}
+
 }  // namespace
 }  // namespace tcmalloc_internal
 }  // namespace tcmalloc

tcmalloc/huge_page_options.h

@@ -59,6 +59,11 @@ enum class MadviseRegionsNoHugepage : bool {
   kEnabled = true,
 };
 
+enum class TrackBrokenHugepages : bool {
+  kDisabled = false,
+  kEnabled = true,
+};
+
 }  // namespace tcmalloc::tcmalloc_internal
 GOOGLE_MALLOC_SECTION_END