From 4fb363efd72fdeeac5bae5d4f5585e36a120ffb9 Mon Sep 17 00:00:00 2001 From: Ben Reisner Date: Fri, 10 Jul 2026 13:34:58 -0700 Subject: [PATCH] Simpler approach maybe PiperOrigin-RevId: 945868435 --- tcmalloc/common.h | 3 +- tcmalloc/huge_page_filler.h | 50 ++++++------------ tcmalloc/huge_page_filler_test.cc | 4 +- tcmalloc/internal/range_tracker.h | 26 ++++++++++ tcmalloc/internal/range_tracker_test.cc | 68 +++++++++++++++++++++++++ 5 files changed, 113 insertions(+), 38 deletions(-) diff --git a/tcmalloc/common.h b/tcmalloc/common.h index 4e7165c73..ad2ccaff6 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,7 +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) inline constexpr size_t kAlignmentShift = diff --git a/tcmalloc/huge_page_filler.h b/tcmalloc/huge_page_filler.h index ffab64510..8ee578718 100644 --- a/tcmalloc/huge_page_filler.h +++ b/tcmalloc/huge_page_filler.h @@ -1620,47 +1620,27 @@ 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; + tcmalloc_internal::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.Contract( + kHugePageSize / GetPageSize()); + 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( diff --git a/tcmalloc/huge_page_filler_test.cc b/tcmalloc/huge_page_filler_test.cc index cb6cfbec3..51ca777de 100644 --- a/tcmalloc/huge_page_filler_test.cc +++ b/tcmalloc/huge_page_filler_test.cc @@ -344,7 +344,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; }; @@ -1428,7 +1428,7 @@ TEST_F(FillerTest, ReleaseFreePagesWhenAnyPageIsSwapped) { EXPECT_FALSE(pageflags.IsHugepageBacked(pa.p.start_addr()).value()); Bitmap unbacked, swapped; - swapped.SetRange(/*index=*/1, /*n=*/1); + swapped.SetRange(/*index=*/0, /*n=*/2); residency.SetUnbackedAndSwappedBitmaps(pa.p.start_addr(), unbacked, swapped); } diff --git a/tcmalloc/internal/range_tracker.h b/tcmalloc/internal/range_tracker.h index e7958af23..2352d05b9 100644 --- a/tcmalloc/internal/range_tracker.h +++ b/tcmalloc/internal/range_tracker.h @@ -61,6 +61,32 @@ class Bitmap { // Clears the lowest set bit. Special case is faster than more flexible code. void ClearLowestBit(); + template + Bitmap Contract(size_t src_len) const { + 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); + + 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 (GetBit(src_idx + j)) { + count++; + } + } + if (count == src_per_dst) { + res.SetRange(dst_idx, dst_per_src); + } + dst_idx += dst_per_src; + } + return res; + } + template size_t PopBatch(Visitor visitor, size_t limit) { size_t count = 0; diff --git a/tcmalloc/internal/range_tracker_test.cc b/tcmalloc/internal/range_tracker_test.cc index 521e1d9c9..91a4d7b94 100644 --- a/tcmalloc/internal/range_tracker_test.cc +++ b/tcmalloc/internal/range_tracker_test.cc @@ -347,6 +347,74 @@ TEST_F(BitmapTest, PopBatch) { EXPECT_TRUE(map.IsZero()); } +#ifndef NDEBUG +TEST_F(BitmapTest, ContractAssertionFailures) { + Bitmap<64> map1; + // 1. src_len > N (64) + EXPECT_DEATH(map1.Contract<64>(128), ""); + // 2. src_len (64) doesn't divide M (10) and vice versa + EXPECT_DEATH(map1.Contract<10>(64), ""); +} +#endif + +TEST_F(BitmapTest, Contract) { + // Test case 1: N == M (equal) + { + Bitmap<64> map1; + map1.SetBit(3); + map1.SetBit(17); + map1.SetBit(63); + + auto map2 = map1.Contract<64>(64); + for (size_t i = 0; i < 64; ++i) { + EXPECT_EQ(map1.GetBit(i), map2.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 + map1.SetBit(2); + + // Group 7: bits 14, 15 -> both set + map1.SetBit(14); + map1.SetBit(15); + + auto map2 = map1.Contract<8>(16); + EXPECT_TRUE(map2.GetBit(0)); + EXPECT_FALSE(map2.GetBit(1)); + EXPECT_FALSE(map2.GetBit(2)); + EXPECT_FALSE(map2.GetBit(3)); + EXPECT_FALSE(map2.GetBit(4)); + EXPECT_FALSE(map2.GetBit(5)); + EXPECT_FALSE(map2.GetBit(6)); + EXPECT_TRUE(map2.GetBit(7)); + } + + // Test case 3: N < M (expanding) + { + Bitmap<64> map1; + // We expand by 2 bits. + map1.SetBit(0); + map1.SetBit(7); + + auto map2 = map1.Contract<16>(8); + EXPECT_TRUE(map2.GetBit(0)); + EXPECT_TRUE(map2.GetBit(1)); + EXPECT_FALSE(map2.GetBit(2)); + EXPECT_FALSE(map2.GetBit(3)); + + EXPECT_TRUE(map2.GetBit(14)); + EXPECT_TRUE(map2.GetBit(15)); + } +} + class RangeTrackerTest : public ::testing::Test { protected: std::vector> FreeRanges() {