Add tests for PageFlags::GetSinglePageBitmaps.

tcmalloc/internal/BUILD

@@ -520,6 +520,7 @@ cc_test(
     deps = [
         ":allocation_guard",
         ":pageflags",
+        ":range_tracker",
         ":strerror",
         ":util",
         "@com_github_google_benchmark//:benchmark",

tcmalloc/internal/pageflags_test.cc

@@ -51,6 +51,7 @@
 #include "absl/time/clock.h"
 #include "absl/time/time.h"
 #include "tcmalloc/internal/allocation_guard.h"
+#include "tcmalloc/internal/range_tracker.h"
 #include "tcmalloc/internal/strerror.h"
 #include "tcmalloc/internal/util.h"
 
@@ -80,6 +81,11 @@ class PageFlagsFriend {
     return r_.IsHugepageBacked(addr);
   }
 
+  decltype(auto) GetSinglePageBitmaps(const void* addr,
+                                      ResidencyBitmap& stale) {
+    return r_.GetSinglePageBitmaps(addr, stale);
+  }
+
   void SetCachedScanSeconds(
       decltype(PageFlags::cached_scan_seconds_) scan_seconds) {
     r_.cached_scan_seconds_ = scan_seconds;
@@ -594,6 +600,69 @@ TEST(StaleSeconds, TextOverflow) {
   ExpectStaleSecondsFailedReadFrom(fake_stale_seconds);
 }
 
+TEST(PageFlagsTest, GetSinglePageBitmapsErrorCases) {
+  {
+    PageFlagsFriend s;
+    ResidencyBitmap stale;
+    EXPECT_EQ(s.GetSinglePageBitmaps(reinterpret_cast<const void*>(1), stale),
+              absl::StatusCode::kFailedPrecondition);
+  }
+  {
+    PageFlagsFriend s("/dev/null/impossible");
+    ResidencyBitmap stale;
+    EXPECT_EQ(s.GetSinglePageBitmaps(nullptr, stale),
+              absl::StatusCode::kUnavailable);
+  }
+  {
+    std::string fake_pageflags =
+        absl::StrCat(testing::TempDir(), "/fake_pageflags_short");
+    SetContents(fake_pageflags, "x");
+    PageFlagsFriend s(fake_pageflags);
+    ResidencyBitmap stale;
+    EXPECT_EQ(s.GetSinglePageBitmaps(nullptr, stale),
+              absl::StatusCode::kUnavailable);
+  }
+  {
+    std::string fake_pageflags =
+        absl::StrCat(testing::TempDir(), "/fake_pageflags_tail");
+    const size_t kPageSize = getpagesize();
+    const size_t kNativePagesInHugePage = kHugePageSize / kPageSize;
+    std::vector<uint64_t> data(kNativePagesInHugePage, kPageTail);
+
+    std::string content(reinterpret_cast<const char*>(data.data()),
+                        data.size() * sizeof(uint64_t));
+    SetContents(fake_pageflags, content);
+
+    PageFlagsFriend s(fake_pageflags);
+    ResidencyBitmap stale;
+    EXPECT_EQ(s.GetSinglePageBitmaps(nullptr, stale),
+              absl::StatusCode::kFailedPrecondition);
+  }
+}
+
+TEST(PageFlagsTest, GetSinglePageBitmapsSuccess) {
+  std::string fake_pageflags =
+      absl::StrCat(testing::TempDir(), "/fake_pageflags_correct");
+  std::vector<uint64_t> data(kMaxResidencyBits, 0);
+
+  for (size_t i = 0; i < kMaxResidencyBits; ++i) {
+    if (i % 2 == 0) {
+      data[i] |= kPageStale;
+    }
+  }
+
+  std::string content(reinterpret_cast<const char*>(data.data()),
+                      data.size() * sizeof(uint64_t));
+  SetContents(fake_pageflags, content);
+
+  PageFlagsFriend s(fake_pageflags);
+  ResidencyBitmap stale;
+  stale.Clear();
+
+  EXPECT_EQ(s.GetSinglePageBitmaps(nullptr, stale), absl::StatusCode::kOk);
+  EXPECT_EQ(stale.CountBits(), kMaxResidencyBits / 2);
+}
+
 }  // namespace
 }  // namespace tcmalloc_internal
 }  // namespace tcmalloc