From f2d16f312427e59603751b1c845a275c2c2b4a39 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/internal/pageflags.h | 11 ++++ tcmalloc/internal/pageflags_test.cc | 60 ++++++++++++++++++++ tcmalloc/internal/range_tracker.h | 25 +++++++++ tcmalloc/internal/range_tracker_test.cc | 68 ++++++++++++++++++++++ tcmalloc/internal/residency.h | 19 ++++++- tcmalloc/internal/residency_test.cc | 75 +++++++++++++++++++++++++ 6 files changed, 255 insertions(+), 3 deletions(-) diff --git a/tcmalloc/internal/pageflags.h b/tcmalloc/internal/pageflags.h index c4a62aba1..8ee7fdc4d 100644 --- a/tcmalloc/internal/pageflags.h +++ b/tcmalloc/internal/pageflags.h @@ -70,6 +70,17 @@ class PageFlagsBase { virtual std::optional Get(const void* addr, size_t size) = 0; virtual absl::StatusCode GetSinglePageBitmaps(const void* addr, ResidencyBitmap& stale) = 0; + template + absl::StatusCode GetSinglePageBitmapsWrapper(const void* addr, + Bitmap& stale) { + ResidencyBitmap temp; + auto res = GetSinglePageBitmaps(addr, temp); + if (res != absl::StatusCode::kOk) { + return res; + } + stale = Contract(temp, kHugePageSize / GetPageSize()); + return res; + } }; // PageFlags offers a look at kernel page flags to identify pieces of memory as diff --git a/tcmalloc/internal/pageflags_test.cc b/tcmalloc/internal/pageflags_test.cc index e6e67a75c..9cc2c1d5c 100644 --- a/tcmalloc/internal/pageflags_test.cc +++ b/tcmalloc/internal/pageflags_test.cc @@ -33,6 +33,7 @@ #include #include #include +#include #include #include @@ -86,6 +87,12 @@ class PageFlagsFriend { return r_.GetSinglePageBitmaps(addr, stale); } + template + absl::StatusCode GetSinglePageBitmapsWrapper(const void* addr, + Bitmap& stale) { + return r_.GetSinglePageBitmapsWrapper(addr, stale); + } + void SetCachedScanSeconds( decltype(PageFlags::cached_scan_seconds_) scan_seconds) { r_.cached_scan_seconds_ = scan_seconds; @@ -663,6 +670,59 @@ TEST(PageFlagsTest, GetSinglePageBitmapsSuccess) { EXPECT_EQ(stale.CountBits(), kMaxResidencyBits / 2); } +template +class GetSinglePageBitmapsWrapperTest : public ::testing::Test {}; + +using WrapperTestSizes = + ::testing::Types, + std::integral_constant, + std::integral_constant>; + +TYPED_TEST_SUITE(GetSinglePageBitmapsWrapperTest, WrapperTestSizes); + +TYPED_TEST(GetSinglePageBitmapsWrapperTest, AlternatingStale) { + constexpr size_t N = TypeParam::value; + + std::string fake_pageflags = + absl::StrCat(testing::TempDir(), "/fake_pageflags_wrapper_", N); + const size_t kHardwarePageSize = getpagesize(); + const size_t kNativePagesInHugePage = kHugePageSize / kHardwarePageSize; + std::vector data(kNativePagesInHugePage, 0); + + data[0] |= kPageStale; + data[1] |= kPageStale; + data[2] |= kPageStale; + + std::string content(reinterpret_cast(data.data()), + data.size() * sizeof(uint64_t)); + SetContents(fake_pageflags, content); + + PageFlagsFriend s(fake_pageflags); + Bitmap stale; + stale.Clear(); + + EXPECT_EQ(s.GetSinglePageBitmapsWrapper(nullptr, stale), + absl::StatusCode::kOk); + + if constexpr (N == kMaxResidencyBits) { + EXPECT_EQ(stale.CountBits(), 3); + EXPECT_TRUE(stale.GetBit(0)); + EXPECT_TRUE(stale.GetBit(1)); + EXPECT_TRUE(stale.GetBit(2)); + EXPECT_FALSE(stale.GetBit(3)); + } else if constexpr (N == kMaxResidencyBits / 2) { + EXPECT_EQ(stale.CountBits(), 1); + EXPECT_TRUE(stale.GetBit(0)); + EXPECT_FALSE(stale.GetBit(1)); + } else if constexpr (N == kMaxResidencyBits * 2) { + EXPECT_EQ(stale.CountBits(), 6); + for (size_t i = 0; i < 6; ++i) { + EXPECT_TRUE(stale.GetBit(i)) << i; + } + EXPECT_FALSE(stale.GetBit(6)); + } +} + } // namespace } // namespace tcmalloc_internal } // namespace tcmalloc diff --git a/tcmalloc/internal/range_tracker.h b/tcmalloc/internal/range_tracker.h index e7958af23..5d8e81cf3 100644 --- a/tcmalloc/internal/range_tracker.h +++ b/tcmalloc/internal/range_tracker.h @@ -118,6 +118,31 @@ class Bitmap { ssize_t FindValueBackwards(size_t index) const; }; +template +Bitmap Contract(const Bitmap& src, size_t src_len) { + 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 (src.GetBit(src_idx + j)) { + count++; + } + } + if (count == src_per_dst) { + res.SetRange(dst_idx, dst_per_src); + } + 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..efb7c3c58 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 (64) doesn't divide M (10) and vice versa + EXPECT_DEATH(Contract<10>(map1, 64), ""); + // 2. src_len > N (64) + EXPECT_DEATH(Contract<256>(map1, 128), ""); +} +#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 = Contract<64>(map1, 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 = Contract<8>(map1, 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 = Contract<16>(map1, 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() { diff --git a/tcmalloc/internal/residency.h b/tcmalloc/internal/residency.h index 440008405..f7d2547d3 100644 --- a/tcmalloc/internal/residency.h +++ b/tcmalloc/internal/residency.h @@ -47,18 +47,31 @@ class Residency { virtual inline size_t GetNativePagesInHugePage() const = 0; // Struct is ordered with bitmaps first to optimize cacheline usage. - struct SinglePageBitmaps { - Bitmap unbacked; - Bitmap swapped; + template + struct SinglePageBitmapsWrapper { + Bitmap unbacked; + Bitmap swapped; absl::StatusCode status; }; + using SinglePageBitmaps = SinglePageBitmapsWrapper; + // Using a hugepage-aligned address, parse through /proc/self/pagemap // 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; + + template + SinglePageBitmapsWrapper GetUnbackedAndSwappedBitmapsWrapper( + const void* addr) { + SinglePageBitmaps res = GetUnbackedAndSwappedBitmaps(addr); + return SinglePageBitmapsWrapper{ + .unbacked = Contract(res.unbacked, GetNativePagesInHugePage()), + .swapped = Contract(res.swapped, GetNativePagesInHugePage()), + .status = res.status}; + } }; // Residency offers information about memory residency: whether or not specific diff --git a/tcmalloc/internal/residency_test.cc b/tcmalloc/internal/residency_test.cc index c93c2d906..92e2d9f85 100644 --- a/tcmalloc/internal/residency_test.cc +++ b/tcmalloc/internal/residency_test.cc @@ -25,6 +25,7 @@ #include #include #include +#include #include #include @@ -369,6 +370,80 @@ TEST(PageMapIntegrationTest, WorksOnActualData) { EXPECT_TRUE(res.swapped.IsZero()); } +template +class PageMapIntegrationTest : public ::testing::Test {}; + +using WorksOnActualDataTestSizes = + ::testing::Types, + std::integral_constant, + std::integral_constant>; + +TYPED_TEST_SUITE(PageMapIntegrationTest, WorksOnActualDataTestSizes); + +TYPED_TEST(PageMapIntegrationTest, WorksOnActualDataWrapper) { + constexpr size_t N = TypeParam::value; + std::optional g; + void* addr = mmap(nullptr, 4 << 20, PROT_WRITE, + MAP_ANONYMOUS | MAP_POPULATE | MAP_PRIVATE, -1, 0); + ASSERT_NE(addr, MAP_FAILED) << errno; + auto position = reinterpret_cast(addr); + if ((position & (kHugePageSize - 1)) != 0) { + position |= kHugePageSize - 1; + position++; + addr = reinterpret_cast(position); + } + g.emplace(); + ResidencyPageMap r; + auto res = r.GetUnbackedAndSwappedBitmapsWrapper(addr); + g.reset(); + ASSERT_EQ(res.status, absl::StatusCode::kOk); + EXPECT_TRUE(res.unbacked.IsZero()); + EXPECT_TRUE(res.swapped.IsZero()); + // Unmap consecutive pages starting at page 2, and page 16 + ASSERT_EQ(munmap(reinterpret_cast(addr) + 2 * 4096, 2 * 4096), 0) + << errno; + ASSERT_EQ(munmap(reinterpret_cast(addr) + 16 * 4096, 2 * 4096), 0) + << errno; + + g.emplace(); + res = r.GetUnbackedAndSwappedBitmapsWrapper(addr); + g.reset(); + ASSERT_EQ(res.status, absl::StatusCode::kOk); + EXPECT_FALSE(res.unbacked.IsZero()); + + if constexpr (N == kMaxResidencyBits) { + ASSERT_TRUE(res.unbacked.GetBit(2)); + ASSERT_TRUE(res.unbacked.GetBit(3)); + ASSERT_TRUE(res.unbacked.GetBit(16)); + ASSERT_TRUE(res.unbacked.GetBit(17)); + res.unbacked.ClearBit(2); + res.unbacked.ClearBit(3); + res.unbacked.ClearBit(16); + res.unbacked.ClearBit(17); + } else if constexpr (N == kMaxResidencyBits / 2) { + ASSERT_TRUE(res.unbacked.GetBit(1)); + ASSERT_TRUE(res.unbacked.GetBit(8)); + res.unbacked.ClearBit(1); + res.unbacked.ClearBit(8); + } else if constexpr (N == kMaxResidencyBits * 2) { + ASSERT_TRUE(res.unbacked.GetBit(4)); + ASSERT_TRUE(res.unbacked.GetBit(5)); + ASSERT_TRUE(res.unbacked.GetBit(32)); + ASSERT_TRUE(res.unbacked.GetBit(33)); + res.unbacked.ClearBit(4); + res.unbacked.ClearBit(5); + res.unbacked.ClearBit(6); + res.unbacked.ClearBit(7); + res.unbacked.ClearBit(32); + res.unbacked.ClearBit(33); + res.unbacked.ClearBit(34); + res.unbacked.ClearBit(35); + } + + EXPECT_TRUE(res.unbacked.IsZero()); + EXPECT_TRUE(res.swapped.IsZero()); +} + } // namespace } // namespace tcmalloc_internal } // namespace tcmalloc