From 43b94ea697a796be820f3b78f098ba1bf36b6329 Mon Sep 17 00:00:00 2001 From: Vaibhav Gogte Date: Thu, 9 Jul 2026 14:06:43 -0700 Subject: [PATCH] Refactor residency to return bitmaps in TCMalloc page granularity. * Update Residency::GetUnbackedAndSwappedBitmaps to take native_pages_per_element, allowing callers to request bitmaps in TCMalloc page granularity directly. * Remove bitmap conversion loop in PageTracker::MarkSubreleased. * Update call sites and telemetry recording in HugePageFiller to use element-granularity bitmaps. * Update mocks and test fakes to aggregate native bitmaps according to native_pages_per_element. PiperOrigin-RevId: 945306649 --- tcmalloc/huge_page_filler.h | 86 +++++++++-------------- tcmalloc/huge_page_filler_fuzz.cc | 34 ++++++++- tcmalloc/huge_page_filler_test.cc | 35 ++++++++-- tcmalloc/internal/BUILD | 1 + tcmalloc/internal/residency.cc | 32 ++++++--- tcmalloc/internal/residency.h | 10 ++- tcmalloc/internal/residency_test.cc | 105 +++++++++++++++------------- 7 files changed, 182 insertions(+), 121 deletions(-) diff --git a/tcmalloc/huge_page_filler.h b/tcmalloc/huge_page_filler.h index ffab64510..75d03d983 100644 --- a/tcmalloc/huge_page_filler.h +++ b/tcmalloc/huge_page_filler.h @@ -218,8 +218,7 @@ class PageTracker : public TList::Elem { Length ReleaseFree(MemoryModifyFunction& unback) ABSL_EXCLUSIVE_LOCKS_REQUIRED(pageheap_lock); - Length MarkSubreleased(Bitmap unbacked, - int native_pages_per_tcmalloc_page) + Length MarkSubreleased(Bitmap unbacked) ABSL_EXCLUSIVE_LOCKS_REQUIRED(pageheap_lock); Bitmap released_by_page() const { @@ -702,11 +701,15 @@ class UsageInfo { } else { records.num_free_non_hugepage_backed += (free - pt.released_pages()); records.num_used_non_hugepage_backed += pt.used_pages(); + const size_t native_pages_per_tcmalloc_page = + kPagesPerHugePage.raw_num() > 0 + ? (kHugePageSize / GetPageSize() / kPagesPerHugePage.raw_num()) + : 0; Residency::SinglePageBitmaps bitmaps = hugepage_residency_state.bitmaps; ++records.unbacked_histo[NativePageBucketNum( - bitmaps.unbacked.CountBits())]; - ++records - .swapped_histo[NativePageBucketNum(bitmaps.swapped.CountBits())]; + bitmaps.unbacked.CountBits() * native_pages_per_tcmalloc_page)]; + ++records.swapped_histo[NativePageBucketNum( + bitmaps.swapped.CountBits() * native_pages_per_tcmalloc_page)]; ++records.stale_histo[NativePageBucketNum( hugepage_residency_state.stale.CountBits())]; @@ -1313,6 +1316,10 @@ class HugePageFiller { Bitmap unbacked) ABSL_EXCLUSIVE_LOCKS_REQUIRED(pageheap_lock); + int native_pages_per_tcmalloc_page() const { + return native_pages_per_tcmalloc_page_; + } + private: // This class wraps an array of N TrackerLists and a Bitmap storing which // elements are non-empty. @@ -1536,8 +1543,8 @@ inline PageTracker::NativePageResidencyInfo PageTracker::CountInfoInHugePage( const int scaled_idx = i >> shift_bits; bool is_free = !free.GetBit(scaled_idx); - n_swapped[is_free] += swapped.GetBit(i); - n_unbacked[is_free] += unbacked.GetBit(i); + n_swapped[is_free] += swapped.GetBit(scaled_idx); + n_unbacked[is_free] += unbacked.GetBit(scaled_idx); n_stale[is_free] += stale.GetBit(i); } @@ -1607,49 +1614,19 @@ inline Length PageTracker::ReleaseFree(MemoryModifyFunction& unback) { return Length(count); } -inline Length PageTracker::MarkSubreleased(Bitmap unbacked, - int native_pages_per_tcmalloc_page) { - // TODO(b/525422238): pass the bitmap in TCMalloc pages directly to avoid this - // conversion. - - // When TCMalloc page size is smaller than the native page size, mark the page - // as broken. We should additionally mark unbacked pages as subreleased. The - // later case isn't done currently. - if (native_pages_per_tcmalloc_page == 0) { +inline Length PageTracker::MarkSubreleased(Bitmap unbacked) { + const size_t kNativePagesInHugePage = kHugePageSize / GetPageSize(); + if (kNativePagesInHugePage < kPagesPerHugePage.raw_num()) { unbroken_ = false; return Length(0); } - TC_ASSERT_GT(native_pages_per_tcmalloc_page, 0); Bitmap free = free_.bits(); size_t count = 0; for (size_t i = 0; i < kPagesPerHugePage.raw_num(); ++i) { - if (!free.GetBit(i) && !released_by_page_.GetBit(i)) { - // Check if all native pages corresponding to this TCMalloc page are - // unbacked. - bool is_unbacked = true; - for (int j = 0; j < native_pages_per_tcmalloc_page; ++j) { - const int bitmap_index = i * native_pages_per_tcmalloc_page + j; - TC_ASSERT_LT(bitmap_index, kMaxResidencyBits); - if (!unbacked.GetBit(bitmap_index)) { - is_unbacked = false; - break; - } - } - // If all native pages corresponding to this TCMalloc page are unbacked, - // and as the TCMalloc page was freed but not released, mark this page - // as released now. - // - // TODO(b/525422238): The residency bitmap was captured outside of the - // lock. So, in a rare case, it's possible that the page was allocated, - // backed and then freed. So, the free page here is actually backed. - // While we currently ignore this case (resulting in underestimating - // RSS), we can potentially fix this by re-investigating the bitmaps - // and marking the pages back to backed to eventually fix this. - if (is_unbacked) { - released_by_page_.SetBit(i); - ++count; - } + if (!free.GetBit(i) && !released_by_page_.GetBit(i) && unbacked.GetBit(i)) { + released_by_page_.SetBit(i); + ++count; } } @@ -2711,8 +2688,9 @@ class HugePageUnbackedTrackerTreatment final : public HugePageTreatment { // information. state.maybe_hugepage_backed = is_hugepage; if (!is_hugepage) { - state.bitmaps = - res->GetUnbackedAndSwappedBitmaps(tracker->location().start_addr()); + state.bitmaps = res->GetUnbackedAndSwappedBitmaps( + tracker->location().start_addr(), + page_filler_.native_pages_per_tcmalloc_page()); if (pf) { pf->GetSinglePageBitmaps(tracker->location().start_addr(), state.stale); @@ -2721,12 +2699,15 @@ class HugePageUnbackedTrackerTreatment final : public HugePageTreatment { const bool backoff = treatment_stats_.collapse_time_max_cycles > max_collapse_cycles; if (enable_collapse_ == EnableCollapse::kEnabled && !backoff) { - bool should_collapse = enable_unfiltered_collapse_ == - EnableUnfilteredCollapse::kEnabled || - (state.bitmaps.swapped.CountBits() < - kMaxSwappedPagesForCollapse && - state.bitmaps.unbacked.CountBits() < - kMaxUnbackedPagesForCollapse); + bool should_collapse = + enable_unfiltered_collapse_ == + EnableUnfilteredCollapse::kEnabled || + (state.bitmaps.swapped.CountBits() * + page_filler_.native_pages_per_tcmalloc_page() < + kMaxSwappedPagesForCollapse && + state.bitmaps.unbacked.CountBits() * + page_filler_.native_pages_per_tcmalloc_page() < + kMaxUnbackedPagesForCollapse); if (should_collapse) { state.maybe_hugepage_backed = TryUserspaceCollapse(tracker); } else { @@ -2991,8 +2972,7 @@ template inline Length HugePageFiller::HandleUnbackedHugePage( PageTracker* tracker, Bitmap unbacked) { RemoveFromFillerList(tracker); - Length unmapped_length = - tracker->MarkSubreleased(unbacked, native_pages_per_tcmalloc_page_); + Length unmapped_length = tracker->MarkSubreleased(unbacked); subrelease_stats_.total_pages_subreleased += unmapped_length; unmapped_ += unmapped_length; unmapping_unaccounted_ += unmapped_length; diff --git a/tcmalloc/huge_page_filler_fuzz.cc b/tcmalloc/huge_page_filler_fuzz.cc index f23dac5dd..9215e2441 100644 --- a/tcmalloc/huge_page_filler_fuzz.cc +++ b/tcmalloc/huge_page_filler_fuzz.cc @@ -141,11 +141,39 @@ class FakeResidency : public Residency { return std::nullopt; }; - SinglePageBitmaps GetUnbackedAndSwappedBitmaps(const void* addr) override { - return {unbacked_bitmap, swapped_bitmap, absl::StatusCode::kOk}; + SinglePageBitmaps GetUnbackedAndSwappedBitmaps( + const void* addr, size_t native_pages_per_element) override { + SinglePageBitmaps native_res = {unbacked_bitmap, swapped_bitmap, + absl::StatusCode::kOk}; + if (native_pages_per_element <= 1) { + return native_res; + } + SinglePageBitmaps result; + result.status = native_res.status; + size_t num_elements = kNativePagesInHugePage / native_pages_per_element; + for (size_t idx = 0; idx < num_elements; ++idx) { + bool all_unbacked = true; + bool any_swapped = false; + for (size_t j = 0; j < native_pages_per_element; ++j) { + size_t native_idx = idx * native_pages_per_element + j; + if (!native_res.unbacked.GetBit(native_idx)) { + all_unbacked = false; + } + if (native_res.swapped.GetBit(native_idx)) { + any_swapped = true; + } + } + if (all_unbacked) { + result.unbacked.SetBit(idx); + } + if (any_swapped) { + result.swapped.SetBit(idx); + } + } + return result; }; - const size_t kNativePagesInHugePage = kHugePageSize / kPageSize; + const size_t kNativePagesInHugePage = kHugePageSize / GetPageSize(); size_t GetNativePagesInHugePage() const override { return kNativePagesInHugePage; }; diff --git a/tcmalloc/huge_page_filler_test.cc b/tcmalloc/huge_page_filler_test.cc index cb6cfbec3..08fd221a8 100644 --- a/tcmalloc/huge_page_filler_test.cc +++ b/tcmalloc/huge_page_filler_test.cc @@ -328,11 +328,38 @@ class FakeResidency : public Residency { return std::nullopt; }; - SinglePageBitmaps GetUnbackedAndSwappedBitmaps(const void* addr) override { + SinglePageBitmaps GetUnbackedAndSwappedBitmaps( + const void* addr, size_t native_pages_per_element) override { PageId p = PageIdContaining(addr); HugePage hp = HugePageContaining(p); EXPECT_TRUE(residency_bitmaps_.contains(hp.start_addr())); - return residency_bitmaps_[hp.start_addr()]; + SinglePageBitmaps native_res = residency_bitmaps_[hp.start_addr()]; + if (native_pages_per_element <= 1) { + return native_res; + } + SinglePageBitmaps result; + result.status = native_res.status; + size_t num_elements = kNativePagesInHugePage / native_pages_per_element; + for (size_t idx = 0; idx < num_elements; ++idx) { + bool all_unbacked = true; + bool any_swapped = false; + for (size_t j = 0; j < native_pages_per_element; ++j) { + size_t native_idx = idx * native_pages_per_element + j; + if (!native_res.unbacked.GetBit(native_idx)) { + all_unbacked = false; + } + if (native_res.swapped.GetBit(native_idx)) { + any_swapped = true; + } + } + if (all_unbacked) { + result.unbacked.SetBit(idx); + } + if (any_swapped) { + result.swapped.SetBit(idx); + } + } + return result; }; void SetUnbackedAndSwappedBitmaps(const void* addr, @@ -344,7 +371,7 @@ class FakeResidency : public Residency { absl::StatusCode::kOk}; } - const size_t kNativePagesInHugePage = kHugePageSize / kPageSize; + const size_t kNativePagesInHugePage = kHugePageSize / GetPageSize(); size_t GetNativePagesInHugePage() const override { return kNativePagesInHugePage; }; @@ -2462,7 +2489,7 @@ TEST_F(FillerTestWithSubreleaseUnbacked, SubreleaseUnbackedPartial) { // free TCMalloc page are unbacked, the TCMalloc page is NOT marked as // subreleased. TEST_F(FillerTestWithSubreleaseUnbacked, SubreleaseUnbackedPartialNativePages) { - const size_t kNativePagesInHugePage = kHugePageSize / kPageSize; + const size_t kNativePagesInHugePage = kHugePageSize / GetPageSize(); const int shift = kNativePagesInHugePage / kPagesPerHugePage.raw_num(); if (shift <= 1) { return; diff --git a/tcmalloc/internal/BUILD b/tcmalloc/internal/BUILD index e1c1c7dbb..aab03c43c 100644 --- a/tcmalloc/internal/BUILD +++ b/tcmalloc/internal/BUILD @@ -601,6 +601,7 @@ cc_test( ":range_tracker", ":residency", ":util", + "//tcmalloc:common_8k_pages", "@com_github_google_benchmark//:benchmark", "@com_google_absl//absl/log:check", "@com_google_absl//absl/status", diff --git a/tcmalloc/internal/residency.cc b/tcmalloc/internal/residency.cc index 7b92740e9..4135fc276 100644 --- a/tcmalloc/internal/residency.cc +++ b/tcmalloc/internal/residency.cc @@ -187,7 +187,7 @@ std::optional ResidencyPageMap::Get(const void* const addr, } Residency::SinglePageBitmaps ResidencyPageMap::GetUnbackedAndSwappedBitmaps( - const void* const addr) { + const void* const addr, size_t native_pages_per_element) { Bitmap page_unbacked; Bitmap page_swapped; uintptr_t currPage = reinterpret_cast(addr); @@ -212,16 +212,30 @@ Residency::SinglePageBitmaps ResidencyPageMap::GetUnbackedAndSwappedBitmaps( absl::StatusCode::kUnavailable}; } - for (int native_page_idx = 0; native_page_idx < kNativePagesInHugePage; - ++native_page_idx) { - uint64_t page_map = buf_[native_page_idx]; - if (!PagePresent(page_map)) { - if (PageSwapped(page_map)) { - page_swapped.SetBit(native_page_idx); - } else { - page_unbacked.SetBit(native_page_idx); + if (native_pages_per_element == 0) { + return SinglePageBitmaps{page_unbacked, page_swapped, + absl::StatusCode::kOk}; + } + const size_t num_elements = kNativePagesInHugePage / native_pages_per_element; + + for (size_t idx = 0; idx < num_elements; ++idx) { + bool all_unbacked = true; + bool any_swapped = false; + for (size_t j = 0; j < native_pages_per_element; ++j) { + uint64_t page_map = buf_[idx * native_pages_per_element + j]; + if (PagePresent(page_map)) { + all_unbacked = false; + } else if (PageSwapped(page_map)) { + all_unbacked = false; + any_swapped = true; } } + if (all_unbacked) { + page_unbacked.SetBit(idx); + } + if (any_swapped) { + page_swapped.SetBit(idx); + } } return SinglePageBitmaps{page_unbacked, page_swapped, absl::StatusCode::kOk}; } diff --git a/tcmalloc/internal/residency.h b/tcmalloc/internal/residency.h index 440008405..b6bcb77d9 100644 --- a/tcmalloc/internal/residency.h +++ b/tcmalloc/internal/residency.h @@ -47,6 +47,7 @@ class Residency { virtual inline size_t GetNativePagesInHugePage() const = 0; // Struct is ordered with bitmaps first to optimize cacheline usage. + // The bitmaps are indexed by TCMalloc pages (0 .. kPagesPerHugePage - 1). struct SinglePageBitmaps { Bitmap unbacked; Bitmap swapped; @@ -57,8 +58,10 @@ class Residency { // to output two bitmaps - one for pages that are unbacked and one for pages // that are swapped. Hugepage-sized regions are assumed to be 2MiB in size. A // SinglePageBitmaps struct is returned with the status, the page_unbacked - // bitmap, and the page_swapped bitmap. - virtual SinglePageBitmaps GetUnbackedAndSwappedBitmaps(const void* addr) = 0; + // bitmap, and the page_swapped bitmap. Each bit represents + // native_pages_per_element native OS pages. + virtual SinglePageBitmaps GetUnbackedAndSwappedBitmaps( + const void* addr, size_t native_pages_per_element) = 0; }; // Residency offers information about memory residency: whether or not specific @@ -95,7 +98,8 @@ class ResidencyPageMap : public Residency { // that are swapped. Hugepage-sized regions are assumed to be 2MiB in size. A // SinglePageBitmaps struct is returned with the status, the page_unbacked // bitmap, and the page_swapped bitmap. - SinglePageBitmaps GetUnbackedAndSwappedBitmaps(const void* addr) override; + SinglePageBitmaps GetUnbackedAndSwappedBitmaps( + const void* addr, size_t native_pages_per_element) override; private: // This helper seeks the internal file to the correct location for the given diff --git a/tcmalloc/internal/residency_test.cc b/tcmalloc/internal/residency_test.cc index c93c2d906..d7803e8ac 100644 --- a/tcmalloc/internal/residency_test.cc +++ b/tcmalloc/internal/residency_test.cc @@ -35,6 +35,7 @@ #include "absl/status/status.h" #include "absl/strings/str_cat.h" #include "absl/strings/string_view.h" +#include "tcmalloc/huge_pages.h" #include "tcmalloc/internal/allocation_guard.h" #include "tcmalloc/internal/config.h" #include "tcmalloc/internal/page_size.h" @@ -55,8 +56,9 @@ class ResidencySpouse { return r_.Get(std::forward(args)...); } - decltype(auto) GetUnbackedAndSwappedBitmaps(const void* const addr) { - return r_.GetUnbackedAndSwappedBitmaps(addr); + decltype(auto) GetUnbackedAndSwappedBitmaps( + const void* const addr, size_t native_pages_per_element = 1) { + return r_.GetUnbackedAndSwappedBitmaps(addr, native_pages_per_element); } private: @@ -208,51 +210,49 @@ void GenerateHolesInSinglePage(absl::string_view filename, int case_num, CHECK_EQ(close(write_fd), 0) << errno; } -Residency::SinglePageBitmaps GenerateExpectedSinglePageBitmaps(int case_num) { +Residency::SinglePageBitmaps GenerateExpectedSinglePageBitmaps( + int case_num, size_t native_pages_per_element = 1) { Bitmap<512> expected_unbacked; Bitmap<512> expected_swapped; - switch (case_num) { - case 0: - // All Pages are present. Both bitmaps are 0 - break; - case 1: - // All Pages are swapped. Swapped bitmap is all 1 - expected_swapped.SetRange(0, 512); - break; - case 2: - // All pages are holes. Unbacked bitmap is all 1 - expected_unbacked.SetRange(0, 512); - break; - case 3: - // Every other page is a hole, rest are present. - // Both bitmaps are 0 and 1 alternating - for (int idx = 0; idx < 512; idx += 2) { - if (idx % 2 == 0) { - expected_unbacked.SetBit(idx); - } - } - break; - case 4: - // Every other page is swapped, rest are present, - // Swapped bitmap are 0 and 1 alternating - for (int idx = 0; idx < 512; idx += 2) { - if (idx % 2 == 0) { - expected_swapped.SetBit(idx); - } - } - break; - case 5: - // Every other page is swapped, rest are holes, - // Swapped bitmaps are 0 and 1 alternating - // Unbacked bitmaps are 1 and 0 alternating - for (int idx = 0; idx < 512; idx++) { - if (idx % 2 == 0) { - expected_swapped.SetBit(idx); - } else { - expected_unbacked.SetBit(idx); - } + const size_t kNativePagesInHugePage = kHugePageSize / GetPageSize(); + const size_t num_elements = kNativePagesInHugePage / native_pages_per_element; + + for (size_t i = 0; i < num_elements; ++i) { + bool all_unbacked = true; + bool any_swapped = false; + for (size_t j = 0; j < native_pages_per_element; ++j) { + size_t idx = i * native_pages_per_element + j; + switch (case_num) { + case 0: + all_unbacked = false; + break; + case 1: + all_unbacked = false; + any_swapped = true; + break; + case 2: + break; + case 3: + if (idx % 2 == 1) all_unbacked = false; + break; + case 4: + if (idx % 2 == 0) { + all_unbacked = false; + any_swapped = true; + } else { + all_unbacked = false; + } + break; + case 5: + if (idx % 2 == 0) { + all_unbacked = false; + any_swapped = true; + } + break; } - break; + } + if (all_unbacked) expected_unbacked.SetBit(i); + if (any_swapped) expected_swapped.SetBit(i); } return Residency::SinglePageBitmaps{expected_unbacked, expected_swapped, absl::StatusCode::kOk}; @@ -344,20 +344,27 @@ TEST(PageMapIntegrationTest, WorksOnActualData) { position++; addr = reinterpret_cast(position); } + const size_t native_pages_per_tcmalloc = + (kHugePageSize / GetPageSize()) / kPagesPerHugePage.raw_num(); g.emplace(); ResidencyPageMap r; - Residency::SinglePageBitmaps res = r.GetUnbackedAndSwappedBitmaps(addr); + Residency::SinglePageBitmaps res = + r.GetUnbackedAndSwappedBitmaps(addr, native_pages_per_tcmalloc); g.reset(); ASSERT_EQ(res.status, absl::StatusCode::kOk); EXPECT_TRUE(res.unbacked.IsZero()); - EXPECT_TRUE(res.swapped.IsZero()); - ASSERT_EQ(munmap(reinterpret_cast(addr) + 1 * 4096, 4096), 0) + const size_t tcmalloc_page_size = GetPageSize() * native_pages_per_tcmalloc; + ASSERT_EQ(munmap(reinterpret_cast(addr) + 1 * tcmalloc_page_size, + tcmalloc_page_size), + 0) << errno; - ASSERT_EQ(munmap(reinterpret_cast(addr) + 17 * 4096, 4096), 0) + ASSERT_EQ(munmap(reinterpret_cast(addr) + 17 * tcmalloc_page_size, + tcmalloc_page_size), + 0) << errno; g.emplace(); - res = r.GetUnbackedAndSwappedBitmaps(addr); + res = r.GetUnbackedAndSwappedBitmaps(addr, native_pages_per_tcmalloc); g.reset(); ASSERT_EQ(res.status, absl::StatusCode::kOk); EXPECT_FALSE(res.unbacked.IsZero());