From 6b1116c6a737797447371ace3f014b0992e87778 Mon Sep 17 00:00:00 2001 From: Ben Reisner Date: Mon, 13 Jul 2026 09:46:18 -0700 Subject: [PATCH] Introduce new functions to return pageflags/residency bitmaps in terms of TCMalloc pages instead of native pages. PiperOrigin-RevId: 947092650 --- tcmalloc/common.h | 2 + tcmalloc/huge_page_filler.h | 65 +++++++----------- tcmalloc/internal/range_tracker.h | 33 ++++++++++ tcmalloc/internal/range_tracker_test.cc | 88 +++++++++++++++++++++++++ 4 files changed, 147 insertions(+), 41 deletions(-) diff --git a/tcmalloc/common.h b/tcmalloc/common.h index 4e7165c73..f7fae9bf5 100644 --- a/tcmalloc/common.h +++ b/tcmalloc/common.h @@ -35,6 +35,7 @@ #include "tcmalloc/internal/logging.h" #include "tcmalloc/internal/memory_tag.h" #include "tcmalloc/internal/optimization.h" +#include "tcmalloc/internal/range_tracker.h" #include "tcmalloc/malloc_extension.h" GOOGLE_MALLOC_SECTION_BEGIN @@ -208,6 +209,7 @@ inline constexpr size_t kMinObjectsToMove = 2; inline constexpr size_t kMaxObjectsToMove = 128; inline constexpr size_t kPageSize = 1 << kPageShift; +using PageBitmap = Bitmap; inline constexpr std::align_val_t kAlignment{8}; // log2 (kAlignment) diff --git a/tcmalloc/huge_page_filler.h b/tcmalloc/huge_page_filler.h index 853705e98..da9cc5710 100644 --- a/tcmalloc/huge_page_filler.h +++ b/tcmalloc/huge_page_filler.h @@ -218,7 +218,7 @@ class PageTracker : public TList::Elem { Length ReleaseFree(MemoryModifyFunction& unback) ABSL_EXCLUSIVE_LOCKS_REQUIRED(pageheap_lock); - Length MarkSubreleased(Bitmap unbacked, + Length MarkSubreleased(PageBitmap unbacked, int native_pages_per_tcmalloc_page) ABSL_EXCLUSIVE_LOCKS_REQUIRED(pageheap_lock); @@ -250,6 +250,9 @@ class PageTracker : public TList::Elem { bool being_collapsed = false; // Records swap and unbacked bitmaps for this hugepage. Residency::SinglePageBitmaps bitmaps; + // Records the unbacked bitmap for this hugepage. In terms of TCMalloc + // pages. scaled via `ReductionOp::kAll`. + PageBitmap unbacked; ResidencyBitmap stale; // Records whether collapse was skipped due to threshold constraints. bool collapse_skipped = false; @@ -1311,8 +1314,7 @@ class HugePageFiller { // Utility function to handle a non-hugepage backed `page_tracker` and // mark its unmapped pages appropriately. - Length HandleUnbackedHugePage(PageTracker* page_tracker, - Bitmap unbacked) + Length HandleUnbackedHugePage(PageTracker* page_tracker, PageBitmap unbacked) ABSL_EXCLUSIVE_LOCKS_REQUIRED(pageheap_lock); private: @@ -1609,7 +1611,7 @@ inline Length PageTracker::ReleaseFree(MemoryModifyFunction& unback) { return Length(count); } -inline Length PageTracker::MarkSubreleased(Bitmap unbacked, +inline Length PageTracker::MarkSubreleased(PageBitmap unbacked, int native_pages_per_tcmalloc_page) { // TODO(b/525422238): pass the bitmap in TCMalloc pages directly to avoid this // conversion. @@ -1622,47 +1624,25 @@ inline Length PageTracker::MarkSubreleased(Bitmap unbacked, 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; - } - } - } - - released_count_ += count; + PageBitmap free = free_.bits(); + + // 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. + auto to_release = (~free) & (~released_by_page_) & unbacked; + released_by_page_ = released_by_page_ | to_release; + + released_count_ += to_release.CountBits(); // Mark this is unbroken regardless of whether it had any unbacked free // TCMalloc pages. Marking this will move this tracker to one of the // released lists. unbroken_ = false; TC_ASSERT_LE(Length(released_count_), kPagesPerHugePage); TC_ASSERT_EQ(released_by_page_.CountBits(), released_count_); - return Length(count); + return Length(to_release.CountBits()); } inline MemoryModifyStatus PageTracker::Collapse( @@ -2715,6 +2695,9 @@ class HugePageUnbackedTrackerTreatment final : public HugePageTreatment { if (!is_hugepage) { state.bitmaps = res->GetUnbackedAndSwappedBitmaps(tracker->location().start_addr()); + state.unbacked = Scale( + state.bitmaps.unbacked, kHugePageSize / GetPageSize(), + ReductionOp::kAll); if (pf) { pf->GetSinglePageBitmaps(tracker->location().start_addr(), state.stale); @@ -2776,7 +2759,7 @@ class HugePageUnbackedTrackerTreatment final : public HugePageTreatment { if (subrelease_unbacked_mode_ == SubreleaseUnbackedMode::kEnabled) { Length released_length = page_filler_.HandleUnbackedHugePage( - tracker, residency_states_[i].tracker_state.bitmaps.unbacked); + tracker, residency_states_[i].tracker_state.unbacked); if (released_length > Length(0)) { treatment_stats_.treated_pages_unbacked_subreleased += released_length.raw_num(); @@ -2991,7 +2974,7 @@ inline void HugePageFiller::OnCollapseSuccess(TrackerType* pt) { template inline Length HugePageFiller::HandleUnbackedHugePage( - PageTracker* tracker, Bitmap unbacked) { + PageTracker* tracker, PageBitmap unbacked) { RemoveFromFillerList(tracker); Length unmapped_length = tracker->MarkSubreleased(unbacked, native_pages_per_tcmalloc_page_); diff --git a/tcmalloc/internal/range_tracker.h b/tcmalloc/internal/range_tracker.h index e7958af23..8df4f9b05 100644 --- a/tcmalloc/internal/range_tracker.h +++ b/tcmalloc/internal/range_tracker.h @@ -118,6 +118,39 @@ class Bitmap { ssize_t FindValueBackwards(size_t index) const; }; +enum class ReductionOp { + kAny, + kAll, +}; + +template +Bitmap Scale(const Bitmap& src, size_t src_len, ReductionOp op) { + TC_ASSERT_LE(src_len, N); + TC_ASSERT(src_len % M == 0 || M % src_len == 0); + + Bitmap res; + const size_t src_per_dst = std::max(1, src_len / M); + const size_t dst_per_src = std::max(1, M / src_len); + + const size_t required_count = (op == ReductionOp::kAll) ? src_per_dst : 1; + + size_t dst_idx = 0; + for (size_t src_idx = 0; src_idx < src_len; src_idx += src_per_dst) { + size_t count = 0; + for (size_t j = 0; j < src_per_dst; ++j) { + if (src.GetBit(src_idx + j)) { + count++; + if (count == required_count) { + res.SetRange(dst_idx, dst_per_src); + break; + } + } + } + dst_idx += dst_per_src; + } + return res; +} + constexpr size_t kMaxResidencyBits = 512; using ResidencyBitmap = Bitmap; // Tracks allocations in a range of items of fixed size. Supports diff --git a/tcmalloc/internal/range_tracker_test.cc b/tcmalloc/internal/range_tracker_test.cc index 521e1d9c9..ad58cbdcd 100644 --- a/tcmalloc/internal/range_tracker_test.cc +++ b/tcmalloc/internal/range_tracker_test.cc @@ -347,6 +347,94 @@ TEST_F(BitmapTest, PopBatch) { EXPECT_TRUE(map.IsZero()); } +#ifndef NDEBUG +TEST_F(BitmapTest, ScaleAssertionFailures) { + Bitmap<64> map1; + // 1. src_len (64) doesn't divide M (10) and vice versa + EXPECT_DEATH(Scale<10>(map1, 64, ReductionOp::kAll), ""); + // 2. src_len > N (64) + EXPECT_DEATH(Scale<256>(map1, 128, ReductionOp::kAll), ""); +} +#endif + +TEST_F(BitmapTest, Scale) { + // Test case 1: N == M (identity) + { + Bitmap<64> map1; + map1.SetBit(3); + map1.SetBit(17); + map1.SetBit(63); + + auto map_all = Scale<64>(map1, 64, ReductionOp::kAll); + auto map_any = Scale<64>(map1, 64, ReductionOp::kAny); + for (size_t i = 0; i < 64; ++i) { + EXPECT_EQ(map1.GetBit(i), map_all.GetBit(i)); + EXPECT_EQ(map1.GetBit(i), map_any.GetBit(i)); + } + } + + // Test case 2: N > M (contracting) + { + Bitmap<64> map1; + // We group by 2 bits. + // Group 0: bits 0, 1 -> both must be set + map1.SetBit(0); + map1.SetBit(1); + + // Group 1: bits 2, 3 -> only one set + // Should not be set for kAll, should be set for kAny + map1.SetBit(2); + + // Group 7: bits 14, 15 -> both set + map1.SetBit(14); + map1.SetBit(15); + + auto map_all = Scale<8>(map1, 16, ReductionOp::kAll); + auto map_any = Scale<8>(map1, 16, ReductionOp::kAny); + + EXPECT_TRUE(map_all.GetBit(0)); + EXPECT_TRUE(map_any.GetBit(0)); + + EXPECT_FALSE(map_all.GetBit(1)); + EXPECT_TRUE(map_any.GetBit(1)); + + for (size_t i = 2; i <= 6; ++i) { + EXPECT_FALSE(map_all.GetBit(i)); + EXPECT_FALSE(map_any.GetBit(i)); + } + + EXPECT_TRUE(map_all.GetBit(7)); + EXPECT_TRUE(map_any.GetBit(7)); + } + + // Test case 3: N < M (expanding) + { + Bitmap<64> map1; + // We expand by 2 bits. + map1.SetBit(0); + map1.SetBit(7); + + auto map_all = Scale<16>(map1, 8, ReductionOp::kAll); + auto map_any = Scale<16>(map1, 8, ReductionOp::kAny); + + // Both operations behave identically for expanding. + EXPECT_TRUE(map_all.GetBit(0)); + EXPECT_TRUE(map_all.GetBit(1)); + EXPECT_TRUE(map_any.GetBit(0)); + EXPECT_TRUE(map_any.GetBit(1)); + + EXPECT_FALSE(map_all.GetBit(2)); + EXPECT_FALSE(map_all.GetBit(3)); + EXPECT_FALSE(map_any.GetBit(2)); + EXPECT_FALSE(map_any.GetBit(3)); + + EXPECT_TRUE(map_all.GetBit(14)); + EXPECT_TRUE(map_all.GetBit(15)); + EXPECT_TRUE(map_any.GetBit(14)); + EXPECT_TRUE(map_any.GetBit(15)); + } +} + class RangeTrackerTest : public ::testing::Test { protected: std::vector> FreeRanges() {